1111 NodeMigratedNotification ,
1212 NodeMigratingNotification ,
1313 NodeMovingNotification ,
14+ OSSNodeMigratedNotification ,
15+ OSSNodeMigratingNotification ,
1416)
17+ from redis .utils import safe_str
1518
1619if sys .version_info .major >= 3 and sys .version_info .minor >= 11 :
1720 from asyncio import timeout as async_timeout
@@ -178,16 +181,56 @@ async def read_response(
178181class MaintenanceNotificationsParser :
179182 """Protocol defining maintenance push notification parsing functionality"""
180183
184+ @staticmethod
185+ def parse_oss_maintenance_start_msg (response ):
186+ # Expected message format is:
187+ # SMIGRATING <seq_number> <slot, range1-range2,...>
188+ id = response [1 ]
189+ slots = safe_str (response [2 ])
190+ return OSSNodeMigratingNotification (id , slots )
191+
192+ @staticmethod
193+ def parse_oss_maintenance_completed_msg (response ):
194+ # Expected message format is:
195+ # SMIGRATED <seq_number> [[<src_host:port> <dest_host:port> <slot_range>], ...]
196+ id = response [1 ]
197+ nodes_to_slots_mapping_data = response [2 ]
198+ # Build the nodes_to_slots_mapping dict structure:
199+ # {
200+ # "src_host:port": [
201+ # {"dest_host:port": "slot_range"},
202+ # ...
203+ # ],
204+ # ...
205+ # }
206+ nodes_to_slots_mapping = {}
207+ for src_node , dest_node , slots in nodes_to_slots_mapping_data :
208+ src_node_str = safe_str (src_node )
209+ dest_node_str = safe_str (dest_node )
210+ slots_str = safe_str (slots )
211+
212+ if src_node_str not in nodes_to_slots_mapping :
213+ nodes_to_slots_mapping [src_node_str ] = []
214+ nodes_to_slots_mapping [src_node_str ].append ({dest_node_str : slots_str })
215+
216+ return OSSNodeMigratedNotification (id , nodes_to_slots_mapping )
217+
181218 @staticmethod
182219 def parse_maintenance_start_msg (response , notification_type ):
183220 # Expected message format is: <notification_type> <seq_number> <time>
221+ # Examples:
222+ # MIGRATING 1 10
223+ # FAILING_OVER 2 20
184224 id = response [1 ]
185225 ttl = response [2 ]
186226 return notification_type (id , ttl )
187227
188228 @staticmethod
189229 def parse_maintenance_completed_msg (response , notification_type ):
190230 # Expected message format is: <notification_type> <seq_number>
231+ # Examples:
232+ # MIGRATED 1
233+ # FAILED_OVER 2
191234 id = response [1 ]
192235 return notification_type (id )
193236
@@ -199,9 +242,7 @@ def parse_moving_msg(response):
199242 if response [3 ] is None :
200243 host , port = None , None
201244 else :
202- value = response [3 ]
203- if isinstance (value , bytes ):
204- value = value .decode ()
245+ value = safe_str (response [3 ])
205246 host , port = value .split (":" )
206247 port = int (port ) if port is not None else None
207248
@@ -214,12 +255,15 @@ def parse_moving_msg(response):
214255_MIGRATED_MESSAGE = "MIGRATED"
215256_FAILING_OVER_MESSAGE = "FAILING_OVER"
216257_FAILED_OVER_MESSAGE = "FAILED_OVER"
258+ _SMIGRATING_MESSAGE = "SMIGRATING"
259+ _SMIGRATED_MESSAGE = "SMIGRATED"
217260
218261_MAINTENANCE_MESSAGES = (
219262 _MIGRATING_MESSAGE ,
220263 _MIGRATED_MESSAGE ,
221264 _FAILING_OVER_MESSAGE ,
222265 _FAILED_OVER_MESSAGE ,
266+ _SMIGRATING_MESSAGE ,
223267)
224268
225269MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING : dict [
@@ -245,6 +289,14 @@ def parse_moving_msg(response):
245289 NodeMovingNotification ,
246290 MaintenanceNotificationsParser .parse_moving_msg ,
247291 ),
292+ _SMIGRATING_MESSAGE : (
293+ OSSNodeMigratingNotification ,
294+ MaintenanceNotificationsParser .parse_oss_maintenance_start_msg ,
295+ ),
296+ _SMIGRATED_MESSAGE : (
297+ OSSNodeMigratedNotification ,
298+ MaintenanceNotificationsParser .parse_oss_maintenance_completed_msg ,
299+ ),
248300}
249301
250302
@@ -255,6 +307,7 @@ class PushNotificationsParser(Protocol):
255307 invalidation_push_handler_func : Optional [Callable ] = None
256308 node_moving_push_handler_func : Optional [Callable ] = None
257309 maintenance_push_handler_func : Optional [Callable ] = None
310+ oss_cluster_maint_push_handler_func : Optional [Callable ] = None
258311
259312 def handle_pubsub_push_response (self , response ):
260313 """Handle pubsub push responses"""
@@ -269,6 +322,7 @@ def handle_push_response(self, response, **kwargs):
269322 _INVALIDATION_MESSAGE ,
270323 * _MAINTENANCE_MESSAGES ,
271324 _MOVING_MESSAGE ,
325+ _SMIGRATED_MESSAGE ,
272326 ):
273327 return self .pubsub_push_handler_func (response )
274328
@@ -291,13 +345,30 @@ def handle_push_response(self, response, **kwargs):
291345 parser_function = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING [
292346 msg_type
293347 ][1 ]
294- notification_type = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING [
295- msg_type
296- ][0 ]
297- notification = parser_function (response , notification_type )
348+ if msg_type == _SMIGRATING_MESSAGE :
349+ notification = parser_function (response )
350+ else :
351+ notification_type = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING [
352+ msg_type
353+ ][0 ]
354+ notification = parser_function (response , notification_type )
298355
299356 if notification is not None :
300357 return self .maintenance_push_handler_func (notification )
358+ if msg_type == _SMIGRATED_MESSAGE and (
359+ self .oss_cluster_maint_push_handler_func
360+ or self .maintenance_push_handler_func
361+ ):
362+ parser_function = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING [
363+ msg_type
364+ ][1 ]
365+ notification = parser_function (response )
366+
367+ if notification is not None :
368+ if self .maintenance_push_handler_func :
369+ self .maintenance_push_handler_func (notification )
370+ if self .oss_cluster_maint_push_handler_func :
371+ self .oss_cluster_maint_push_handler_func (notification )
301372 except Exception as e :
302373 logger .error (
303374 "Error handling {} message ({}): {}" .format (msg_type , response , e )
@@ -317,6 +388,9 @@ def set_node_moving_push_handler(self, node_moving_push_handler_func):
317388 def set_maintenance_push_handler (self , maintenance_push_handler_func ):
318389 self .maintenance_push_handler_func = maintenance_push_handler_func
319390
391+ def set_oss_cluster_maint_push_handler (self , oss_cluster_maint_push_handler_func ):
392+ self .oss_cluster_maint_push_handler_func = oss_cluster_maint_push_handler_func
393+
320394
321395class AsyncPushNotificationsParser (Protocol ):
322396 """Protocol defining async RESP3-specific parsing functionality"""
@@ -325,6 +399,7 @@ class AsyncPushNotificationsParser(Protocol):
325399 invalidation_push_handler_func : Optional [Callable ] = None
326400 node_moving_push_handler_func : Optional [Callable [..., Awaitable [None ]]] = None
327401 maintenance_push_handler_func : Optional [Callable [..., Awaitable [None ]]] = None
402+ oss_cluster_maint_push_handler_func : Optional [Callable [..., Awaitable [None ]]] = None
328403
329404 async def handle_pubsub_push_response (self , response ):
330405 """Handle pubsub push responses asynchronously"""
@@ -341,6 +416,7 @@ async def handle_push_response(self, response, **kwargs):
341416 _INVALIDATION_MESSAGE ,
342417 * _MAINTENANCE_MESSAGES ,
343418 _MOVING_MESSAGE ,
419+ _SMIGRATED_MESSAGE ,
344420 ):
345421 return await self .pubsub_push_handler_func (response )
346422
@@ -365,13 +441,26 @@ async def handle_push_response(self, response, **kwargs):
365441 parser_function = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING [
366442 msg_type
367443 ][1 ]
368- notification_type = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING [
369- msg_type
370- ][0 ]
371- notification = parser_function (response , notification_type )
444+ if msg_type == _SMIGRATING_MESSAGE :
445+ notification = parser_function (response )
446+ else :
447+ notification_type = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING [
448+ msg_type
449+ ][0 ]
450+ notification = parser_function (response , notification_type )
372451
373452 if notification is not None :
374453 return await self .maintenance_push_handler_func (notification )
454+ if (
455+ msg_type == _SMIGRATED_MESSAGE
456+ and self .oss_cluster_maint_push_handler_func
457+ ):
458+ parser_function = MSG_TYPE_TO_MAINT_NOTIFICATION_PARSER_MAPPING [
459+ msg_type
460+ ][1 ]
461+ notification = parser_function (response )
462+ if notification is not None :
463+ return await self .oss_cluster_maint_push_handler_func (notification )
375464 except Exception as e :
376465 logger .error (
377466 "Error handling {} message ({}): {}" .format (msg_type , response , e )
@@ -393,6 +482,9 @@ def set_node_moving_push_handler(self, node_moving_push_handler_func):
393482 def set_maintenance_push_handler (self , maintenance_push_handler_func ):
394483 self .maintenance_push_handler_func = maintenance_push_handler_func
395484
485+ def set_oss_cluster_maint_push_handler (self , oss_cluster_maint_push_handler_func ):
486+ self .oss_cluster_maint_push_handler_func = oss_cluster_maint_push_handler_func
487+
396488
397489class _AsyncRESPBase (AsyncBaseParser ):
398490 """Base class for async resp parsing"""
0 commit comments