Adding oss maint notifications handler configurations to parsers. Placeholder for smigrated handler in OSSMaintNotificationsHandler class

Author: petyaslavova

redis/_parsers/base.py

@@ -11,6 +11,8 @@
     NodeMigratedNotification,
     NodeMigratingNotification,
     NodeMovingNotification,
+    OSSNodeMigratedNotification,
+    OSSNodeMigratingNotification,
 )
 
 if sys.version_info.major >= 3 and sys.version_info.minor >= 11:
@@ -179,6 +181,34 @@ async def read_response(
 class MaintenanceNotificationsParser:
     """Protocol defining maintenance push notification parsing functionality"""
 
+    @staticmethod
+    def parse_oss_maintenance_start_msg(response):
+        # TODO This format is not the final - will be changed later
+        # Expected message format is: SMIGRATING <seq_number> <src_node> <dest_node> <slots>
+        id = response[1]
+
+        address_value = response[3]
+        if isinstance(address_value, bytes):
+            address_value = address_value.decode()
+        if response[2] in ("TO", b"TO"):
+            dest_node = address_value
+            src_node = None
+        else:
+            dest_node = None
+            src_node = address_value
+
+        slots = response[4]
+        return OSSNodeMigratingNotification(id, src_node, dest_node, slots)
+
+    @staticmethod
+    def parse_oss_maintenance_completed_msg(response):
+        # TODO This format is not the final - will be changed later
+        # Expected message format is: SMIGRATED <seq_number> <node_address> <slots>
+        id = response[1]
+        node_address = response[2]
+        slots = response[3]
+        return OSSNodeMigratedNotification(id, node_address, slots)
+
     @staticmethod
     def parse_maintenance_start_msg(response, notification_type):
         # Expected message format is: <notification_type> <seq_number> <time>
@@ -215,12 +245,15 @@ def parse_moving_msg(response):
 _MIGRATED_MESSAGE = "MIGRATED"
 _FAILING_OVER_MESSAGE = "FAILING_OVER"
 _FAILED_OVER_MESSAGE = "FAILED_OVER"
+_SMIGRATING_MESSAGE = "SMIGRATING"
+_SMIGRATED_MESSAGE = "SMIGRATED"
 
 _MAINTENANCE_MESSAGES = (
     _MIGRATING_MESSAGE,
     _MIGRATED_MESSAGE,
     _FAILING_OVER_MESSAGE,
     _FAILED_OVER_MESSAGE,
+    _SMIGRATING_MESSAGE,
 )
 
 MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING: dict[
@@ -246,6 +279,14 @@ def parse_moving_msg(response):
         NodeMovingNotification,
         MaintenanceNotificationsParser.parse_moving_msg,
     ),
+    _SMIGRATING_MESSAGE: (
+        OSSNodeMigratingNotification,
+        MaintenanceNotificationsParser.parse_oss_maintenance_start_msg,
+    ),
+    _SMIGRATED_MESSAGE: (
+        OSSNodeMigratedNotification,
+        MaintenanceNotificationsParser.parse_oss_maintenance_completed_msg,
+    ),
 }
 
 
@@ -256,6 +297,7 @@ class PushNotificationsParser(Protocol):
     invalidation_push_handler_func: Optional[Callable] = None
     node_moving_push_handler_func: Optional[Callable] = None
     maintenance_push_handler_func: Optional[Callable] = None
+    oss_maintenance_push_handler_func: Optional[Callable] = None
 
     def handle_pubsub_push_response(self, response):
         """Handle pubsub push responses"""
@@ -270,6 +312,7 @@ def handle_push_response(self, response, **kwargs):
             _INVALIDATION_MESSAGE,
             *_MAINTENANCE_MESSAGES,
             _MOVING_MESSAGE,
+            _SMIGRATED_MESSAGE,
         ):
             return self.pubsub_push_handler_func(response)
 
@@ -292,13 +335,27 @@ def handle_push_response(self, response, **kwargs):
                 parser_function = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING[
                     msg_type
                 ][1]
-                notification_type = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING[
-                    msg_type
-                ][0]
-                notification = parser_function(response, notification_type)
+                if msg_type == _SMIGRATING_MESSAGE:
+                    notification = parser_function(response)
+                else:
+                    notification_type = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING[
+                        msg_type
+                    ][0]
+                    notification = parser_function(response, notification_type)
 
                 if notification is not None:
                     return self.maintenance_push_handler_func(notification)
+            if (
+                msg_type == _SMIGRATED_MESSAGE
+                and self.oss_maintenance_push_handler_func
+            ):
+                parser_function = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING[
+                    msg_type
+                ][1]
+                notification = parser_function(response)
+
+                if notification is not None:
+                    return self.oss_maintenance_push_handler_func(notification)
         except Exception as e:
             logger.error(
                 "Error handling {} message ({}): {}".format(msg_type, response, e)
@@ -318,6 +375,9 @@ def set_node_moving_push_handler(self, node_moving_push_handler_func):
     def set_maintenance_push_handler(self, maintenance_push_handler_func):
         self.maintenance_push_handler_func = maintenance_push_handler_func
 
+    def set_oss_maintenance_push_handler(self, oss_maintenance_push_handler_func):
+        self.oss_maintenance_push_handler_func = oss_maintenance_push_handler_func
+
 
 class AsyncPushNotificationsParser(Protocol):
     """Protocol defining async RESP3-specific parsing functionality"""
@@ -326,6 +386,7 @@ class AsyncPushNotificationsParser(Protocol):
     invalidation_push_handler_func: Optional[Callable] = None
     node_moving_push_handler_func: Optional[Callable[..., Awaitable[None]]] = None
     maintenance_push_handler_func: Optional[Callable[..., Awaitable[None]]] = None
+    oss_maintenance_push_handler_func: Optional[Callable[..., Awaitable[None]]] = None
 
     async def handle_pubsub_push_response(self, response):
         """Handle pubsub push responses asynchronously"""
@@ -342,6 +403,7 @@ async def handle_push_response(self, response, **kwargs):
             _INVALIDATION_MESSAGE,
             *_MAINTENANCE_MESSAGES,
             _MOVING_MESSAGE,
+            _SMIGRATED_MESSAGE,
         ):
             return await self.pubsub_push_handler_func(response)
 
@@ -366,13 +428,28 @@ async def handle_push_response(self, response, **kwargs):
                 parser_function = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING[
                     msg_type
                 ][1]
-                notification_type = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING[
-                    msg_type
-                ][0]
-                notification = parser_function(response, notification_type)
+                if msg_type == _SMIGRATING_MESSAGE:
+                    notification = parser_function(response)
+                else:
+                    notification_type = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING[
+                        msg_type
+                    ][0]
+                    notification = parser_function(response, notification_type)
 
                 if notification is not None:
                     return await self.maintenance_push_handler_func(notification)
+                if (
+                    msg_type == _SMIGRATED_MESSAGE
+                    and self.oss_maintenance_push_handler_func
+                ):
+                    parser_function = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING[
+                        msg_type
+                    ][1]
+                    notification = parser_function(response)
+                    if notification is not None:
+                        return await self.oss_maintenance_push_handler_func(
+                            notification
+                        )
         except Exception as e:
             logger.error(
                 "Error handling {} message ({}): {}".format(msg_type, response, e)
@@ -394,6 +471,9 @@ def set_node_moving_push_handler(self, node_moving_push_handler_func):
     def set_maintenance_push_handler(self, maintenance_push_handler_func):
         self.maintenance_push_handler_func = maintenance_push_handler_func
 
+    def set_oss_maintenance_push_handler(self, oss_maintenance_push_handler_func):
+        self.oss_maintenance_push_handler_func = oss_maintenance_push_handler_func
+
 
 class _AsyncRESPBase(AsyncBaseParser):
     """Base class for async resp parsing"""

redis/_parsers/hiredis.py

@@ -49,6 +49,7 @@ def __init__(self, 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.oss_maintenance_push_handler_func = None
         self.invalidation_push_handler_func = None
         self._hiredis_PushNotificationType = None
 

redis/_parsers/resp3.py

@@ -20,6 +20,7 @@ def __init__(self, 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.oss_maintenance_push_handler_func = None
         self.invalidation_push_handler_func = None
 
     def handle_pubsub_push_response(self, response):

redis/maint_notifications.py

@@ -9,6 +9,9 @@
 
 from redis.typing import Number
 
+if TYPE_CHECKING:
+    from redis.cluster import NodesManager
+
 
 class MaintenanceState(enum.Enum):
     NONE = "none"
@@ -886,6 +889,7 @@ class MaintNotificationsConnectionHandler:
     _NOTIFICATION_TYPES: dict[type["MaintenanceNotification"], int] = {
         NodeMigratingNotification: 1,
         NodeFailingOverNotification: 1,
+        OSSNodeMigratingNotification: 1,
         NodeMigratedNotification: 0,
         NodeFailedOverNotification: 0,
     }
@@ -939,3 +943,31 @@ def handle_maintenance_completed_notification(self):
         # timeouts by providing -1 as the relaxed timeout
         self.connection.update_current_socket_timeout(-1)
         self.connection.maintenance_state = MaintenanceState.NONE
+
+
+class OSSMaintNotificationsHandler:
+    def __init__(
+        self, nodes_manager: "NodesManager", config: MaintNotificationsConfig
+    ) -> None:
+        self.nodes_manager = nodes_manager
+        self.config = config
+        self._processed_notifications = set()
+        self._lock = threading.RLock()
+
+    def remove_expired_notifications(self):
+        with self._lock:
+            for notification in tuple(self._processed_notifications):
+                if notification.is_expired():
+                    self._processed_notifications.remove(notification)
+
+    def handle_notification(self, notification: MaintenanceNotification):
+        if isinstance(notification, OSSNodeMigratedNotification):
+            self.handle_oss_maintenance_completed_notification(notification)
+        else:
+            logging.error(f"Unhandled notification type: {notification}")
+
+    def handle_oss_maintenance_completed_notification(
+        self, notification: OSSNodeMigratedNotification
+    ):
+        self.remove_expired_notifications()
+        logging.info(f"Received OSS maintenance completed notification: {notification}")

tests/maint_notifications/proxy_server_helpers.py

@@ -1,5 +1,6 @@
 import base64
 import logging
+import re
 from typing import Union
 
 from redis.http.http_client import HttpClient, HttpError
@@ -10,6 +11,19 @@
 class RespTranslator:
     """Helper class to translate between RESP and other encodings."""
 
+    @staticmethod
+    def str_or_list_to_resp(txt: str) -> str:
+        """
+        Convert specific string or list to RESP format.
+        """
+        if re.match(r"^<.*>$", txt):
+            items = txt[1:-1].split(",")
+            return f"*{len(items)}\r\n" + "\r\n".join(
+                f"${len(x)}\r\n{x}" for x in items
+            )
+        else:
+            return f"${len(txt)}\r\n{txt}"
+
     @staticmethod
     def cluster_slots_to_resp(resp: str) -> str:
         """Convert query to RESP format."""
@@ -24,7 +38,9 @@ def smigrating_to_resp(resp: str) -> str:
         """Convert query to RESP format."""
         return (
             f">{len(resp.split())}\r\n"
-            + "\r\n".join(f"${len(x)}\r\n{x}" for x in resp.split())
+            + "\r\n".join(
+                f"{RespTranslator.str_or_list_to_resp(x)}" for x in resp.split()
+            )
             + "\r\n"
         )
 
@@ -118,6 +134,8 @@ def get_stats(self) -> dict:
 
         try:
             response = self.http_client.get(url)
+            if isinstance(response, dict):
+                return response
             return response.json()
 
         except HttpError as e:
@@ -134,6 +152,8 @@ def get_connections(self) -> dict:
 
         try:
             response = self.http_client.get(url)
+            if isinstance(response, dict):
+                return response
             return response.json()
         except HttpError as e:
             raise RuntimeError(f"Failed to get connections: {e}")
@@ -192,7 +212,9 @@ def send_notification(
         data = base64.b64encode(notification.encode("utf-8"))
 
         try:
-            response = self.http_client.post(url, json_body=data)
+            response = self.http_client.post(url, data=data)
+            if isinstance(response, dict):
+                return response
             results = response.json()
         except HttpError as e:
             results = {"error": str(e)}