added wrapper for PipelineRedis

Author: Matias Melograno

splitio/storage/adapters/redis.py

@@ -53,25 +53,18 @@ class SentinelConfigurationException(Exception):
     pass
 
 
-class RedisAdapter(object):  # pylint: disable=too-many-public-methods
-    """
-    Instance decorator for Redis clients such as StrictRedis.
-
-    Adds an extra layer handling addition/removal of user prefix when handling
-    keys
-    """
+class PrefixTrait(object):
+    """Prefix generator."""
 
-    def __init__(self, decorated, prefix=None):
+    def __init__(self, prefix=None):
         """
-        Store the user prefix and the redis client instance.
+        Class constructor.
 
-        :param decorated: Instance of redis cache client to decorate.
         :param prefix: User prefix to add.
         """
         self._prefix = prefix
-        self._decorated = decorated
 
-    def _add_prefix(self, k):
+    def add_prefix(self, k):
         """
         Add a prefix to the contents of k.
 
@@ -105,7 +98,7 @@ def _add_prefix(self, k):
             "Cannot append prefix correctly. Wrong type for key(s) provided"
         )
 
-    def _remove_prefix(self, k):
+    def remove_prefix(self, k):
         """
         Remove the user prefix from a key before handling it back to the requester.
 
@@ -127,6 +120,25 @@ def _remove_prefix(self, k):
             "Cannot remove prefix correctly. Wrong type for key(s) provided"
         )
 
+
+class RedisAdapter(object):  # pylint: disable=too-many-public-methods
+    """
+    Instance decorator for Redis clients such as StrictRedis.
+
+    Adds an extra layer handling addition/removal of user prefix when handling
+    keys
+    """
+
+    def __init__(self, decorated, prefix=None):
+        """
+        Store the user prefix and the redis client instance.
+
+        :param decorated: Instance of redis cache client to decorate.
+        :param prefix: User prefix to add.
+        """
+        self._decorated = decorated
+        self._prefix_trait = PrefixTrait(prefix)
+
     # Below starts a list of methods that implement the interface of a standard
     # redis client.
 
@@ -135,7 +147,7 @@ def keys(self, pattern):
         try:
             return [
                 _bytes_to_string(key)
-                for key in self._remove_prefix(self._decorated.keys(self._add_prefix(pattern)))
+                for key in self._prefix_trait.remove_prefix(self._decorated.keys(self._prefix_trait.add_prefix(pattern)))
             ]
         except RedisError as exc:
             raise_from(RedisAdapterException('Failed to execute keys operation'), exc)
@@ -144,43 +156,43 @@ def set(self, name, value, *args, **kwargs):
         """Mimic original redis function but using user custom prefix."""
         try:
             return self._decorated.set(
-                self._add_prefix(name), value, *args, **kwargs
+                self._prefix_trait.add_prefix(name), value, *args, **kwargs
             )
         except RedisError as exc:
             raise RedisAdapterException('Failed to execute set operation', exc)
 
     def get(self, name):
         """Mimic original redis function but using user custom prefix."""
         try:
-            return _bytes_to_string(self._decorated.get(self._add_prefix(name)))
+            return _bytes_to_string(self._decorated.get(self._prefix_trait.add_prefix(name)))
         except RedisError as exc:
             raise RedisAdapterException('Error executing get operation', exc)
 
     def setex(self, name, time, value):
         """Mimic original redis function but using user custom prefix."""
         try:
-            return self._decorated.setex(self._add_prefix(name), time, value)
+            return self._decorated.setex(self._prefix_trait.add_prefix(name), time, value)
         except RedisError as exc:
             raise RedisAdapterException('Error executing setex operation', exc)
 
     def delete(self, *names):
         """Mimic original redis function but using user custom prefix."""
         try:
-            return self._decorated.delete(*self._add_prefix(list(names)))
+            return self._decorated.delete(*self._prefix_trait.add_prefix(list(names)))
         except RedisError as exc:
             raise_from(RedisAdapterException('Error executing delete operation'), exc)
 
     def exists(self, name):
         """Mimic original redis function but using user custom prefix."""
         try:
-            return self._decorated.exists(self._add_prefix(name))
+            return self._decorated.exists(self._prefix_trait.add_prefix(name))
         except RedisError as exc:
             raise_from(RedisAdapterException('Error executing exists operation'), exc)
 
     def lrange(self, key, start, end):
         """Mimic original redis function but using user custom prefix."""
         try:
-            return self._decorated.lrange(self._add_prefix(key), start, end)
+            return self._decorated.lrange(self._prefix_trait.add_prefix(key), start, end)
         except RedisError as exc:
             raise_from(RedisAdapterException('Error executing exists operation'), exc)
 
@@ -189,7 +201,7 @@ def mget(self, names):
         try:
             return [
                 _bytes_to_string(item)
-                for item in self._decorated.mget(self._add_prefix(names))
+                for item in self._decorated.mget(self._prefix_trait.add_prefix(names))
             ]
         except RedisError as exc:
             raise_from(RedisAdapterException('Error executing mget operation'), exc)
@@ -199,110 +211,140 @@ def smembers(self, name):
         try:
             return [
                 _bytes_to_string(item)
-                for item in self._decorated.smembers(self._add_prefix(name))
+                for item in self._decorated.smembers(self._prefix_trait.add_prefix(name))
             ]
         except RedisError as exc:
             raise_from(RedisAdapterException('Error executing smembers operation'), exc)
 
     def sadd(self, name, *values):
         """Mimic original redis function but using user custom prefix."""
         try:
-            return self._decorated.sadd(self._add_prefix(name), *values)
+            return self._decorated.sadd(self._prefix_trait.add_prefix(name), *values)
         except RedisError as exc:
             raise_from(RedisAdapterException('Error executing sadd operation'), exc)
 
     def srem(self, name, *values):
         """Mimic original redis function but using user custom prefix."""
         try:
-            return self._decorated.srem(self._add_prefix(name), *values)
+            return self._decorated.srem(self._prefix_trait.add_prefix(name), *values)
         except RedisError as exc:
             raise_from(RedisAdapterException('Error executing srem operation'), exc)
 
     def sismember(self, name, value):
         """Mimic original redis function but using user custom prefix."""
         try:
-            return self._decorated.sismember(self._add_prefix(name), value)
+            return self._decorated.sismember(self._prefix_trait.add_prefix(name), value)
         except RedisError as exc:
             raise_from(RedisAdapterException('Error executing sismember operation'), exc)
 
     def eval(self, script, number_of_keys, *keys):
         """Mimic original redis function but using user custom prefix."""
         try:
-            return self._decorated.eval(script, number_of_keys, *self._add_prefix(list(keys)))
+            return self._decorated.eval(script, number_of_keys, *self._prefix_trait.add_prefix(list(keys)))
         except RedisError as exc:
             raise_from(RedisAdapterException('Error executing eval operation'), exc)
 
     def hset(self, name, key, value):
         """Mimic original redis function but using user custom prefix."""
         try:
-            return self._decorated.hset(self._add_prefix(name), key, value)
+            return self._decorated.hset(self._prefix_trait.add_prefix(name), key, value)
         except RedisError as exc:
             raise_from(RedisAdapterException('Error executing hset operation'), exc)
 
     def hget(self, name, key):
         """Mimic original redis function but using user custom prefix."""
         try:
-            return _bytes_to_string(self._decorated.hget(self._add_prefix(name), key))
+            return _bytes_to_string(self._decorated.hget(self._prefix_trait.add_prefix(name), key))
         except RedisError as exc:
             raise_from(RedisAdapterException('Error executing hget operation'), exc)
 
     def incr(self, name, amount=1):
         """Mimic original redis function but using user custom prefix."""
         try:
-            return self._decorated.incr(self._add_prefix(name), amount)
+            return self._decorated.incr(self._prefix_trait.add_prefix(name), amount)
         except RedisError as exc:
             raise_from(RedisAdapterException('Error executing incr operation'), exc)
 
     def getset(self, name, value):
         """Mimic original redis function but using user custom prefix."""
         try:
-            return _bytes_to_string(self._decorated.getset(self._add_prefix(name), value))
+            return _bytes_to_string(self._decorated.getset(self._prefix_trait.add_prefix(name), value))
         except RedisError as exc:
             raise_from(RedisAdapterException('Error executing getset operation'), exc)
 
     def rpush(self, key, *values):
         """Mimic original redis function but using user custom prefix."""
         try:
-            return self._decorated.rpush(self._add_prefix(key), *values)
+            return self._decorated.rpush(self._prefix_trait.add_prefix(key), *values)
         except RedisError as exc:
             raise_from(RedisAdapterException('Error executing rpush operation'), exc)
 
     def expire(self, key, value):
         """Mimic original redis function but using user custom prefix."""
         try:
-            return self._decorated.expire(self._add_prefix(key), value)
+            return self._decorated.expire(self._prefix_trait.add_prefix(key), value)
         except RedisError as exc:
             raise_from(RedisAdapterException('Error executing expire operation'), exc)
 
     def rpop(self, key):
         """Mimic original redis function but using user custom prefix."""
         try:
-            return _bytes_to_string(self._decorated.rpop(self._add_prefix(key)))
+            return _bytes_to_string(self._decorated.rpop(self._prefix_trait.add_prefix(key)))
         except RedisError as exc:
             raise_from(RedisAdapterException('Error executing rpop operation'), exc)
 
     def ttl(self, key):
         """Mimic original redis function but using user custom prefix."""
         try:
-            return self._decorated.ttl(self._add_prefix(key))
+            return self._decorated.ttl(self._prefix_trait.add_prefix(key))
         except RedisError as exc:
             raise_from(RedisAdapterException('Error executing ttl operation'), exc)
 
     def lpop(self, key):
         """Mimic original redis function but using user custom prefix."""
         try:
-            return self._decorated.lpop(self._add_prefix(key))
+            return self._decorated.lpop(self._prefix_trait.add_prefix(key))
         except RedisError as exc:
             raise_from(RedisAdapterException('Error executing lpop operation'), exc)
 
     def pipeline(self):
         """Mimic original redis pipeline."""
         try:
-            return self._decorated.pipeline()
+            return RedisPipelineAdapter(self._decorated, self._prefix_trait)
         except RedisError as exc:
             raise_from(RedisAdapterException('Error executing ttl operation'), exc)
 
 
+class RedisPipelineAdapter(object):
+    """
+    Instance decorator for Redis Pipeline.
+
+    Adds an extra layer handling addition/removal of user prefix when handling
+    keys
+    """
+    def __init__(self, decorated, prefix_trait):
+        """
+        Store the user prefix and the redis client instance.
+
+        :param decorated: Instance of redis cache client to decorate.
+        :param prefix_trait: Prefix Trait utility
+        """
+        self._prefix_trait = prefix_trait
+        self._pipe = decorated.pipeline()
+
+    def rpush(self, key, *values):
+        """Mimic original redis function but using user custom prefix."""
+        self._pipe.rpush(self._prefix_trait.add_prefix(key), *values)
+
+    def incr(self, name, amount=1):
+        """Mimic original redis function but using user custom prefix."""
+        self._pipe.incr(self._prefix_trait.add_prefix(name), amount)
+
+    def execute(self):
+        """Mimic original redis function but using user custom prefix."""
+        return self._pipe.execute()
+
+
 def _build_default_client(config):  # pylint: disable=too-many-locals
     """
     Build a redis adapter.

tests/integration/test_client_e2e.py

@@ -743,6 +743,7 @@ def teardown_method(self):
         """Clear redis cache."""
         keys_to_delete = [
             "SPLITIO/python-1.2.3/some_ip/latency.sdk.getTreatment.bucket.0",
+            "SPLITIO/python-1.2.3/some_ip/latency.sdk.getTreatmentWithConfig.bucket.0",
             "SPLITIO.segment.human_beigns",
             "SPLITIO.segment.employees.till",
             "SPLITIO.split.sample_feature",
@@ -752,6 +753,7 @@ def teardown_method(self):
             "SPLITIO.split.whitelist_feature",
             "SPLITIO.segment.employees",
             "SPLITIO/python-1.2.3/some_ip/latency.sdk.getTreatments.bucket.0",
+            "SPLITIO/python-1.2.3/some_ip/latency.sdk.getTreatmentsWithConfig.bucket.0",
             "SPLITIO.split.regex_test",
             "SPLITIO.segment.human_beigns.till",
             "SPLITIO.split.boolean_test",
@@ -805,7 +807,7 @@ def setup_method(self):
         self.factory = SplitFactory('some_api_key', storages, True, recorder)  # pylint:disable=attribute-defined-outside-init
 
 
-class LocalhostIntegrationTests(object):  #pylint: disable=too-few-public-methods
+class LocalhostIntegrationTests(object):  # pylint: disable=too-few-public-methods
     """Client & Manager integration tests."""
 
     def test_localhost_e2e(self):

tests/storage/adapters/test_redis_adapter.py

@@ -2,9 +2,10 @@
 
 import pytest
 from splitio.storage.adapters import redis
-from redis import StrictRedis
+from redis import StrictRedis, Redis
 from redis.sentinel import Sentinel
 
+
 class RedisStorageAdapterTests(object):
     """Redis storage adapter test cases."""
 
@@ -178,3 +179,21 @@ def test_sentinel_ssl_fails(self):
                 'redisSentinels': ['a', 'b'],
                 'redisSsl': True,
             })
+
+
+class RedisPipelineAdapterTests(object):
+    """Redis pipelined adapter test cases."""
+
+    def test_forwarding(self, mocker):
+        """Test that all redis functions forward prefix appropriately."""
+        redis_mock = mocker.Mock(StrictRedis)
+        redis_mock_2 = mocker.Mock(Redis)
+        redis_mock.pipeline.return_value = redis_mock_2
+        prefix_trait = redis.PrefixTrait('some_prefix')
+        adapter = redis.RedisPipelineAdapter(redis_mock, prefix_trait)
+
+        adapter.rpush('key1', 'value1', 'value2')
+        assert redis_mock_2.rpush.mock_calls[0] == mocker.call('some_prefix.key1', 'value1', 'value2')
+
+        adapter.incr('key1')
+        assert redis_mock_2.incr.mock_calls[0] == mocker.call('some_prefix.key1', 1)