Added redis impressions async storage

chillaq · chillaq · commit 630f89055411 · 2023-07-06T11:11:53.000-07:00
diff --git a/splitio/storage/redis.py b/splitio/storage/redis.py
@@ -385,24 +385,12 @@ def get_segments_keys_count(self):
         """
         return 0
 
-class RedisImpressionsStorage(ImpressionStorage, ImpressionPipelinedStorage):
-    """Redis based event storage class."""
+class RedisImpressionsStorageBase(ImpressionStorage, ImpressionPipelinedStorage):
+    """Redis based event storage base class."""
 
     IMPRESSIONS_QUEUE_KEY = 'SPLITIO.impressions'
     IMPRESSIONS_KEY_DEFAULT_TTL = 3600
 
-    def __init__(self, redis_client, sdk_metadata):
-        """
-        Class constructor.
-
-        :param redis_client: Redis client or compliant interface.
-        :type redis_client: splitio.storage.adapters.redis.RedisAdapter
-        :param sdk_metadata: SDK & Machine information.
-        :type sdk_metadata: splitio.client.util.SdkMetadata
-        """
-        self._redis = redis_client
-        self._sdk_metadata = sdk_metadata
-
     def _wrap_impressions(self, impressions):
         """
         Wrap impressions to be stored in redis
@@ -444,8 +432,7 @@ def expire_key(self, total_keys, inserted):
         :param inserted: added keys.
         :type inserted: int
         """
-        if total_keys == inserted:
-            self._redis.expire(self.IMPRESSIONS_QUEUE_KEY, self.IMPRESSIONS_KEY_DEFAULT_TTL)
+        pass
 
     def add_impressions_to_pipe(self, impressions, pipe):
         """
@@ -461,6 +448,61 @@ def add_impressions_to_pipe(self, impressions, pipe):
         _LOGGER.debug(bulk_impressions)
         pipe.rpush(self.IMPRESSIONS_QUEUE_KEY, *bulk_impressions)
 
+    def put(self, impressions):
+        """
+        Add an impression to the redis storage.
+
+        :param impressions: Impression to add to the queue.
+        :type impressions: splitio.models.impressions.Impression
+
+        :return: Whether the impression has been added or not.
+        :rtype: bool
+        """
+        pass
+
+    def pop_many(self, count):
+        """
+        Pop the oldest N events from storage.
+
+        :param count: Number of events to pop.
+        :type count: int
+        """
+        raise NotImplementedError('Only redis-consumer mode is supported.')
+
+    def clear(self):
+        """
+        Clear data.
+        """
+        raise NotImplementedError('Not supported for redis.')
+
+
+class RedisImpressionsStorage(RedisImpressionsStorageBase):
+    """Redis based event storage class."""
+
+    def __init__(self, redis_client, sdk_metadata):
+        """
+        Class constructor.
+
+        :param redis_client: Redis client or compliant interface.
+        :type redis_client: splitio.storage.adapters.redis.RedisAdapter
+        :param sdk_metadata: SDK & Machine information.
+        :type sdk_metadata: splitio.client.util.SdkMetadata
+        """
+        self._redis = redis_client
+        self._sdk_metadata = sdk_metadata
+
+    def expire_key(self, total_keys, inserted):
+        """
+        Set expire
+
+        :param total_keys: length of keys.
+        :type total_keys: int
+        :param inserted: added keys.
+        :type inserted: int
+        """
+        if total_keys == inserted:
+            self._redis.expire(self.IMPRESSIONS_QUEUE_KEY, self.IMPRESSIONS_KEY_DEFAULT_TTL)
+
     def put(self, impressions):
         """
         Add an impression to the redis storage.
@@ -483,20 +525,55 @@ def put(self, impressions):
             _LOGGER.error('Error: ', exc_info=True)
             return False
 
-    def pop_many(self, count):
+
+class RedisImpressionsStorageAsync(RedisImpressionsStorageBase):
+    """Redis based event storage async class."""
+
+    def __init__(self, redis_client, sdk_metadata):
         """
-        Pop the oldest N events from storage.
+        Class constructor.
 
-        :param count: Number of events to pop.
-        :type count: int
+        :param redis_client: Redis client or compliant interface.
+        :type redis_client: splitio.storage.adapters.redis.RedisAdapter
+        :param sdk_metadata: SDK & Machine information.
+        :type sdk_metadata: splitio.client.util.SdkMetadata
         """
-        raise NotImplementedError('Only redis-consumer mode is supported.')
+        self._redis = redis_client
+        self._sdk_metadata = sdk_metadata
 
-    def clear(self):
+    async def expire_key(self, total_keys, inserted):
         """
-        Clear data.
+        Set expire
+
+        :param total_keys: length of keys.
+        :type total_keys: int
+        :param inserted: added keys.
+        :type inserted: int
         """
-        raise NotImplementedError('Not supported for redis.')
+        if total_keys == inserted:
+            await self._redis.expire(self.IMPRESSIONS_QUEUE_KEY, self.IMPRESSIONS_KEY_DEFAULT_TTL)
+
+    async def put(self, impressions):
+        """
+        Add an impression to the redis storage.
+
+        :param impressions: Impression to add to the queue.
+        :type impressions: splitio.models.impressions.Impression
+
+        :return: Whether the impression has been added or not.
+        :rtype: bool
+        """
+        bulk_impressions = self._wrap_impressions(impressions)
+        try:
+            _LOGGER.debug("Adding Impressions to redis key %s" % (self.IMPRESSIONS_QUEUE_KEY))
+            _LOGGER.debug(bulk_impressions)
+            inserted = await self._redis.rpush(self.IMPRESSIONS_QUEUE_KEY, *bulk_impressions)
+            await self.expire_key(inserted, len(bulk_impressions))
+            return True
+        except RedisAdapterException:
+            _LOGGER.error('Something went wrong when trying to add impression to redis')
+            _LOGGER.error('Error: ', exc_info=True)
+            return False
 
 
 class RedisEventsStorage(EventStorage):
diff --git a/tests/storage/test_redis.py b/tests/storage/test_redis.py
@@ -6,10 +6,11 @@
 import unittest.mock as mock
 import pytest
 
+from splitio.optional.loaders import asyncio
 from splitio.client.util import get_metadata, SdkMetadata
-from splitio.storage.redis import RedisEventsStorage, RedisImpressionsStorage, \
+from splitio.storage.redis import RedisEventsStorage, RedisImpressionsStorage, RedisImpressionsStorageAsync, \
     RedisSegmentStorage, RedisSplitStorage, RedisTelemetryStorage
-from splitio.storage.adapters.redis import RedisAdapter, RedisAdapterException, build
+from splitio.storage.adapters.redis import RedisAdapter, RedisAdapterAsync, RedisAdapterException, build
 from splitio.models.segments import Segment
 from splitio.models.impressions import Impression
 from splitio.models.events import Event, EventWrapper
@@ -334,6 +335,124 @@ def test_add_impressions_to_pipe(self, mocker):
         storage.add_impressions_to_pipe(impressions, adapter)
         assert adapter.rpush.mock_calls == [mocker.call('SPLITIO.impressions', *to_validate)]
 
+class RedisImpressionsStorageAsyncTests(object):  # pylint: disable=too-few-public-methods
+    """Redis Impressions async storage test cases."""
+
+    def test_wrap_impressions(self, mocker):
+        """Test wrap impressions."""
+        adapter = mocker.Mock(spec=RedisAdapterAsync)
+        metadata = get_metadata({})
+        storage = RedisImpressionsStorageAsync(adapter, metadata)
+
+        impressions = [
+            Impression('key1', 'feature1', 'on', 'some_label', 123456, 'buck1', 321654),
+            Impression('key2', 'feature2', 'on', 'some_label', 123456, 'buck1', 321654),
+            Impression('key3', 'feature2', 'on', 'some_label', 123456, 'buck1', 321654),
+            Impression('key4', 'feature1', 'on', 'some_label', 123456, 'buck1', 321654)
+        ]
+
+        to_validate = [json.dumps({
+            'm': {  # METADATA PORTION
+                's': metadata.sdk_version,
+                'n': metadata.instance_name,
+                'i': metadata.instance_ip,
+            },
+            'i': {  # IMPRESSION PORTION
+                'k': impression.matching_key,
+                'b': impression.bucketing_key,
+                'f': impression.feature_name,
+                't': impression.treatment,
+                'r': impression.label,
+                'c': impression.change_number,
+                'm': impression.time,
+            }
+        }) for impression in impressions]
+
+        assert storage._wrap_impressions(impressions) == to_validate
+
+    @pytest.mark.asyncio
+    async def test_add_impressions(self, mocker):
+        """Test that adding impressions to storage works."""
+        adapter = mocker.Mock(spec=RedisAdapterAsync)
+        metadata = get_metadata({})
+        storage = RedisImpressionsStorageAsync(adapter, metadata)
+
+        impressions = [
+            Impression('key1', 'feature1', 'on', 'some_label', 123456, 'buck1', 321654),
+            Impression('key2', 'feature2', 'on', 'some_label', 123456, 'buck1', 321654),
+            Impression('key3', 'feature2', 'on', 'some_label', 123456, 'buck1', 321654),
+            Impression('key4', 'feature1', 'on', 'some_label', 123456, 'buck1', 321654)
+        ]
+        self.key = None
+        self.imps = None
+        async def rpush(key, *imps):
+            self.key = key
+            self.imps = imps
+
+        adapter.rpush = rpush
+        assert await storage.put(impressions) is True
+
+        to_validate = [json.dumps({
+            'm': {  # METADATA PORTION
+                's': metadata.sdk_version,
+                'n': metadata.instance_name,
+                'i': metadata.instance_ip,
+            },
+            'i': {  # IMPRESSION PORTION
+                'k': impression.matching_key,
+                'b': impression.bucketing_key,
+                'f': impression.feature_name,
+                't': impression.treatment,
+                'r': impression.label,
+                'c': impression.change_number,
+                'm': impression.time,
+            }
+        }) for impression in impressions]
+
+        assert self.key == 'SPLITIO.impressions'
+        assert self.imps == tuple(to_validate)
+
+        # Assert that if an exception is thrown it's caught and False is returned
+        adapter.reset_mock()
+
+        async def rpush2(key, *imps):
+            raise RedisAdapterException('something')
+        adapter.rpush = rpush2
+        assert await storage.put(impressions) is False
+
+    def test_add_impressions_to_pipe(self, mocker):
+        """Test that adding impressions to storage works."""
+        adapter = mocker.Mock(spec=RedisAdapterAsync)
+        metadata = get_metadata({})
+        storage = RedisImpressionsStorageAsync(adapter, metadata)
+
+        impressions = [
+            Impression('key1', 'feature1', 'on', 'some_label', 123456, 'buck1', 321654),
+            Impression('key2', 'feature2', 'on', 'some_label', 123456, 'buck1', 321654),
+            Impression('key3', 'feature2', 'on', 'some_label', 123456, 'buck1', 321654),
+            Impression('key4', 'feature1', 'on', 'some_label', 123456, 'buck1', 321654)
+        ]
+
+        to_validate = [json.dumps({
+            'm': {  # METADATA PORTION
+                's': metadata.sdk_version,
+                'n': metadata.instance_name,
+                'i': metadata.instance_ip,
+            },
+            'i': {  # IMPRESSION PORTION
+                'k': impression.matching_key,
+                'b': impression.bucketing_key,
+                'f': impression.feature_name,
+                't': impression.treatment,
+                'r': impression.label,
+                'c': impression.change_number,
+                'm': impression.time,
+            }
+        }) for impression in impressions]
+
+        storage.add_impressions_to_pipe(impressions, adapter)
+        assert adapter.rpush.mock_calls == [mocker.call('SPLITIO.impressions', *to_validate)]
+
 
 class RedisEventsStorageTests(object):  # pylint: disable=too-few-public-methods
     """Redis Impression storage test cases."""
