Added MultiDbClient support with OSS Cluster API

redis/client.py

@@ -1217,6 +1217,8 @@ def run_in_thread(
         sleep_time: float = 0.0,
         daemon: bool = False,
         exception_handler: Optional[Callable] = None,
+        pubsub = None,
+        sharded_pubsub: bool = False,
     ) -> "PubSubWorkerThread":
         for channel, handler in self.channels.items():
             if handler is None:
@@ -1230,8 +1232,9 @@ def run_in_thread(
                     f"Shard Channel: '{s_channel}' has no handler registered"
                 )
 
+        pubsub = self if pubsub is None else pubsub
         thread = PubSubWorkerThread(
-            self, sleep_time, daemon=daemon, exception_handler=exception_handler
+            pubsub, sleep_time, daemon=daemon, exception_handler=exception_handler, sharded_pubsub=sharded_pubsub
         )
         thread.start()
         return thread
@@ -1246,12 +1249,14 @@ def __init__(
         exception_handler: Union[
             Callable[[Exception, "PubSub", "PubSubWorkerThread"], None], None
         ] = None,
+        sharded_pubsub: bool = False,
     ):
         super().__init__()
         self.daemon = daemon
         self.pubsub = pubsub
         self.sleep_time = sleep_time
         self.exception_handler = exception_handler
+        self.sharded_pubsub = sharded_pubsub
         self._running = threading.Event()
 
     def run(self) -> None:
@@ -1262,7 +1267,10 @@ def run(self) -> None:
         sleep_time = self.sleep_time
         while self._running.is_set():
             try:
-                pubsub.get_message(ignore_subscribe_messages=True, timeout=sleep_time)
+                if not self.sharded_pubsub:
+                    pubsub.get_message(ignore_subscribe_messages=True, timeout=sleep_time)
+                else:
+                    pubsub.get_sharded_message(ignore_subscribe_messages=True, timeout=sleep_time)
             except BaseException as e:
                 if self.exception_handler is None:
                     raise

redis/cluster.py

@@ -3154,7 +3154,8 @@ def _reinitialize_on_error(self, error):
                 self._nodes_manager.initialize()
                 self.reinitialize_counter = 0
             else:
-                self._nodes_manager.update_moved_exception(error)
+                if type(error) == MovedError:
+                    self._nodes_manager.update_moved_exception(error)
 
         self._executing = False
 

redis/multidb/client.py

@@ -1,14 +1,15 @@
 import threading
 import socket
-from typing import List, Any, Callable
+from typing import List, Any, Callable, Optional
 
 from redis.background import BackgroundScheduler
+from redis.client import PubSubWorkerThread
 from redis.exceptions import ConnectionError, TimeoutError
 from redis.commands import RedisModuleCommands, CoreCommands
 from redis.multidb.command_executor import DefaultCommandExecutor
 from redis.multidb.config import MultiDbConfig, DEFAULT_GRACE_PERIOD
 from redis.multidb.circuit import State as CBState, CircuitBreaker
-from redis.multidb.database import State as DBState, Database, AbstractDatabase, Databases
+from redis.multidb.database import Database, AbstractDatabase, Databases
 from redis.multidb.exception import NoValidDatabaseException
 from redis.multidb.failure_detector import FailureDetector
 from redis.multidb.healthcheck import HealthCheck
@@ -70,13 +71,8 @@ def raise_exception_on_failed_hc(error):
 
             # Set states according to a weights and circuit state
             if database.circuit.state == CBState.CLOSED and not is_active_db_found:
-                database.state = DBState.ACTIVE
                 self.command_executor.active_database = database
                 is_active_db_found = True
-            elif database.circuit.state == CBState.CLOSED and is_active_db_found:
-                database.state = DBState.PASSIVE
-            else:
-                database.state = DBState.DISCONNECTED
 
         if not is_active_db_found:
             raise NoValidDatabaseException('Initial connection failed - no active database found')
@@ -107,8 +103,6 @@ def set_active_database(self, database: AbstractDatabase) -> None:
 
         if database.circuit.state == CBState.CLOSED:
             highest_weighted_db, _ = self._databases.get_top_n(1)[0]
-            highest_weighted_db.state = DBState.PASSIVE
-            database.state = DBState.ACTIVE
             self.command_executor.active_database = database
             return
 
@@ -130,9 +124,7 @@ def add_database(self, database: AbstractDatabase):
 
     def _change_active_database(self, new_database: AbstractDatabase, highest_weight_database: AbstractDatabase):
         if new_database.weight > highest_weight_database.weight and new_database.circuit.state == CBState.CLOSED:
-            new_database.state = DBState.ACTIVE
             self.command_executor.active_database = new_database
-            highest_weight_database.state = DBState.PASSIVE
 
     def remove_database(self, database: Database):
         """
@@ -142,7 +134,6 @@ def remove_database(self, database: Database):
         highest_weighted_db, highest_weight = self._databases.get_top_n(1)[0]
 
         if highest_weight <= weight and highest_weighted_db.circuit.state == CBState.CLOSED:
-            highest_weighted_db.state = DBState.ACTIVE
             self.command_executor.active_database = highest_weighted_db
 
     def update_database_weight(self, database: AbstractDatabase, weight: float):
@@ -201,6 +192,17 @@ def transaction(self, func: Callable[["Pipeline"], None], *watches, **options):
 
         return self.command_executor.execute_transaction(func, *watches, *options)
 
+    def pubsub(self, **kwargs):
+        """
+        Return a Publish/Subscribe object. With this object, you can
+        subscribe to channels and listen for messages that get published to
+        them.
+        """
+        if not self.initialized:
+            self.initialize()
+
+        return PubSub(self, **kwargs)
+
     def _check_db_health(self, database: AbstractDatabase, on_error: Callable[[Exception], None] = None) -> None:
         """
         Runs health checks on the given database until first failure.
@@ -221,7 +223,7 @@ def _check_db_health(self, database: AbstractDatabase, on_error: Callable[[Excep
                         database.circuit.state = CBState.OPEN
                     elif is_healthy and database.circuit.state != CBState.CLOSED:
                         database.circuit.state = CBState.CLOSED
-                except (ConnectionError, TimeoutError, socket.timeout, ConnectionRefusedError) as e:
+                except (ConnectionError, TimeoutError, socket.timeout, ConnectionRefusedError, ValueError) as e:
                     if database.circuit.state != CBState.OPEN:
                         database.circuit.state = CBState.OPEN
                     is_healthy = False
@@ -311,3 +313,133 @@ def execute(self) -> List[Any]:
             return self._client.command_executor.execute_pipeline(tuple(self._command_stack))
         finally:
             self.reset()
+
+class PubSub:
+    """
+    PubSub object for multi database client.
+    """
+    def __init__(self, client: MultiDBClient, **kwargs):
+        self._client = client
+        self._client.command_executor.pubsub(**kwargs)
+
+    def __enter__(self) -> "PubSub":
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        self.reset()
+
+    def __del__(self) -> None:
+        try:
+            # if this object went out of scope prior to shutting down
+            # subscriptions, close the connection manually before
+            # returning it to the connection pool
+            self.reset()
+        except Exception:
+            pass
+
+    def reset(self) -> None:
+        pass
+
+    def close(self) -> None:
+        self.reset()
+
+    @property
+    def subscribed(self) -> bool:
+        return self._client.command_executor.active_pubsub.subscribed
+
+    def psubscribe(self, *args, **kwargs):
+        """
+        Subscribe to channel patterns. Patterns supplied as keyword arguments
+        expect a pattern name as the key and a callable as the value. A
+        pattern's callable will be invoked automatically when a message is
+        received on that pattern rather than producing a message via
+        ``listen()``.
+        """
+        return self._client.command_executor.execute_pubsub_method('psubscribe', *args, **kwargs)
+
+    def punsubscribe(self, *args):
+        """
+        Unsubscribe from the supplied patterns. If empty, unsubscribe from
+        all patterns.
+        """
+        return self._client.command_executor.execute_pubsub_method('punsubscribe', *args)
+
+    def subscribe(self, *args, **kwargs):
+        """
+        Subscribe to channels. Channels supplied as keyword arguments expect
+        a channel name as the key and a callable as the value. A channel's
+        callable will be invoked automatically when a message is received on
+        that channel rather than producing a message via ``listen()`` or
+        ``get_message()``.
+        """
+        return self._client.command_executor.execute_pubsub_method('subscribe', *args, **kwargs)
+
+    def unsubscribe(self, *args):
+        """
+        Unsubscribe from the supplied channels. If empty, unsubscribe from
+        all channels
+        """
+        return self._client.command_executor.execute_pubsub_method('unsubscribe', *args)
+
+    def ssubscribe(self, *args, **kwargs):
+        """
+        Subscribes the client to the specified shard channels.
+        Channels supplied as keyword arguments expect a channel name as the key
+        and a callable as the value. A channel's callable will be invoked automatically
+        when a message is received on that channel rather than producing a message via
+        ``listen()`` or ``get_sharded_message()``.
+        """
+        return self._client.command_executor.execute_pubsub_method('ssubscribe', *args, **kwargs)
+
+    def sunsubscribe(self, *args):
+        """
+        Unsubscribe from the supplied shard_channels. If empty, unsubscribe from
+        all shard_channels
+        """
+        return self._client.command_executor.execute_pubsub_method('sunsubscribe', *args)
+
+    def get_message(
+        self, ignore_subscribe_messages: bool = False, timeout: float = 0.0
+    ):
+        """
+        Get the next message if one is available, otherwise None.
+
+        If timeout is specified, the system will wait for `timeout` seconds
+        before returning. Timeout should be specified as a floating point
+        number, or None, to wait indefinitely.
+        """
+        return self._client.command_executor.execute_pubsub_method(
+            'get_message',
+            ignore_subscribe_messages=ignore_subscribe_messages, timeout=timeout
+        )
+
+    def get_sharded_message(
+        self, ignore_subscribe_messages: bool = False, timeout: float = 0.0
+    ):
+        """
+        Get the next message if one is available in a sharded channel, otherwise None.
+
+        If timeout is specified, the system will wait for `timeout` seconds
+        before returning. Timeout should be specified as a floating point
+        number, or None, to wait indefinitely.
+        """
+        return self._client.command_executor.execute_pubsub_method(
+            'get_sharded_message',
+            ignore_subscribe_messages=ignore_subscribe_messages, timeout=timeout
+        )
+
+    def run_in_thread(
+        self,
+        sleep_time: float = 0.0,
+        daemon: bool = False,
+        exception_handler: Optional[Callable] = None,
+        sharded_pubsub: bool = False,
+    ) -> "PubSubWorkerThread":
+        return self._client.command_executor.execute_pubsub_run_in_thread(
+            sleep_time=sleep_time,
+            daemon=daemon,
+            exception_handler=exception_handler,
+            pubsub=self,
+            sharded_pubsub=sharded_pubsub,
+        )
+