Hitless-Upgrade: Add handling of MOVING push notification with "null" host:port info.

redis/_parsers/base.py

@@ -1,8 +1,17 @@
+import logging
 import sys
 from abc import ABC
 from asyncio import IncompleteReadError, StreamReader, TimeoutError
 from typing import Callable, List, Optional, Protocol, Union
 
+from redis.maintenance_events import (
+    NodeFailedOverEvent,
+    NodeFailingOverEvent,
+    NodeMigratedEvent,
+    NodeMigratingEvent,
+    NodeMovingEvent,
+)
+
 if sys.version_info.major >= 3 and sys.version_info.minor >= 11:
     from asyncio import timeout as async_timeout
 else:
@@ -50,6 +59,8 @@
     "Client sent AUTH, but no password is set": AuthenticationError,
 }
 
+logger = logging.getLogger(__name__)
+
 
 class BaseParser(ABC):
     EXCEPTION_CLASSES = {
@@ -158,48 +169,163 @@ async def read_response(
         raise NotImplementedError()
 
 
-_INVALIDATION_MESSAGE = [b"invalidate", "invalidate"]
+_INVALIDATION_MESSAGE = (b"invalidate", "invalidate")
+_MOVING_MESSAGE = (b"MOVING", "MOVING")
+_MIGRATING_MESSAGE = (b"MIGRATING", "MIGRATING")
+_MIGRATED_MESSAGE = (b"MIGRATED", "MIGRATED")
+_FAILING_OVER_MESSAGE = (b"FAILING_OVER", "FAILING_OVER")
+_FAILED_OVER_MESSAGE = (b"FAILED_OVER", "FAILED_OVER")
+
+_MAINTENANCE_MESSAGES = (
+    *_MIGRATING_MESSAGE,
+    *_MIGRATED_MESSAGE,
+    *_FAILING_OVER_MESSAGE,
+    *_FAILED_OVER_MESSAGE,
+)
 
 
 class PushNotificationsParser(Protocol):
     """Protocol defining RESP3-specific parsing functionality"""
 
     pubsub_push_handler_func: Callable
     invalidation_push_handler_func: Optional[Callable] = None
+    node_moving_push_handler_func: Optional[Callable] = None
+    maintenance_push_handler_func: Optional[Callable] = None
 
     def handle_pubsub_push_response(self, response):
         """Handle pubsub push responses"""
         raise NotImplementedError()
 
     def handle_push_response(self, response, **kwargs):
-        if response[0] not in _INVALIDATION_MESSAGE:
+        msg_type = response[0]
+        if msg_type not in (
+            *_INVALIDATION_MESSAGE,
+            *_MAINTENANCE_MESSAGES,
+            *_MOVING_MESSAGE,
+        ):
             return self.pubsub_push_handler_func(response)
-        if self.invalidation_push_handler_func:
-            return self.invalidation_push_handler_func(response)
+
+        try:
+            if (
+                msg_type in _INVALIDATION_MESSAGE
+                and self.invalidation_push_handler_func
+            ):
+                return self.invalidation_push_handler_func(response)
+            if msg_type in _MOVING_MESSAGE and self.node_moving_push_handler_func:
+                # Expected message format is: MOVING <seq_number> <time> <endpoint>
+                id = response[1]
+                ttl = response[2]
+                host, port = response[3].decode().split(":")
+                notification = NodeMovingEvent(id, host, port, ttl)
+                return self.node_moving_push_handler_func(notification)
+
+            if msg_type in _MAINTENANCE_MESSAGES and self.maintenance_push_handler_func:
+                notification = None
+
+                if msg_type in _MIGRATING_MESSAGE:
+                    # Expected message format is: MIGRATING <seq_number> <time> <shard_id-s>
+                    id = response[1]
+                    ttl = response[2]
+                    notification = NodeMigratingEvent(id, ttl)
+                elif msg_type in _MIGRATED_MESSAGE:
+                    id = response[1]
+                    notification = NodeMigratedEvent(id)
+                elif msg_type in _FAILING_OVER_MESSAGE:
+                    id = response[1]
+                    ttl = response[2]
+                    notification = NodeFailingOverEvent(id, ttl)
+                elif msg_type in _FAILED_OVER_MESSAGE:
+                    id = response[1]
+                    notification = NodeFailedOverEvent(id)
+
+                if notification is not None:
+                    return self.maintenance_push_handler_func(notification)
+        except Exception as e:
+            logger.error(
+                "Error handling {} message ({}): {}".format(msg_type, response, e)
+            )
+
+        return None
 
     def set_pubsub_push_handler(self, pubsub_push_handler_func):
         self.pubsub_push_handler_func = pubsub_push_handler_func
 
     def set_invalidation_push_handler(self, invalidation_push_handler_func):
         self.invalidation_push_handler_func = invalidation_push_handler_func
 
+    def set_node_moving_push_handler(self, node_moving_push_handler_func):
+        self.node_moving_push_handler_func = node_moving_push_handler_func
+
+    def set_maintenance_push_handler(self, maintenance_push_handler_func):
+        self.maintenance_push_handler_func = maintenance_push_handler_func
+
 
 class AsyncPushNotificationsParser(Protocol):
     """Protocol defining async RESP3-specific parsing functionality"""
 
     pubsub_push_handler_func: Callable
     invalidation_push_handler_func: Optional[Callable] = None
+    node_moving_push_handler_func: Optional[Callable] = None
+    maintenance_push_handler_func: Optional[Callable] = None
 
     async def handle_pubsub_push_response(self, response):
         """Handle pubsub push responses asynchronously"""
         raise NotImplementedError()
 
     async def handle_push_response(self, response, **kwargs):
         """Handle push responses asynchronously"""
-        if response[0] not in _INVALIDATION_MESSAGE:
+
+        msg_type = response[0]
+        if msg_type not in (
+            *_INVALIDATION_MESSAGE,
+            *_MAINTENANCE_MESSAGES,
+            *_MOVING_MESSAGE,
+        ):
             return await self.pubsub_push_handler_func(response)
-        if self.invalidation_push_handler_func:
-            return await self.invalidation_push_handler_func(response)
+
+        try:
+            if (
+                msg_type in _INVALIDATION_MESSAGE
+                and self.invalidation_push_handler_func
+            ):
+                return await self.invalidation_push_handler_func(response)
+            if msg_type in _MOVING_MESSAGE and self.node_moving_push_handler_func:
+                # push notification from enterprise cluster for node moving
+                id = response[1]
+                ttl = response[2]
+                if response[3] != "null":
+                    host, port = response[3].split(":")
+                else:
+                    host, port = None, None
+                notification = NodeMovingEvent(id, host, port, ttl)
+                return await self.node_moving_push_handler_func(notification)
+
+            if msg_type in _MAINTENANCE_MESSAGES and self.maintenance_push_handler_func:
+                notification = None
+
+                if msg_type in _MIGRATING_MESSAGE:
+                    id = response[1]
+                    ttl = response[2]
+                    notification = NodeMigratingEvent(id, ttl)
+                elif msg_type in _MIGRATED_MESSAGE:
+                    id = response[1]
+                    notification = NodeMigratedEvent(id)
+                elif msg_type in _FAILING_OVER_MESSAGE:
+                    id = response[1]
+                    ttl = response[2]
+                    notification = NodeFailingOverEvent(id, ttl)
+                elif msg_type in _FAILED_OVER_MESSAGE:
+                    id = response[1]
+                    notification = NodeFailedOverEvent(id)
+
+                if notification is not None:
+                    return await self.maintenance_push_handler_func(notification)
+        except Exception as e:
+            logger.error(
+                "Error handling {} message ({}): {}".format(msg_type, response, e)
+            )
+
+        return None
 
     def set_pubsub_push_handler(self, pubsub_push_handler_func):
         """Set the pubsub push handler function"""
@@ -209,6 +335,12 @@ def set_invalidation_push_handler(self, invalidation_push_handler_func):
         """Set the invalidation push handler function"""
         self.invalidation_push_handler_func = invalidation_push_handler_func
 
+    def set_node_moving_push_handler(self, node_moving_push_handler_func):
+        self.node_moving_push_handler_func = node_moving_push_handler_func
+
+    def set_maintenance_push_handler(self, maintenance_push_handler_func):
+        self.maintenance_push_handler_func = maintenance_push_handler_func
+
 
 class _AsyncRESPBase(AsyncBaseParser):
     """Base class for async resp parsing"""

redis/_parsers/hiredis.py

@@ -47,6 +47,8 @@ def __init__(self, socket_read_size):
         self.socket_read_size = socket_read_size
         self._buffer = bytearray(socket_read_size)
         self.pubsub_push_handler_func = self.handle_pubsub_push_response
+        self.node_moving_push_handler_func = None
+        self.maintenance_push_handler_func = None
         self.invalidation_push_handler_func = None
         self._hiredis_PushNotificationType = None
 
@@ -141,12 +143,15 @@ def read_response(self, disable_decoding=False, push_request=False):
                 response, self._hiredis_PushNotificationType
             ):
                 response = self.handle_push_response(response)
-                if not push_request:
-                    return self.read_response(
-                        disable_decoding=disable_decoding, push_request=push_request
-                    )
-                else:
+
+                # if this is a push request return the push response
+                if push_request:
                     return response
+
+                return self.read_response(
+                    disable_decoding=disable_decoding,
+                    push_request=push_request,
+                )
             return response
 
         if disable_decoding:
@@ -169,12 +174,13 @@ def read_response(self, disable_decoding=False, push_request=False):
             response, self._hiredis_PushNotificationType
         ):
             response = self.handle_push_response(response)
-            if not push_request:
-                return self.read_response(
-                    disable_decoding=disable_decoding, push_request=push_request
-                )
-            else:
+            if push_request:
                 return response
+            return self.read_response(
+                disable_decoding=disable_decoding,
+                push_request=push_request,
+            )
+
         elif (
             isinstance(response, list)
             and response

redis/_parsers/resp3.py

@@ -18,6 +18,8 @@ class _RESP3Parser(_RESPBase, PushNotificationsParser):
     def __init__(self, socket_read_size):
         super().__init__(socket_read_size)
         self.pubsub_push_handler_func = self.handle_pubsub_push_response
+        self.node_moving_push_handler_func = None
+        self.maintenance_push_handler_func = None
         self.invalidation_push_handler_func = None
 
     def handle_pubsub_push_response(self, response):
@@ -117,17 +119,21 @@ def _read_response(self, disable_decoding=False, push_request=False):
                 for _ in range(int(response))
             ]
             response = self.handle_push_response(response)
-            if not push_request:
-                return self._read_response(
-                    disable_decoding=disable_decoding, push_request=push_request
-                )
-            else:
+
+            # if this is a push request return the push response
+            if push_request:
                 return response
+
+            return self._read_response(
+                disable_decoding=disable_decoding,
+                push_request=push_request,
+            )
         else:
             raise InvalidResponse(f"Protocol Error: {raw!r}")
 
         if isinstance(response, bytes) and disable_decoding is False:
             response = self.encoder.decode(response)
+
         return response