Merge pull request #398 from splitio/async-redis-impression-storage

chillaq · web-flow · commit bcf8ead46b4c · 2023-07-11T14:26:41.000-07:00
Added redis impressions async storage
diff --git a/splitio/storage/redis.py b/splitio/storage/redis.py
@@ -562,24 +562,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
@@ -621,8 +609,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):
         """
@@ -638,6 +625,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.
@@ -660,20 +702,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
@@ -8,7 +8,7 @@
 
 from splitio.client.util import get_metadata, SdkMetadata
 from splitio.optional.loaders import asyncio
-from splitio.storage.redis import RedisEventsStorage, RedisImpressionsStorage, \
+from splitio.storage.redis import RedisEventsStorage, RedisImpressionsStorage, RedisImpressionsStorageAsync\
     RedisSegmentStorage, RedisSplitStorage, RedisSplitStorageAsync, RedisTelemetryStorage
 from splitio.storage.adapters.redis import RedisAdapter, RedisAdapterException, build
 from redis.asyncio.client import Redis as aioredis
@@ -590,6 +590,167 @@ 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)]
 
+    def test_expire_key(self, mocker):
+        adapter = mocker.Mock(spec=RedisAdapter)
+        metadata = get_metadata({})
+        storage = RedisImpressionsStorage(adapter, metadata)
+
+        self.key = None
+        self.ttl = None
+        def expire(key, ttl):
+            self.key = key
+            self.ttl = ttl
+        adapter.expire = expire
+
+        storage.expire_key(2, 2)
+        assert self.key == 'SPLITIO.impressions'
+        assert self.ttl == 3600
+
+        self.key = None
+        storage.expire_key(2, 1)
+        assert self.key == None
+
+
+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)]
+
+    @pytest.mark.asyncio
+    async def test_expire_key(self, mocker):
+        adapter = mocker.Mock(spec=RedisAdapterAsync)
+        metadata = get_metadata({})
+        storage = RedisImpressionsStorageAsync(adapter, metadata)
+
+        self.key = None
+        self.ttl = None
+        async def expire(key, ttl):
+            self.key = key
+            self.ttl = ttl
+        adapter.expire = expire
+
+        await storage.expire_key(2, 2)
+        assert self.key == 'SPLITIO.impressions'
+        assert self.ttl == 3600
+
+        self.key = None
+        await storage.expire_key(2, 1)
+        assert self.key == None
+
+
 
 class RedisEventsStorageTests(object):  # pylint: disable=too-few-public-methods
     """Redis Impression storage test cases."""
