Added support for cluster client

Author: vladvildanov

redis/cluster.py

@@ -29,7 +29,7 @@
     AfterPooledConnectionsInstantiationEvent,
     AfterPubSubConnectionInstantiationEvent,
     ClientType,
-    EventDispatcher,
+    EventDispatcher, AfterCommandExecutionEvent,
 )
 from redis.exceptions import (
     AskError,
@@ -984,6 +984,7 @@ def pipeline(self, transaction=None, shard_hint=None):
             retry=self.retry,
             lock=self._lock,
             transaction=transaction,
+            event_dispatcher=self._event_dispatcher
         )
 
     def lock(
@@ -1367,6 +1368,9 @@ def _execute_command(self, target_node, *args, **kwargs):
         moved = False
         ttl = int(self.RedisClusterRequestTTL)
 
+        # Start timing for observability
+        start_time = time.monotonic()
+
         while ttl > 0:
             ttl -= 1
             try:
@@ -1401,14 +1405,32 @@ def _execute_command(self, target_node, *args, **kwargs):
                     response = self.cluster_response_callbacks[command](
                         response, **kwargs
                     )
+
+                self._emit_after_command_execution_event(
+                    command_name=command,
+                    duration_seconds=time.monotonic() - start_time,
+                    connection=connection,
+                )
                 return response
-            except AuthenticationError:
+            except AuthenticationError as e:
+                self._emit_after_command_execution_event(
+                    command_name=command,
+                    duration_seconds=time.monotonic() - start_time,
+                    connection=connection,
+                    error=e,
+                )
                 raise
-            except MaxConnectionsError:
+            except MaxConnectionsError as e:
                 # MaxConnectionsError indicates client-side resource exhaustion
                 # (too many connections in the pool), not a node failure.
                 # Don't treat this as a node failure - just re-raise the error
                 # without reinitializing the cluster.
+                self._emit_after_command_execution_event(
+                    command_name=command,
+                    duration_seconds=time.monotonic() - start_time,
+                    connection=connection,
+                    error=e,
+                )
                 raise
             except (ConnectionError, TimeoutError) as e:
                 # ConnectionError can also be raised if we couldn't get a
@@ -1423,6 +1445,12 @@ def _execute_command(self, target_node, *args, **kwargs):
                 # Reset the cluster node's connection
                 target_node.redis_connection = None
                 self.nodes_manager.initialize()
+                self._emit_after_command_execution_event(
+                    command_name=command,
+                    duration_seconds=time.monotonic() - start_time,
+                    connection=connection,
+                    error=e,
+                )
                 raise e
             except MovedError as e:
                 # First, we will try to patch the slots/nodes cache with the
@@ -1441,13 +1469,33 @@ def _execute_command(self, target_node, *args, **kwargs):
                 else:
                     self.nodes_manager.update_moved_exception(e)
                 moved = True
-            except TryAgainError:
+                self._emit_after_command_execution_event(
+                    command_name=command,
+                    duration_seconds=time.monotonic() - start_time,
+                    connection=connection,
+                    error=e,
+                )
+            except TryAgainError as e:
                 if ttl < self.RedisClusterRequestTTL / 2:
                     time.sleep(0.05)
+
+                self._emit_after_command_execution_event(
+                    command_name=command,
+                    duration_seconds=time.monotonic() - start_time,
+                    connection=connection,
+                    error=e,
+                )
             except AskError as e:
                 redirect_addr = get_node_name(host=e.host, port=e.port)
                 asking = True
-            except (ClusterDownError, SlotNotCoveredError):
+
+                self._emit_after_command_execution_event(
+                    command_name=command,
+                    duration_seconds=time.monotonic() - start_time,
+                    connection=connection,
+                    error=e,
+                )
+            except (ClusterDownError, SlotNotCoveredError) as e:
                 # ClusterDownError can occur during a failover and to get
                 # self-healed, we will try to reinitialize the cluster layout
                 # and retry executing the command
@@ -1459,19 +1507,60 @@ def _execute_command(self, target_node, *args, **kwargs):
 
                 time.sleep(0.25)
                 self.nodes_manager.initialize()
+
+                self._emit_after_command_execution_event(
+                    command_name=command,
+                    duration_seconds=time.monotonic() - start_time,
+                    connection=connection,
+                    error=e,
+                )
                 raise
-            except ResponseError:
+            except ResponseError as e:
+                self._emit_after_command_execution_event(
+                    command_name=command,
+                    duration_seconds=time.monotonic() - start_time,
+                    connection=connection,
+                    error=e,
+                )
                 raise
             except Exception as e:
                 if connection:
                     connection.disconnect()
+
+                self._emit_after_command_execution_event(
+                    command_name=command,
+                    duration_seconds=time.monotonic() - start_time,
+                    connection=connection,
+                    error=e,
+                )
                 raise e
             finally:
                 if connection is not None:
                     redis_node.connection_pool.release(connection)
 
         raise ClusterError("TTL exhausted.")
 
+    def _emit_after_command_execution_event(
+            self,
+            command_name: str,
+            duration_seconds: float,
+            connection: Connection,
+            error=None
+    ):
+        """
+        Triggers AfterCommandExecutionEvent emit.
+        """
+        self._event_dispatcher.dispatch(
+            AfterCommandExecutionEvent(
+                command_name=command_name,
+                duration_seconds=duration_seconds,
+                server_address=connection.host,
+                server_port=connection.port,
+                db_namespace=str(connection.db),
+                error=error,
+            )
+        )
+
     def close(self) -> None:
         try:
             with self._lock:
@@ -2326,6 +2415,7 @@ def __init__(
         lock=None,
         transaction=False,
         policy_resolver: PolicyResolver = StaticPolicyResolver(),
+        event_dispatcher: Optional["EventDispatcher"] = None,
         **kwargs,
     ):
         """ """
@@ -2395,6 +2485,11 @@ def __init__(
 
         self._policy_resolver = policy_resolver
 
+        if event_dispatcher is None:
+            self._event_dispatcher = EventDispatcher()
+        else:
+            self._event_dispatcher = event_dispatcher
+
     def __repr__(self):
         """ """
         return f"{type(self).__name__}"
@@ -2889,14 +2984,24 @@ def __init__(self, pipe: ClusterPipeline):
     def execute_command(self, *args, **kwargs):
         return self.pipeline_execute_command(*args, **kwargs)
 
-    def _raise_first_error(self, stack):
+    def _raise_first_error(self, stack, start_time):
         """
         Raise the first exception on the stack
         """
         for c in stack:
             r = c.result
             if isinstance(r, Exception):
                 self.annotate_exception(r, c.position + 1, c.args)
+
+                self._pipe._event_dispatcher.dispatch(
+                    AfterCommandExecutionEvent(
+                        command_name="PIPELINE",
+                        duration_seconds=time.monotonic() - start_time,
+                        batch_size=len(stack),
+                        error=r,
+                    )
+                )
+
                 raise r
 
     def execute(self, raise_on_error: bool = True) -> List[Any]:
@@ -3076,13 +3181,28 @@ def _send_cluster_commands(
         # so that we can read them from different sockets as they come back.
         # we dont' multiplex on the sockets as they come available,
         # but that shouldn't make too much difference.
+
+        # Start timing for observability
+        start_time = time.monotonic()
+
         try:
             node_commands = nodes.values()
             for n in node_commands:
                 n.write()
 
             for n in node_commands:
                 n.read()
+
+                self._pipe._event_dispatcher.dispatch(
+                    AfterCommandExecutionEvent(
+                        command_name="PIPELINE",
+                        duration_seconds=time.monotonic() - start_time,
+                        server_address=n.connection.host,
+                        server_port=n.connection.port,
+                        db_namespace=str(n.connection.db),
+                        batch_size=len(n.commands),
+                    )
+                )
         finally:
             # release all of the redis connections we allocated earlier
             # back into the connection pool.
@@ -3168,7 +3288,7 @@ def _send_cluster_commands(
             response.append(c.result)
 
         if raise_on_error:
-            self._raise_first_error(stack)
+            self._raise_first_error(stack, start_time)
 
         return response
 
@@ -3373,9 +3493,38 @@ def _immediate_execute_command(self, *args, **options):
 
     def _get_connection_and_send_command(self, *args, **options):
         redis_node, connection = self._get_client_and_connection_for_transaction()
-        return self._send_command_parse_response(
-            connection, redis_node, args[0], *args, **options
-        )
+
+        # Start timing for observability
+        start_time = time.monotonic()
+
+        try:
+            response = self._send_command_parse_response(
+                connection, redis_node, args[0], *args, **options
+            )
+
+            self._pipe._event_dispatcher.dispatch(
+                AfterCommandExecutionEvent(
+                    command_name=args[0],
+                    duration_seconds=time.monotonic() - start_time,
+                    server_address=connection.host,
+                    server_port=connection.port,
+                    db_namespace=str(connection.db),
+                )
+            )
+
+            return response
+        except Exception as e:
+            self._pipe._event_dispatcher.dispatch(
+                AfterCommandExecutionEvent(
+                    command_name=args[0],
+                    duration_seconds=time.monotonic() - start_time,
+                    server_address=connection.host,
+                    server_port=connection.port,
+                    db_namespace=str(connection.db),
+                    error=e,
+                )
+            )
+            raise
 
     def _send_command_parse_response(
         self, conn, redis_node: Redis, command_name, *args, **options
@@ -3413,13 +3562,24 @@ def _reinitialize_on_error(self, error):
 
         self._executing = False
 
-    def _raise_first_error(self, responses, stack):
+    def _raise_first_error(self, responses, stack, start_time):
         """
         Raise the first exception on the stack
         """
         for r, cmd in zip(responses, stack):
             if isinstance(r, Exception):
                 self.annotate_exception(r, cmd.position + 1, cmd.args)
+
+                self._pipe._event_dispatcher.dispatch(
+                    AfterCommandExecutionEvent(
+                        command_name='TRANSACTION',
+                        duration_seconds=time.monotonic() - start_time,
+                        server_address=self._transaction_connection.host,
+                        server_port=self._transaction_connection.port,
+                        db_namespace=str(self._transaction_connection.db),
+                    )
+                )
+
                 raise r
 
     def execute(self, raise_on_error: bool = True) -> List[Any]:
@@ -3456,6 +3616,10 @@ def _execute_transaction(
         )
         commands = [c.args for c in stack if EMPTY_RESPONSE not in c.options]
         packed_commands = connection.pack_commands(commands)
+
+        # Start timing for observability
+        start_time = time.monotonic()
+
         connection.send_packed_command(packed_commands)
         errors = []
 
@@ -3500,6 +3664,17 @@ def _execute_transaction(
 
         self._executing = False
 
+        self._pipe._event_dispatcher.dispatch(
+            AfterCommandExecutionEvent(
+                command_name='TRANSACTION',
+                duration_seconds=time.monotonic() - start_time,
+                server_address=connection.host,
+                server_port=connection.port,
+                db_namespace=str(connection.db),
+                batch_size=len(self._command_queue),
+            )
+        )
+
         # EXEC clears any watched keys
         self._watching = False
 
@@ -3523,6 +3698,7 @@ def _execute_transaction(
             self._raise_first_error(
                 response,
                 self._command_queue,
+                start_time,
             )
 
         # We have to run response callbacks manually

redis/event.py

@@ -311,6 +311,7 @@ class AfterCommandExecutionEvent:
     error: Optional[Exception] = None
     is_blocking: Optional[bool] = None
     batch_size: Optional[int] = None
+    retry_attempts: Optional[int] = None
 
 class AsyncOnCommandsFailEvent(OnCommandsFailEvent):
     pass
@@ -497,4 +498,5 @@ def listen(self, event: AfterCommandExecutionEvent):
             error=event.error,
             is_blocking=event.is_blocking,
             batch_size=event.batch_size,
+            retry_attempts=event.retry_attempts,
         )

redis/observability/recorder.py

@@ -39,6 +39,7 @@ def record_operation_duration(
         error: Optional[Exception] = None,
         is_blocking: Optional[bool] = None,
         batch_size: Optional[int] = None,
+        retry_attempts: Optional[int] = None,
 ) -> None:
     """
     Record a Redis command execution duration.
@@ -55,6 +56,7 @@ def record_operation_duration(
         error: Exception if command failed, None if successful
         is_blocking: Whether the operation is a blocking command
         batch_size: Number of commands in batch (for pipelines/transactions)
+        retry_attempts: Number of retry attempts made
 
     Example:
         >>> start = time.monotonic()
@@ -89,6 +91,7 @@ def record_operation_duration(
         network_peer_port=server_port,
         is_blocking=is_blocking,
         batch_size=batch_size,
+        retry_attempts=retry_attempts,
     )
     # except Exception:
     #     # Don't let metric recording errors break Redis operations