fix:create type for BrokerMetaData

Author: cunla

scheduler/broker_types.py

@@ -1,4 +1,5 @@
 # This is a helper module to obfuscate types used by different broker implementations.
+from collections import namedtuple
 from typing import Union, Dict, Tuple, Type
 
 import redis
@@ -18,10 +19,12 @@
 PipelineType = Union[redis.client.Pipeline, valkey.client.Pipeline]
 SentinelType = Union[redis.sentinel.Sentinel, valkey.sentinel.Sentinel]
 
-BrokerMetaData: Dict[Tuple[Broker, bool], Tuple[Type[ConnectionType], Type[SentinelType], str]] = {
+BrokerMetaDataType = namedtuple("BrokerMetaDataType", ["connection_type", "sentinel_type", "ssl_prefix"])
+
+BrokerMetaData: Dict[Tuple[Broker, bool], BrokerMetaDataType] = {
     # Map of (Broker, Strict flag) => Connection Class, Sentinel Class, SSL Connection Prefix
-    (Broker.REDIS, False): (redis.Redis, redis.sentinel.Sentinel, "rediss"),
-    (Broker.VALKEY, False): (valkey.Valkey, valkey.sentinel.Sentinel, "valkeys"),
-    (Broker.REDIS, True): (redis.StrictRedis, redis.sentinel.Sentinel, "rediss"),
-    (Broker.VALKEY, True): (valkey.StrictValkey, valkey.sentinel.Sentinel, "valkeys"),
+    (Broker.REDIS, False): BrokerMetaDataType(redis.Redis, redis.sentinel.Sentinel, "rediss"),
+    (Broker.VALKEY, False): BrokerMetaDataType(valkey.Valkey, valkey.sentinel.Sentinel, "valkeys"),
+    (Broker.REDIS, True): BrokerMetaDataType(redis.StrictRedis, redis.sentinel.Sentinel, "rediss"),
+    (Broker.VALKEY, True): BrokerMetaDataType(valkey.StrictValkey, valkey.sentinel.Sentinel, "valkeys"),
 }

scheduler/queues.py

@@ -34,10 +34,10 @@ def _get_broker_connection(config, use_strict_broker=False):
 
         broker_cls = fakeredis.FakeRedis if not use_strict_broker else fakeredis.FakeStrictRedis
     else:
-        broker_cls = BrokerMetaData[(SCHEDULER_CONFIG.BROKER, use_strict_broker)][0]
+        broker_cls = BrokerMetaData[(SCHEDULER_CONFIG.BROKER, use_strict_broker)].connection_type
     logger.debug(f"Getting connection for {config}")
     if "URL" in config:
-        ssl_url_protocol = BrokerMetaData[(SCHEDULER_CONFIG.BROKER, use_strict_broker)][2]
+        ssl_url_protocol = BrokerMetaData[(SCHEDULER_CONFIG.BROKER, use_strict_broker)].ssl_prefix
         if config.get("SSL") or config.get("URL").startswith(f"{ssl_url_protocol}://"):
             return broker_cls.from_url(
                 config["URL"],
@@ -61,7 +61,7 @@ def _get_broker_connection(config, use_strict_broker=False):
         }
         connection_kwargs.update(config.get("CONNECTION_KWARGS", {}))
         sentinel_kwargs = config.get("SENTINEL_KWARGS", {})
-        SentinelClass = BrokerMetaData[(SCHEDULER_CONFIG.BROKER, use_strict_broker)][1]
+        SentinelClass = BrokerMetaData[(SCHEDULER_CONFIG.BROKER, use_strict_broker)].sentinel_type
         sentinel = SentinelClass(config["SENTINELS"], sentinel_kwargs=sentinel_kwargs, **connection_kwargs)
         return sentinel.master_for(
             service_name=config["MASTER_NAME"],