wip

Author: cunla

scheduler/helpers/callback.py

@@ -30,7 +30,7 @@ def __init__(self, func: Union[str, Callable[..., Any]], timeout: Optional[int]
     def name(self) -> str:
         return f"{self.func.__module__}.{self.func.__qualname__}"
 
-    def __call__(self, *args, **kwargs):
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
         from scheduler.settings import SCHEDULER_CONFIG
 
         with SCHEDULER_CONFIG.DEATH_PENALTY_CLASS(self.timeout, JobTimeoutException):

scheduler/helpers/queues/getters.py

@@ -1,15 +1,14 @@
 from typing import Set
 
 from scheduler.redis_models.worker import WorkerModel
-from scheduler.settings import SCHEDULER_CONFIG, get_queue_names, get_queue_configuration, QueueConfiguration, logger
-from scheduler.types import ConnectionErrorTypes, BrokerMetaData, Broker
+from scheduler.settings import SCHEDULER_CONFIG, get_queue_names, get_queue_configuration, logger
+from scheduler.types import ConnectionErrorTypes, BrokerMetaData, Broker, ConnectionType, QueueConfiguration
 from .queue_logic import Queue
 
-
 _BAD_QUEUE_CONFIGURATION = set()
 
 
-def _get_connection(config: QueueConfiguration, use_strict_broker=False):
+def _get_connection(config: QueueConfiguration, use_strict_broker: bool = False) -> ConnectionType:
     """Returns a Broker connection to use based on parameters in SCHEDULER_QUEUES"""
     if SCHEDULER_CONFIG.BROKER == Broker.FAKEREDIS:
         import fakeredis
@@ -32,7 +31,7 @@ def _get_connection(config: QueueConfiguration, use_strict_broker=False):
         sentinel_kwargs = config.SENTINEL_KWARGS or {}
         SentinelClass = BrokerMetaData[(SCHEDULER_CONFIG.BROKER, use_strict_broker)].sentinel_type
         sentinel = SentinelClass(config.SENTINELS, sentinel_kwargs=sentinel_kwargs, **connection_kwargs)
-        return sentinel.master_for(
+        return sentinel.master_for(  # type: ignore
             service_name=config.MASTER_NAME,
             redis_class=broker_cls,
         )
@@ -47,7 +46,7 @@ def _get_connection(config: QueueConfiguration, use_strict_broker=False):
     )
 
 
-def get_queue(name="default") -> Queue:
+def get_queue(name: str = "default") -> Queue:
     """Returns an DjangoQueue using parameters defined in `SCHEDULER_QUEUES`"""
     queue_settings = get_queue_configuration(name)
     is_async = queue_settings.ASYNC

scheduler/helpers/queues/queue_logic.py

@@ -19,7 +19,7 @@
 )
 from scheduler.redis_models import JobStatus, SchedulerLock, Result, ResultType, JobModel
 from scheduler.settings import logger, SCHEDULER_CONFIG
-from scheduler.types import ConnectionType, FunctionReferenceType, Self
+from scheduler.types import ConnectionType, FunctionReferenceType, Self, PipelineType
 
 
 class InvalidJobOperation(Exception):
@@ -49,17 +49,17 @@ def perform_job(job_model: JobModel, connection: ConnectionType) -> Any:  # noqa
             coro_result = loop.run_until_complete(result)
             result = coro_result
         if job_model.success_callback:
-            job_model.success_callback(job_model, connection, result)  # type: ignore
+            job_model.success_callback(job_model, connection, result)
         return result
     except:
         if job_model.failure_callback:
-            job_model.failure_callback(job_model, connection, *sys.exc_info())  # type: ignore
+            job_model.failure_callback(job_model, connection, *sys.exc_info())
         raise
     finally:
         assert job_model is _job_stack.pop()
 
 
-_job_stack = []
+_job_stack: List[JobModel] = []
 
 
 class Queue:
@@ -89,11 +89,11 @@ def __init__(self, connection: ConnectionType, name: str, is_async: bool = True)
         self.scheduled_job_registry = ScheduledJobRegistry(connection=self.connection, name=self.name)
         self.canceled_job_registry = CanceledJobRegistry(connection=self.connection, name=self.name)
 
-    def __len__(self):
+    def __len__(self) -> int:
         return self.count
 
     @property
-    def scheduler_pid(self) -> int:
+    def scheduler_pid(self) -> Optional[int]:
         lock = SchedulerLock(self.name)
         pid = lock.value(self.connection)
         return int(pid.decode()) if pid is not None else None
@@ -162,7 +162,7 @@ def count(self) -> int:
     def get_registry(self, name: str) -> JobNamesRegistry:
         name = name.lower()
         if name in Queue.REGISTRIES:
-            return getattr(self, Queue.REGISTRIES[name])
+            return getattr(self, Queue.REGISTRIES[name])  # type: ignore
         raise NoSuchRegistryError(f"Unknown registry name {name}")
 
     def get_all_job_names(self) -> List[str]:
@@ -182,22 +182,21 @@ def get_all_jobs(self) -> List[JobModel]:
     def create_and_enqueue_job(
         self,
         func: FunctionReferenceType,
-        args: Union[Tuple, List, None] = None,
-        kwargs: Optional[Dict] = None,
+        args: Union[Tuple[Any, ...], List[Any], None] = None,
+        kwargs: Optional[Dict[str, Any]] = None,
         when: Optional[datetime] = None,
         timeout: Optional[int] = None,
         result_ttl: Optional[int] = None,
         job_info_ttl: Optional[int] = None,
         description: Optional[str] = None,
         name: Optional[str] = None,
         at_front: bool = False,
-        meta: Optional[Dict] = None,
+        meta: Optional[Dict[str, Any]] = None,
         on_success: Optional[Callback] = None,
         on_failure: Optional[Callback] = None,
         on_stopped: Optional[Callback] = None,
         task_type: Optional[str] = None,
         scheduled_task_id: Optional[int] = None,
-        pipeline: Optional[ConnectionType] = None,
     ) -> JobModel:
         """Creates a job to represent the delayed function call and enqueues it.
         :param when: When to schedule the job (None to enqueue immediately)
@@ -216,7 +215,6 @@ def create_and_enqueue_job(
         :param on_stopped: Callback for on stopped
         :param task_type: The task type
         :param scheduled_task_id: The scheduled task id
-        :param pipeline: The Broker Pipeline
         :returns: The enqueued Job
         """
         status = JobStatus.QUEUED if when is None else JobStatus.SCHEDULED
@@ -240,7 +238,7 @@ def create_and_enqueue_job(
             scheduled_task_id=scheduled_task_id,
         )
         if when is None:
-            job_model = self.enqueue_job(job_model, connection=pipeline, at_front=at_front)
+            job_model = self.enqueue_job(job_model, at_front=at_front)
         elif isinstance(when, datetime):
             job_model.save(connection=self.connection)
             self.scheduled_job_registry.schedule(self.connection, job_model.name, when)
@@ -250,7 +248,7 @@ def create_and_enqueue_job(
 
     def job_handle_success(
         self, job: JobModel, result: Any, job_info_ttl: int, result_ttl: int, connection: ConnectionType
-    ):
+    ) -> None:
         """Saves and cleanup job after successful execution"""
         job.after_execution(
             job_info_ttl,
@@ -268,7 +266,7 @@ def job_handle_success(
             ttl=result_ttl,
         )
 
-    def job_handle_failure(self, status: JobStatus, job: JobModel, exc_string: str, connection: ConnectionType):
+    def job_handle_failure(self, status: JobStatus, job: JobModel, exc_string: str, connection: ConnectionType) -> None:
         # Does not set job status since the job might be stopped
         job.after_execution(
             SCHEDULER_CONFIG.DEFAULT_FAILURE_TTL,
@@ -308,10 +306,7 @@ def run_sync(self, job: JobModel) -> JobModel:
 
     @classmethod
     def dequeue_any(
-        cls,
-        queues: List[Self],
-        timeout: Optional[int],
-        connection: Optional[ConnectionType] = None,
+        cls, queues: List[Self], timeout: Optional[int], connection: ConnectionType
     ) -> Tuple[Optional[JobModel], Optional[Self]]:
         """Class method returning a Job instance at the front of the given set of Queues, where the order of the queues
         is important.
@@ -414,19 +409,19 @@ def delete_job(self, job_name: str, expire_job_model: bool = True) -> None:
                 pass
 
     def enqueue_job(
-        self, job_model: JobModel, connection: Optional[ConnectionType] = None, at_front: bool = False
+        self, job_model: JobModel, pipeline: Optional[PipelineType] = None, at_front: bool = False
     ) -> JobModel:
         """Enqueues a job for delayed execution without checking dependencies.
 
         If Queue is instantiated with is_async=False, job is executed immediately.
         :param job_model: The job redis model
-        :param connection: The Redis Pipeline
+        :param pipeline: The Broker Pipeline
         :param at_front: Whether to enqueue the job at the front
 
         :returns: The enqueued JobModel
         """
 
-        pipe = connection if connection is not None else self.connection.pipeline()
+        pipe: PipelineType = pipeline if pipeline is not None else self.connection.pipeline()
         job_model.started_at = None
         job_model.ended_at = None
         job_model.status = JobStatus.QUEUED

scheduler/helpers/utils.py

@@ -1,7 +1,7 @@
 import datetime
 import importlib
 import time
-from typing import Callable
+from typing import Callable, Any
 
 
 def current_timestamp() -> int:
@@ -14,10 +14,10 @@ def utcnow() -> datetime.datetime:
     return datetime.datetime.now(datetime.timezone.utc)
 
 
-def callable_func(callable_str: str) -> Callable:
+def callable_func(callable_str: str) -> Callable[[Any], Any]:
     path = callable_str.split(".")
     module = importlib.import_module(".".join(path[:-1]))
-    func = getattr(module, path[-1])
-    if callable(func) is False:
+    func: Callable[[Any], Any] = getattr(module, path[-1])
+    if not callable(func):
         raise TypeError(f"'{callable_str}' is not callable")
     return func

scheduler/models/args.py

@@ -9,7 +9,7 @@
 
 from scheduler.helpers import utils
 
-ARG_TYPE_TYPES_DICT:Dict[str,Type] = {
+ARG_TYPE_TYPES_DICT: Dict[str, Type] = {
     "str": str,
     "int": int,
     "bool": bool,

scheduler/models/task.py

@@ -33,7 +33,7 @@ def _get_task_for_job(job: JobModel) -> Optional["Task"]:
     return task
 
 
-def failure_callback(job: JobModel, connection:ConnectionType, result: Any, *args: Any, **kwargs: Any) -> None:
+def failure_callback(job: JobModel, connection: ConnectionType, result: Any, *args: Any, **kwargs: Any) -> None:
     task = _get_task_for_job(job)
     if task is None:
         logger.warn(f"Could not find task for job {job.name}")
@@ -176,7 +176,7 @@ class TimeUnits(models.TextChoices):
         ),
     )
 
-    def callable_func(self)->Callable:
+    def callable_func(self) -> Callable:
         """Translate callable string to callable"""
         return utils.callable_func(self.callable)
 
@@ -187,9 +187,9 @@ def is_scheduled(self) -> bool:
             return False
         # check whether job_id is in scheduled/queued/active jobs
         res = (
-                (self.job_name in self.rqueue.scheduled_job_registry.all())
-                or (self.job_name in self.rqueue.queued_job_registry.all())
-                or (self.job_name in self.rqueue.active_job_registry.all())
+            (self.job_name in self.rqueue.scheduled_job_registry.all())
+            or (self.job_name in self.rqueue.queued_job_registry.all())
+            or (self.job_name in self.rqueue.active_job_registry.all())
         )
         # If the job_id is not scheduled/queued/started,
         # update the job_id to None. (The job_id belongs to a previous run which is completed)

scheduler/redis_models/registry/base_registry.py

@@ -74,11 +74,11 @@ def get_first(self) -> Optional[str]:
         first_job = self.connection.zrange(self._key, 0, 0)
         return first_job[0].decode() if first_job else None
 
-    def get_last_timestamp(self) -> Optional[float]:
+    def get_last_timestamp(self) -> Optional[int]:
         """Returns the last timestamp in the registry."""
         self.cleanup(self.connection)
         last_timestamp = self.connection.zrange(self._key, -1, -1, withscores=True)
-        return last_timestamp[0][1] if last_timestamp else None
+        return int(last_timestamp[0][1]) if last_timestamp else None
 
     @property
     def key(self) -> str:

scheduler/redis_models/result.py

@@ -5,6 +5,7 @@
 
 from scheduler.helpers.utils import utcnow
 from scheduler.redis_models.base import StreamModel, decode_dict
+from scheduler.settings import logger
 from scheduler.types import ConnectionType, Self
 
 
@@ -34,19 +35,21 @@ def create(
         cls,
         connection: ConnectionType,
         job_name: str,
-        worker_name: str,
+        worker_name: Optional[str],
         _type: ResultType,
         ttl: int,
         return_value: Any = None,
         exc_string: Optional[str] = None,
     ) -> Self:
+        if worker_name is None:
+            logger.warning(f"Job {job_name} has no worker name, will save result with 'unknown_worker'")
         result = cls(
             parent=job_name,
             ttl=ttl,
             type=_type,
             return_value=return_value,
             exc_string=exc_string,
-            worker_name=worker_name,
+            worker_name=worker_name or "unknown_worker",
         )
         result.save(connection)
         return result

scheduler/types/settings_types.py

@@ -65,7 +65,7 @@ class QueueConfiguration:
     USERNAME: Optional[str] = None
     PASSWORD: Optional[str] = None
 
-    ASYNC: Optional[bool] = True
+    ASYNC: bool = True
 
     SENTINELS: Optional[List[Tuple[str, int]]] = None
     SENTINEL_KWARGS: Optional[Dict[str, str]] = None

scheduler/views/helpers.py

@@ -68,7 +68,7 @@ def _enqueue_multiple_jobs(queue: Queue, job_names: List[str], at_front: bool =
             if job is None:
                 continue
             job.save(connection=pipe)
-            queue.enqueue_job(job, connection=pipe, at_front=at_front)
+            queue.enqueue_job(job, pipeline=pipe, at_front=at_front)
             jobs_requeued += 1
         pipe.execute()
     return jobs_requeued