dev (#49)

* change wrong order of lontitude & latitude in comment of command GEOADD & GEOPOS

* sync optimization from redis-py made by bgreenberg

* update release log

* [fix bug] pipeline callback executed

* sync bug fixed of `geopos` from redis-py made by categulario

* fix error caused by byte decode issues in sentinel

* fix comment of connection pool

* fix sentinel document bug

* add basic transaction support for single node in cluster

* [fix bug] fix bug of get_random_connection reported by myrfy001

* optimize parse_response of cluster transaction

* update cluster scripting test case

Author: NoneGG

aredis/client.py

@@ -465,23 +465,24 @@ async def _execute_command_on_nodes(self, nodes, *args, **kwargs):
                 self.connection_pool.release(connection)
         return self._merge_result(command, res, **kwargs)
 
-    async def pipeline(self, transaction=None, shard_hint=None):
+    async def pipeline(self, transaction=None, shard_hint=None, watches=None):
         """
         Cluster impl:
             Pipelines do not work in cluster mode the same way they do in normal mode.
             Create a clone of this object so that simulating pipelines will work correctly.
             Each command will be called directly when used and when calling execute() will only return the result stack.
+        cluster transaction can only be run with commands in the same node, otherwise error will be raised.
         """
         await self.connection_pool.initialize()
         if shard_hint:
             raise RedisClusterException("shard_hint is deprecated in cluster mode")
 
-        if transaction:
-            raise RedisClusterException("transaction is deprecated in cluster mode")
         from aredis.pipeline import StrictClusterPipeline
         return StrictClusterPipeline(
             connection_pool=self.connection_pool,
             startup_nodes=self.connection_pool.nodes.startup_nodes,
             result_callbacks=self.result_callbacks,
             response_callbacks=self.response_callbacks,
+            transaction=transaction,
+            watches=watches
         )

aredis/commands/transaction.py

@@ -54,8 +54,28 @@ async def unwatch(self):
 
 class ClusterTransactionCommandMixin(TransactionCommandMixin):
 
-    async def transaction(self, *args, **kwargs):
+    async def transaction(self, func, *watches, **kwargs):
         """
-        Transaction is not implemented in cluster mode yet.
+        Convenience method for executing the callable `func` as a transaction
+        while watching all keys specified in `watches`. The 'func' callable
+        should expect a single argument which is a Pipeline object.
+
+        cluster transaction can only be run with commands in the same node,
+        otherwise error will be raised.
         """
-        raise RedisClusterException("method StrictRedisCluster.transaction() is not implemented")
+        shard_hint = kwargs.pop('shard_hint', None)
+        value_from_callable = kwargs.pop('value_from_callable', False)
+        watch_delay = kwargs.pop('watch_delay', None)
+        async with await self.pipeline(True, shard_hint, watches=watches) as pipe:
+            while True:
+                try:
+                    func_value = await func(pipe)
+                    exec_value = await pipe.execute()
+                    return func_value if value_from_callable else exec_value
+                except WatchError:
+                    if watch_delay is not None and watch_delay > 0:
+                        await asyncio.sleep(
+                            watch_delay,
+                            loop=self.connection_pool.loop
+                        )
+                    continue

aredis/exceptions.py

@@ -109,6 +109,12 @@ def __init__(self, resp):
         self.message = resp
 
 
+class ClusterTransactionError(ClusterError):
+
+    def __init__(self, msg):
+        self.msg = msg
+
+
 class AskError(ResponseError):
     """
     src node: MIGRATING to dst node

aredis/nodemanager.py

@@ -100,7 +100,7 @@ def get_redis_link(self, host, port):
             'loop'
         )
         connection_kwargs = {k: v for k, v in self.connection_kwargs.items() if k in allowed_keys}
-        return StrictRedis(host=host, port=port, **connection_kwargs)
+        return StrictRedis(host=host, port=port, decode_responses=True, **connection_kwargs)
 
     async def initialize(self):
         """

aredis/pipeline.py

@@ -1,19 +1,14 @@
-import asyncio
 import sys
 import typing
 from itertools import chain
 
 from aredis.client import (StrictRedisCluster, StrictRedis)
-from aredis.exceptions import (RedisError,
-                               ConnectionError,
-                               TimeoutError,
-                               ResponseError,
-                               WatchError,
-                               ExecAbortError,
-                               MovedError,
-                               AskError,
-                               TryAgainError,
-                               RedisClusterException)
+from aredis.exceptions import (RedisError, ConnectionError,
+                               TimeoutError, ResponseError,
+                               WatchError, ExecAbortError,
+                               MovedError, AskError,
+                               TryAgainError, RedisClusterException,
+                               ClusterTransactionError)
 from aredis.utils import (dict_merge,
                           clusterdown_wrapper)
 
@@ -276,9 +271,9 @@ async def execute(self, raise_on_error=True):
         if self.scripts:
             await self.load_scripts()
         if self.transaction or self.explicit_transaction:
-            execute = self._execute_transaction
+            exec = self._execute_transaction
         else:
-            execute = self._execute_pipeline
+            exec = self._execute_pipeline
 
         conn = self.connection
         if not conn:
@@ -288,7 +283,7 @@ async def execute(self, raise_on_error=True):
             self.connection = conn
 
         try:
-            return await execute(conn, stack, raise_on_error)
+            return await exec(conn, stack, raise_on_error)
         except (ConnectionError, TimeoutError) as e:
             conn.disconnect()
             if not conn.retry_on_timeout and isinstance(e, TimeoutError):
@@ -303,7 +298,7 @@ async def execute(self, raise_on_error=True):
                                  "one or more keys")
             # otherwise, it's safe to retry since the transaction isn't
             # predicated on any state
-            return await execute(conn, stack, raise_on_error)
+            return await exec(conn, stack, raise_on_error)
         finally:
             await self.reset()
 
@@ -325,32 +320,25 @@ class StrictPipeline(BasePipeline, StrictRedis):
 
 class StrictClusterPipeline(StrictRedisCluster):
     def __init__(self, connection_pool, result_callbacks=None,
-                 response_callbacks=None, startup_nodes=None):
-        """
-        """
+                 response_callbacks=None, startup_nodes=None,
+                 transaction=False, watches=None):
         self.command_stack = []
         self.refresh_table_asap = False
         self.connection_pool = connection_pool
         self.result_callbacks = result_callbacks or self.__class__.RESULT_CALLBACKS.copy()
         self.startup_nodes = startup_nodes if startup_nodes else []
         self.nodes_flags = self.__class__.NODES_FLAGS.copy()
         self.response_callbacks = dict_merge(response_callbacks or self.__class__.RESPONSE_CALLBACKS.copy())
+        self.transaction = transaction
+        self.watches = watches or None
+        self.watching = False
+        self.explicit_transaction = False
 
     def __repr__(self):
         """
         """
         return "{0}".format(type(self).__name__)
 
-    def __enter__(self):
-        """
-        """
-        return self
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        """
-        """
-        self.reset()
-
     def __del__(self):
         """
         """
@@ -361,6 +349,12 @@ def __len__(self):
         """
         return len(self.command_stack)
 
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self.reset()
+
     def _determine_slot(self, *args):
         """
         figure out what slot based on command and args
@@ -407,9 +401,12 @@ async def execute(self, raise_on_error=True):
 
         if not stack:
             return []
-
+        if self.transaction:
+            execute = self.send_cluster_transaction
+        else:
+            execute = self.send_cluster_commands
         try:
-            return await self.send_cluster_commands(stack, raise_on_error)
+            return await execute(stack, raise_on_error)
         finally:
             self.reset()
 
@@ -420,30 +417,61 @@ def reset(self):
         self.command_stack = []
 
         self.scripts = set()
-
-        # TODO: Implement
-        # make sure to reset the connection state in the event that we were
-        # watching something
-        # if self.watching and self.connection:
-        #     try:
-        #         # call this manually since our unwatch or
-        #         # immediate_execute_command methods can call reset()
-        #         self.connection.send_command('UNWATCH')
-        #         self.connection.read_response()
-        #     except ConnectionError:
-        #         # disconnect will also remove any previous WATCHes
-        #         self.connection.disconnect()
-
+        self.watches = []
         # clean up the other instance attributes
         self.watching = False
         self.explicit_transaction = False
 
-        # TODO: Implement
-        # we can safely return the connection to the pool here since we're
-        # sure we're no longer WATCHing anything
-        # if self.connection:
-        #     self.connection_pool.release(self.connection)
-        #     self.connection = None
+    @clusterdown_wrapper
+    async def send_cluster_transaction(self, stack, raise_on_error=True):
+        # the first time sending the commands we send all of the commands that were queued up.
+        # if we have to run through it again, we only retry the commands that failed.
+        attempt = sorted(stack, key=lambda x: x.position)
+        node = {}
+
+        # as we move through each command that still needs to be processed,
+        # we figure out the slot number that command maps to, then from the slot determine the node.
+        for c in attempt:
+            # refer to our internal node -> slot table that tells us where a given
+            # command should route to.
+            slot = self._determine_slot(*c.args)
+            hashed_node = self.connection_pool.get_node_by_slot(slot)
+
+            # now that we know the name of the node ( it's just a string in the form of host:port )
+            # we can build a list of commands for each node.
+            if node.get('name') != hashed_node['name']:
+                # raise error if commands in a transaction can not hash to same node
+                if len(node) > 0:
+                    raise ClusterTransactionError("Keys in request don't hash to the same node")
+            node = hashed_node
+        conn = self.connection_pool.get_connection_by_node(node)
+        if self.watches:
+            await self._watch(node, conn, self.watches)
+        node_commands = NodeCommands(self.parse_response, conn, in_transaction=True)
+        node_commands.append(PipelineCommand(('MULTI',)))
+        node_commands.extend(attempt)
+        node_commands.append(PipelineCommand(('EXEC',)))
+        self.explicit_transaction = True
+        await node_commands.write()
+        # todo: make this place clear
+        try:
+            await node_commands.read()
+        except ExecAbortError:
+            if self.explicit_transaction:
+                await conn.send_command('DISCARD')
+                await conn.read_response()
+
+        # If at least one watched key is modified before the EXEC command,
+        # the whole transaction aborts,
+        # and EXEC returns a Null reply to notify that the transaction failed.
+        if node_commands.commands[-1].result is None:
+            raise WatchError
+        self.connection_pool.release(conn)
+        if self.watching:
+            self._unwatch(conn)
+        if raise_on_error:
+            self.raise_first_error(stack)
+
 
     @clusterdown_wrapper
     async def send_cluster_commands(self, stack, raise_on_error=True, allow_redirections=True):
@@ -550,35 +578,38 @@ def _fail_on_redirect(self, allow_redirections):
         if not allow_redirections:
             raise RedisClusterException("ASK & MOVED redirection not allowed in this pipeline")
 
-    def multi(self):
+    def _multi(self):
         """
         """
         raise RedisClusterException("method multi() is not implemented")
 
     def immediate_execute_command(self, *args, **options):
-        """
-        """
         raise RedisClusterException("method immediate_execute_command() is not implemented")
 
-    def _execute_transaction(self, *args, **kwargs):
-        """
-        """
-        raise RedisClusterException("method _execute_transaction() is not implemented")
-
     def load_scripts(self):
         """
         """
         raise RedisClusterException("method load_scripts() is not implemented")
 
-    def watch(self, *names):
-        """
-        """
-        raise RedisClusterException("method watch() is not implemented")
+    async def _watch(self, node, conn, names):
+        "Watches the values at keys ``names``"
+        for name in names:
+            slot = self._determine_slot('WATCH', name)
+            dist_node = self.connection_pool.get_node_by_slot(slot)
+            if node.get('name') != dist_node['name']:
+                # raise error if commands in a transaction can not hash to same node
+                if len(node) > 0:
+                    raise ClusterTransactionError("Keys in request don't hash to the same node")
+        if self.explicit_transaction:
+            raise RedisError('Cannot issue a WATCH after a MULTI')
+        await conn.send_command('WATCH', *names)
+        return await conn.read_response()
 
-    def unwatch(self):
-        """
-        """
-        raise RedisClusterException("method unwatch() is not implemented")
+    async def _unwatch(self, conn):
+        "Unwatches all previously specified keys"
+        await conn.send_command('UNWATCH')
+        res = await conn.read_response()
+        return self.watching and res or True
 
     def script_load_for_pipeline(self, *args, **kwargs):
         """
@@ -691,12 +722,16 @@ class NodeCommands(object):
     """
     """
 
-    def __init__(self, parse_response, connection):
+    def __init__(self, parse_response, connection, in_transaction=False):
         """
         """
         self.parse_response = parse_response
         self.connection = connection
         self.commands = []
+        self.in_transaction = in_transaction
+
+    def extend(self, c):
+        self.commands.extend(c)
 
     def append(self, c):
         """
@@ -724,8 +759,6 @@ async def write(self):
                 c.result = e
 
     async def read(self):
-        """
-        """
         connection = self.connection
         for c in self.commands:
 
@@ -741,7 +774,13 @@ async def read(self):
             # explicitly open the connection and all will be well.
             if c.result is None:
                 try:
-                    c.result = await self.parse_response(connection, c.args[0], **c.options)
+                    if self.in_transaction:
+                        cmd = '_'
+                    else:
+                        cmd = c.args[0]
+                    c.result = await self.parse_response(connection, cmd, **c.options)
+                except ExecAbortError:
+                    raise
                 except (ConnectionError, TimeoutError) as e:
                     for c in self.commands:
                         c.result = e