Fixing linters and applying review comments

Author: petyaslavova

redis/keyspace_notifications.py

@@ -63,6 +63,7 @@
     RedisError,
     TimeoutError,
 )
+from redis.utils import safe_str
 
 # Type alias for channel arguments - can be a string, bytes, or Channel object
 # This is defined here and the actual types are added after class definitions
@@ -296,7 +297,7 @@ class KeyNotification:
     @classmethod
     def from_message(
         cls,
-        message: Dict[str, Any],
+        message: Optional[Dict[str, Any]],
         key_prefix: Optional[Union[str, bytes]] = None,
     ) -> Optional["KeyNotification"]:
         """
@@ -340,10 +341,8 @@ def from_message(
             return None
 
         # Convert bytes to string if needed
-        if isinstance(channel, bytes):
-            channel = channel.decode("utf-8", errors="replace")
-        if isinstance(data, bytes):
-            data = data.decode("utf-8", errors="replace")
+        channel = safe_str(channel)
+        data = safe_str(data)
 
         return cls._parse(channel, data, key_prefix)
 
@@ -368,10 +367,8 @@ def try_parse(
         Returns:
             A KeyNotification if valid, None otherwise.
         """
-        if isinstance(channel, bytes):
-            channel = channel.decode("utf-8", errors="replace")
-        if isinstance(data, bytes):
-            data = data.decode("utf-8", errors="replace")
+        channel = safe_str(channel)
+        data = safe_str(data)
 
         return cls._parse(channel, data, key_prefix)
 
@@ -384,8 +381,7 @@ def _parse(
     ) -> Optional["KeyNotification"]:
         """Internal parsing logic."""
         # Normalize key_prefix
-        if isinstance(key_prefix, bytes):
-            key_prefix = key_prefix.decode("utf-8", errors="replace")
+        key_prefix = safe_str(key_prefix)
 
         # Try keyspace pattern first: __keyspace@<db>__:<key>
         match = cls._KEYSPACE_PATTERN.match(channel)
@@ -433,8 +429,7 @@ def _parse(
 
     def key_starts_with(self, prefix: Union[str, bytes]) -> bool:
         """Check if the key starts with the given prefix."""
-        if isinstance(prefix, bytes):
-            prefix = prefix.decode("utf-8", errors="replace")
+        prefix = safe_str(prefix)
         return self.key.startswith(prefix)
 
 
@@ -454,16 +449,17 @@ class KeyspaceChannel:
     __str__ to return the channel string.
 
     Attributes:
-        key_or_pattern: The key or pattern to monitor
-        db: The database number (None means all databases)
+        key_or_pattern: The key or pattern to monitor (use '*' for wildcards)
+        db: The database number (defaults to 0, the only database in Redis Cluster)
         is_pattern: Whether this channel contains wildcards
 
     Examples:
         >>> channel = KeyspaceChannel("user:123", db=0)
         >>> str(channel)
         '__keyspace@0__:user:123'
 
-        >>> channel = KeyspaceChannel.pattern("user:", db=0)
+        >>> # Pattern subscription (wildcards are auto-detected)
+        >>> channel = KeyspaceChannel("user:*", db=0)
         >>> str(channel)
         '__keyspace@0__:user:*'
 
@@ -599,15 +595,13 @@ def __hash__(self) -> int:
 
 def is_keyspace_channel(channel: Union[str, bytes]) -> bool:
     """Check if a channel is a keyspace notification channel."""
-    if isinstance(channel, bytes):
-        channel = channel.decode("utf-8", errors="replace")
+    channel = safe_str(channel)
     return channel.startswith(KeyspaceChannel.PREFIX)
 
 
 def is_keyevent_channel(channel: Union[str, bytes]) -> bool:
     """Check if a channel is a keyevent notification channel."""
-    if isinstance(channel, bytes):
-        channel = channel.decode("utf-8", errors="replace")
+    channel = safe_str(channel)
     return channel.startswith(KeyeventChannel.PREFIX)
 
 
@@ -616,7 +610,9 @@ def is_keyspace_notification_channel(channel: Union[str, bytes]) -> bool:
     return is_keyspace_channel(channel) or is_keyevent_channel(channel)
 
 
-def _is_pattern(channel: Union[str, bytes, "KeyspaceChannel", "KeyeventChannel"]) -> bool:
+def _is_pattern(
+    channel: Union[str, bytes, "KeyspaceChannel", "KeyeventChannel"],
+) -> bool:
     """
     Check if a channel string contains glob-style pattern characters.
 
@@ -636,8 +632,7 @@ def _is_pattern(channel: Union[str, bytes, "KeyspaceChannel", "KeyeventChannel"]
     # (KeyspaceChannel, KeyeventChannel)
     if hasattr(channel, "_channel_str"):
         channel = channel._channel_str
-    if isinstance(channel, bytes):
-        channel = channel.decode("utf-8", errors="replace")
+    channel = safe_str(channel)
     # Check for unescaped glob pattern characters
     # We look for *, ?, or [ that are not escaped with backslash
     i = 0
@@ -714,9 +709,19 @@ def __init__(
         self._topology_check_interval: float = 1.0
         self._last_topology_check: float = 0.0
 
-        # Enable keyspace notifications on all nodes if requested
+        # Store the notify-keyspace-events configuration for applying to new nodes
+        self._notify_keyspace_events: Optional[str] = None
         if notify_keyspace_events is not None:
-            self._configure_keyspace_notifications(notify_keyspace_events)
+            # Normalize to string value for storage
+            self._notify_keyspace_events = (
+                notify_keyspace_events.value
+                if isinstance(notify_keyspace_events, NotifyKeyspaceEvents)
+                else notify_keyspace_events
+            )
+
+        # Enable keyspace notifications on all nodes if requested
+        if self._notify_keyspace_events is not None:
+            self._configure_keyspace_notifications(self._notify_keyspace_events)
 
     def _configure_keyspace_notifications(
         self, events: Union[str, NotifyKeyspaceEvents]
@@ -729,17 +734,17 @@ def _configure_keyspace_notifications(
                    (e.g., NotifyKeyspaceEvents.ALL or "KEA")
         """
         # Get the string value (handles both enum and plain strings)
-        events_str = events.value if isinstance(events, NotifyKeyspaceEvents) else events
+        events_str = (
+            events.value if isinstance(events, NotifyKeyspaceEvents) else events
+        )
         for node in self._get_all_primary_nodes():
             redis_conn = self.cluster.get_redis_connection(node)
             redis_conn.config_set("notify-keyspace-events", events_str)
 
     def _get_all_primary_nodes(self):
         """Get all primary nodes in the cluster."""
         return [
-            node
-            for node in self.cluster.get_nodes()
-            if node.server_type == "primary"
+            node for node in self.cluster.get_nodes() if node.server_type == "primary"
         ]
 
     def _ensure_node_pubsub(self, node) -> Any:
@@ -752,9 +757,7 @@ def _ensure_node_pubsub(self, node) -> Any:
             self._node_pubsubs[node.name] = pubsub
         return self._node_pubsubs[node.name]
 
-    def _subscribe_to_all_nodes(
-        self, channels: Dict[str, Any], use_psubscribe: bool
-    ):
+    def _subscribe_to_all_nodes(self, channels: Dict[str, Any], use_psubscribe: bool):
         """Subscribe to patterns/channels on all primary nodes."""
         primaries = self._get_all_primary_nodes()
 
@@ -771,9 +774,7 @@ def _subscribe_to_all_nodes(
             else:
                 pubsub.subscribe(**channels)
 
-    def _unsubscribe_from_all_nodes(
-        self, channels: List[str], use_punsubscribe: bool
-    ):
+    def _unsubscribe_from_all_nodes(self, channels: List[str], use_punsubscribe: bool):
         """Unsubscribe from patterns/channels on all nodes."""
         for pubsub in self._node_pubsubs.values():
             if use_punsubscribe:
@@ -815,6 +816,7 @@ def subscribe(
         # Wrap the handler to convert raw messages to KeyNotification objects
         wrapped_handler = None
         if handler is not None:
+
             def wrapped_handler(message):
                 notification = KeyNotification.from_message(message)
                 if notification is not None:
@@ -824,8 +826,12 @@ def wrapped_handler(message):
         exact_channels = {}
 
         for channel in channels:
-            # Convert Channel objects to strings for use as dict keys
-            channel_str = str(channel) if hasattr(channel, "_channel_str") else channel
+            # Convert Channel objects and bytes to strings for use as dict keys
+            # (keyword arguments in **patterns/**exact_channels must be str)
+            if hasattr(channel, "_channel_str"):
+                channel_str = str(channel)
+            else:
+                channel_str = safe_str(channel)
             if _is_pattern(channel):
                 patterns[channel_str] = wrapped_handler
             else:
@@ -854,8 +860,12 @@ def unsubscribe(self, *channels: ChannelT):
         exact_channels = []
 
         for channel in channels:
-            # Convert Channel objects to strings
-            channel_str = str(channel) if hasattr(channel, "_channel_str") else channel
+            # Convert Channel objects and bytes to strings
+            # (must match the keys used in subscribe())
+            if hasattr(channel, "_channel_str"):
+                channel_str = str(channel)
+            else:
+                channel_str = safe_str(channel)
             if _is_pattern(channel):
                 self._subscribed_patterns.pop(channel_str, None)
                 patterns.append(channel_str)
@@ -934,6 +944,51 @@ def _create_pubsub_iterator(self):
                 return
             yield from pubsubs
 
+    def _poll_all_nodes_once(
+        self, ignore_subscribe_messages: bool
+    ) -> Optional[KeyNotification]:
+        """
+        Perform a single non-blocking poll over all node pubsubs.
+
+        This is used when timeout=0 to match the expected semantics of
+        PubSub.get_message(timeout=0) - a non-blocking check for messages.
+
+        Returns:
+            A KeyNotification if one is available, None otherwise.
+        """
+        # Check for topology changes before polling
+        self._check_topology_changed()
+
+        for pubsub in list(self._node_pubsubs.values()):
+            try:
+                message = pubsub.get_message(
+                    ignore_subscribe_messages=ignore_subscribe_messages,
+                    timeout=0.0,
+                )
+            except (MovedError, AskError):
+                self._refresh_subscriptions_on_error()
+                continue
+            except (ConnectionError, TimeoutError, RedisError):
+                self._refresh_subscriptions_on_error()
+                continue
+
+            if message is not None:
+                notification = KeyNotification.from_message(
+                    message, key_prefix=self.key_prefix
+                )
+                if notification is not None:
+                    # Call handler if registered
+                    handler = self._subscribed_patterns.get(
+                        message.get("pattern")
+                    ) or self._subscribed_channels.get(message.get("channel"))
+                    if handler:
+                        handler(notification)
+                        # Continue polling remaining nodes
+                        continue
+                    return notification
+
+        return None
+
     def get_message(
         self,
         ignore_subscribe_messages: bool = False,
@@ -964,6 +1019,11 @@ def get_message(
         if total_nodes == 0:
             return None
 
+        # Handle timeout=0 as a single non-blocking poll over all pubsubs
+        # This matches the expected semantics of PubSub.get_message(timeout=0)
+        if timeout == 0.0:
+            return self._poll_all_nodes_once(ignore_subscribe_messages)
+
         # Calculate per-node timeout for each poll
         # Use a small timeout per node to allow round-robin polling
         per_node_timeout = min(0.1, timeout / max(total_nodes, 1))
@@ -1116,6 +1176,18 @@ def refresh_subscriptions(self):
         new_nodes = set(current_primaries.keys()) - set(self._node_pubsubs.keys())
         for node_name in new_nodes:
             node = current_primaries[node_name]
+
+            # Configure notify-keyspace-events on new node before subscribing
+            if self._notify_keyspace_events is not None:
+                try:
+                    redis_conn = self.cluster.get_redis_connection(node)
+                    redis_conn.config_set(
+                        "notify-keyspace-events", self._notify_keyspace_events
+                    )
+                except Exception:
+                    # Log or ignore - node may not be ready yet
+                    pass
+
             pubsub = self._ensure_node_pubsub(node)
 
             if self._subscribed_patterns:

specs/keyspace-notifications/SPEC.md

@@ -2,7 +2,7 @@
 
 ## Tasks
 
-1. Understand the introduction, basics, and cluster specifics of keyspache notifications in Redis OSS
+1. Understand the introduction, basics, and cluster specifics of keyspace notifications in Redis OSS
 2. Take the requirements into account
 3. Understand the implementation example based on the referenced PR
 4. Implement a similar abstraction for redis-py
@@ -33,7 +33,7 @@ Redis allows to notify clients whenever a key is modified. The mechanism that is
 By using a cluster client, it should be possible to:
 
 - Consume keyspace and key event notifications across the cluster via a new API for subscribing to `keyspace` and `keyevent` channels
-- Support multiple logical 
+- Support multiple logical databases for standalone clients (via `SELECT <db_index>`), while Redis OSS Cluster continues to use database 0 only
 - Take the differences in behavior between normally published messages and keyspace notifications into account - Keyspace notifications aren't propagated between nodes via the cluster bus
 - Automatically react to topology changes (e.g. node added/removed, slot migration, ...) by resubscribing to the relevant channels
 
@@ -49,6 +49,6 @@ Please take a look at https://github.com/StackExchange/StackExchange.Redis/pull/
 - Updated a key on node 1 and verify that the keyspace notification is received
 - Update a key on node 2 and verify that the key event notification is received
 - Delete a key on node 3 and verify that the deletion notification is received
-- Modify a bunch of keys across nodes 1, 2 and 3 that all match the sane pattern and check that all notifications are received
+- Modify a bunch of keys across nodes 1, 2 and 3 that all match the same pattern and check that all notifications are received
 - Move a bunch of slots from node 1 to node 2 that impact some pre-defined keys and ensure that keyspace notifications and key event notifications are received correctly before and after the migration