BrokerMetaData

Author: cunla

scheduler/admin/task_models.py

@@ -1,10 +1,9 @@
-import redis
-import valkey
 from django.contrib import admin, messages
 from django.contrib.contenttypes.admin import GenericStackedInline
 from django.utils.translation import gettext_lazy as _
 
 from scheduler import tools
+from scheduler.broker_types import ConnectionErrorTypes
 from scheduler.models import CronTask, TaskArg, TaskKwarg, RepeatableTask, ScheduledTask
 from scheduler.settings import SCHEDULER_CONFIG, logger
 from scheduler.tools import get_job_executions_for_task
@@ -186,7 +185,7 @@ def change_view(self, request, object_id, form_url="", extra_context=None):
         obj = self.get_object(request, object_id)
         try:
             execution_list = get_job_executions_for_task(obj.queue, obj)
-        except (redis.ConnectionError, valkey.ConnectionError) as e:
+        except ConnectionErrorTypes as e:
             logger.warn(f"Could not get job executions: {e}")
             execution_list = list()
         paginator = self.get_paginator(request, execution_list, SCHEDULER_CONFIG.EXECUTIONS_IN_PAGE)

scheduler/broker_types.py

@@ -0,0 +1,26 @@
+from typing import Union, Dict, Tuple, Type
+
+import redis
+
+try:
+    import valkey
+except ImportError:
+    valkey = redis
+    valkey.Valkey = redis.Redis
+    valkey.StrictValkey = redis.StrictRedis
+
+from scheduler.settings import Broker
+
+ConnectionErrorTypes = (redis.ConnectionError, valkey.ConnectionError)
+ResponseErrorTypes = (redis.ResponseError, valkey.ResponseError)
+ConnectionType = Union[redis.Redis, valkey.Valkey]
+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]] = {
+    # 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"),
+}

scheduler/connection_types.py

@@ -3,12 +3,11 @@
 import sys
 
 import click
-import redis
-import valkey
 from django.core.management.base import BaseCommand
 from django.db import connections
 from rq.logutils import setup_loghandlers
 
+from scheduler.broker_types import ConnectionErrorTypes
 from scheduler.rq_classes import register_sentry
 from scheduler.tools import create_worker
 
@@ -136,6 +135,6 @@ def handle(self, **options):
                 logging_level=log_level,
                 max_jobs=options["max_jobs"],
             )
-        except (redis.ConnectionError, valkey.ConnectionError) as e:
+        except ConnectionErrorTypes as e:
             click.echo(str(e), err=True)
             sys.exit(1)

scheduler/management/commands/rqworker.py

@@ -1,21 +1,10 @@
 from typing import List, Dict, Set
 
-import redis
-
-from .connection_types import RedisSentinel, ValkeySentinel, BrokerConnectionClass
+from .broker_types import ConnectionErrorTypes, BrokerMetaData
 from .rq_classes import JobExecution, DjangoQueue, DjangoWorker
 from .settings import SCHEDULER_CONFIG
 from .settings import logger, Broker
 
-
-try:
-    from valkey import exceptions
-except ImportError:
-    exceptions = ""
-    exceptions.ConnectionError = redis.ConnectionError
-
-ConnectionErrors = (redis.ConnectionError, exceptions.ConnectionError)
-
 _CONNECTION_PARAMS = {
     "URL",
     "DB",
@@ -36,18 +25,6 @@ class QueueNotFoundError(Exception):
     pass
 
 
-ssl_url_protocol = {
-    "valkey": "valkeys",
-    "redis": "rediss",
-    "fakeredis": "rediss"
-}
-
-sentinel_broker = {
-    "valkey": ValkeySentinel,
-    "redis": RedisSentinel,
-}
-
-
 def _get_broker_connection(config, use_strict_broker=False):
     """
     Returns a redis connection from a connection config
@@ -57,10 +34,11 @@ def _get_broker_connection(config, use_strict_broker=False):
 
         broker_cls = fakeredis.FakeRedis if not use_strict_broker else fakeredis.FakeStrictRedis
     else:
-        broker_cls = BrokerConnectionClass[(SCHEDULER_CONFIG.BROKER, use_strict_broker)]
+        broker_cls = BrokerMetaData[(SCHEDULER_CONFIG.BROKER, use_strict_broker)][0]
     logger.debug(f"Getting connection for {config}")
     if "URL" in config:
-        if config.get("SSL") or config.get("URL").startswith(f"{ssl_url_protocol[SCHEDULER_CONFIG.BROKER.value]}://"):
+        ssl_url_protocol = BrokerMetaData[(SCHEDULER_CONFIG.BROKER, use_strict_broker)][2]
+        if config.get("SSL") or config.get("URL").startswith(f"{ssl_url_protocol}://"):
             return broker_cls.from_url(
                 config["URL"],
                 db=config.get("DB"),
@@ -83,9 +61,8 @@ def _get_broker_connection(config, use_strict_broker=False):
         }
         connection_kwargs.update(config.get("CONNECTION_KWARGS", {}))
         sentinel_kwargs = config.get("SENTINEL_KWARGS", {})
-        sentinel = sentinel_broker[SCHEDULER_CONFIG.BROKER.value](
-            config["SENTINELS"], sentinel_kwargs=sentinel_kwargs, **connection_kwargs
-        )
+        SentinelClass = BrokerMetaData[(SCHEDULER_CONFIG.BROKER, use_strict_broker)][1]
+        sentinel = SentinelClass(config["SENTINELS"], sentinel_kwargs=sentinel_kwargs, **connection_kwargs)
         return sentinel.master_for(
             service_name=config["MASTER_NAME"],
             redis_class=broker_cls,
@@ -138,7 +115,7 @@ def get_all_workers() -> Set[DjangoWorker]:
         try:
             curr_workers: Set[DjangoWorker] = set(DjangoWorker.all(connection=connection))
             workers_set.update(curr_workers)
-        except ConnectionErrors as e:
+        except ConnectionErrorTypes as e:
             logger.error(f"Could not connect for queue {queue_name}: {e}")
     return workers_set
 

scheduler/queues.py

@@ -22,7 +22,7 @@
 from rq.worker import WorkerStatus
 
 from scheduler import settings
-from scheduler.connection_types import PipelineType, ConnectionType
+from scheduler.broker_types import PipelineType, ConnectionType
 
 MODEL_NAMES = ["ScheduledTask", "RepeatableTask", "CronTask"]
 

scheduler/rq_classes.py

@@ -13,22 +13,13 @@
 from django.views.decorators.cache import never_cache
 from redis.exceptions import ResponseError
 
-
-from .queues import get_all_workers, get_connection, QueueNotFoundError, ConnectionErrors
+from .broker_types import ConnectionErrorTypes
+from .queues import get_all_workers, get_connection, QueueNotFoundError
 from .queues import get_queue as get_queue_base
 from .rq_classes import JobExecution, DjangoWorker, DjangoQueue, InvalidJobOperation
 from .settings import SCHEDULER_CONFIG, logger
 
 
-try:
-    from valkey import exceptions
-except ImportError:
-    exceptions = ""
-    exceptions.ResponseError = ResponseError
-
-ResponseErrors = (ResponseError, exceptions.ResponseError)
-
-
 def get_queue(queue_name: str) -> DjangoQueue:
     try:
         return get_queue_base(queue_name)
@@ -111,7 +102,7 @@ def get_statistics(run_maintenance_tasks=False):
                 canceled_jobs=len(queue.canceled_job_registry),
             )
             queues.append(queue_data)
-        except ConnectionErrors as e:
+        except ConnectionErrorTypes as e:
             logger.error(f"Could not connect for queue {queue_name}: {e}")
             continue