wip

Author: cunla

scheduler/management/commands/export.py

@@ -4,7 +4,9 @@
 from django.apps import apps
 from django.core.management.base import BaseCommand
 
-from scheduler.tools import MODEL_NAMES
+from scheduler.models import Task
+from scheduler.rq_classes import MODEL_NAMES
+from scheduler.tools import OLD_MODEL_NAMES
 
 
 class Command(BaseCommand):
@@ -43,7 +45,12 @@ def add_arguments(self, parser):
     def handle(self, *args, **options):
         file = open(options.get("filename"), "w") if options.get("filename") else sys.stdout
         res = list()
-        for model_name in MODEL_NAMES:
+        jobs = Task.objects.all()
+        if options.get("enabled"):
+            jobs = jobs.filter(enabled=True)
+        for job in jobs:
+            res.append(job.to_dict())
+        for model_name in OLD_MODEL_NAMES:
             model = apps.get_model(app_label="scheduler", model_name=model_name)
             jobs = model.objects.all()
             if options.get("enabled"):

scheduler/management/commands/import.py

@@ -10,7 +10,7 @@
 
 from scheduler.models import TaskArg, TaskKwarg, Task
 from scheduler.models.task import TaskType
-from scheduler.tools import MODEL_NAMES
+from scheduler.tools import OLD_MODEL_NAMES
 
 
 def job_model_str(model_str: str) -> str:
@@ -133,7 +133,7 @@ def handle(self, *args, **options):
             jobs = yaml.load(file, yaml.SafeLoader)
 
         if options.get("reset"):
-            for model_name in MODEL_NAMES:
+            for model_name in OLD_MODEL_NAMES:
                 model = apps.get_model(app_label="scheduler", model_name=model_name)
                 model.objects.all().delete()
 

scheduler/models/task.py

@@ -51,9 +51,13 @@ def success_callback(job, connection, result, *args, **kwargs):
     model_name = job.meta.get("task_type", None)
     if model_name is None:
         return
-    model = apps.get_model(app_label="scheduler", model_name=model_name)
-    task = model.objects.filter(job_id=job.id).first()
+
+    task = Task.objects.filter(job_id=job.id).first()
     if task is None:
+        model = apps.get_model(app_label="scheduler", model_name=model_name)
+        task = model.objects.filter(job_id=job.id).first()
+    if task is None:
+        logger.warn(f"Could not find task for job {job.id}")
         return
     task.job_id = None
     task.successful_runs += 1

scheduler/rq_classes.py

@@ -24,7 +24,8 @@
 from scheduler import settings
 from scheduler.broker_types import PipelineType, ConnectionType
 
-MODEL_NAMES = ["ScheduledTask", "RepeatableTask", "CronTask", "Task"]
+OLD_MODEL_NAMES = ["ScheduledTask", "RepeatableTask", "CronTask"]
+MODEL_NAMES = ["OnceTask", "RepeatableTask", "CronTask"]
 
 rq_job_decorator = job
 ExecutionStatus = JobStatus
@@ -63,7 +64,8 @@ def is_scheduled_task(self) -> bool:
 
     def is_execution_of(self, task: "Task") -> bool:  # noqa: F821
         return (
-            self.meta.get("task_type", None) == task.task_type and self.meta.get("scheduled_task_id", None) == task.id
+                self.meta.get("task_type", None) == task.task_type and self.meta.get("scheduled_task_id",
+                                                                                     None) == task.id
         )
 
     def stop_execution(self, connection: ConnectionType):
@@ -92,11 +94,11 @@ def __str__(self):
         return f"{self.name}/{','.join(self.queue_names())}"
 
     def _start_scheduler(
-        self,
-        burst: bool = False,
-        logging_level: str = "INFO",
-        date_format: str = "%H:%M:%S",
-        log_format: str = "%(asctime)s %(message)s",
+            self,
+            burst: bool = False,
+            logging_level: str = "INFO",
+            date_format: str = "%H:%M:%S",
+            log_format: str = "%(asctime)s %(message)s",
     ) -> None:
         """Starts the scheduler process.
         This is specifically designed to be run by the worker when running the `work()` method.
@@ -262,7 +264,7 @@ def __init__(self, *args, **kwargs) -> None:
 
     @staticmethod
     def reschedule_all_jobs():
-        for model_name in MODEL_NAMES:
+        for model_name in OLD_MODEL_NAMES:
             model = apps.get_model(app_label="scheduler", model_name=model_name)
             enabled_jobs = model.objects.filter(enabled=True)
             unscheduled_jobs = filter(lambda j: j.ready_for_schedule(), enabled_jobs)

scheduler/tests/test_task_types/test_once_task.py

@@ -0,0 +1,22 @@
+from datetime import timedelta
+
+from django.utils import timezone
+
+from scheduler import settings
+from scheduler.models.task import TaskType
+from scheduler.tests.test_task_types.test_task_model import BaseTestCases
+from scheduler.tests.testtools import task_factory
+
+
+class TestScheduledTask(BaseTestCases.TestSchedulableTask):
+    task_type = TaskType.ONCE
+
+    def test_clean(self):
+        job = task_factory(self.task_type)
+        job.queue = list(settings.QUEUES)[0]
+        job.callable = "scheduler.tests.jobs.test_job"
+        self.assertIsNone(job.clean())
+
+    def test_unschedulable_old_job(self):
+        job = task_factory(self.task_type, scheduled_time=timezone.now() - timedelta(hours=1))
+        self.assertFalse(job.is_scheduled())

scheduler/tests/test_task_types/test_task_model.py

@@ -507,16 +507,3 @@ def test_result_ttl_passthrough(self):
             entry = _get_task_job_execution_from_registry(job)
             self.assertEqual(entry.result_ttl, 500)
 
-
-class TestScheduledTask(BaseTestCases.TestSchedulableTask):
-    task_type = TaskType.ONCE
-
-    def test_clean(self):
-        job = task_factory(self.task_type)
-        job.queue = list(settings.QUEUES)[0]
-        job.callable = "scheduler.tests.jobs.test_job"
-        self.assertIsNone(job.clean())
-
-    def test_unschedulable_old_job(self):
-        job = task_factory(self.task_type, scheduled_time=timezone.now() - timedelta(hours=1))
-        self.assertFalse(job.is_scheduled())

scheduler/tools.py

@@ -10,7 +10,7 @@
 from django.utils.translation import gettext_lazy as _
 
 from scheduler.queues import get_queues, logger, get_queue
-from scheduler.rq_classes import DjangoWorker, MODEL_NAMES, JobExecution
+from scheduler.rq_classes import DjangoWorker, OLD_MODEL_NAMES, JobExecution, MODEL_NAMES
 from scheduler.settings import SCHEDULER_CONFIG, Broker
 
 
@@ -39,15 +39,26 @@ def get_next_cron_time(cron_string: Optional[str]) -> Optional[timezone.datetime
     return next_itr
 
 
-def get_scheduled_task(task_model: Union[TaskType, str], task_id: int) -> "BaseTask":  # noqa: F821
-    if isinstance(task_model, TaskType):
-        model = apps.get_model(app_label="scheduler", model_name="Task")
-        task = model.objects.filter(task_type=task_model, id=task_id).first()
+def get_scheduled_task(task_model: str, task_id: int) -> "BaseTask":  # noqa: F821
+    if isinstance(task_model, str) and task_model not in OLD_MODEL_NAMES and task_model not in MODEL_NAMES:
+        raise ValueError(f"Job Model `{task_model}` does not exist, choices are {OLD_MODEL_NAMES}")
+
+    # Try with new model names
+    model = apps.get_model(app_label="scheduler", model_name="Task")
+    if task_model == "OnceTask":
+        task = model.objects.filter(task_type=TaskType.ONCE, id=task_id).first()
         if task is None:
             raise ValueError(f"Job {task_model}:{task_id} does not exit")
         return task
-    if isinstance(task_model, str) and task_model not in MODEL_NAMES:
-        raise ValueError(f"Job Model `{task_model}` does not exist, choices are {MODEL_NAMES}")
+    elif task_model == "RepeatableTask":
+        task = model.objects.filter(task_type=TaskType.REPEATABLE, id=task_id).first()
+        if task is not None:
+            return task
+    elif task_model == "CronTask":
+        task = model.objects.filter(task_type=TaskType.CRON, id=task_id).first()
+        if task is not None:
+            return task
+
     model = apps.get_model(app_label="scheduler", model_name=task_model)
     task = model.objects.filter(id=task_id).first()
     if task is None:
@@ -81,7 +92,7 @@ def create_worker(*queue_names, **kwargs) -> DjangoWorker:
     queues = get_queues(*queue_names)
     existing_workers = DjangoWorker.all(connection=queues[0].connection)
     existing_worker_names = set(map(lambda w: w.name, existing_workers))
-    kwargs["fork_job_execution"] = SCHEDULER_CONFIG.BROKER != Broker.FAKEREDIS
+    kwargs.setdefault("fork_job_execution", SCHEDULER_CONFIG.BROKER != Broker.FAKEREDIS)
     if kwargs.get("name", None) is None:
         kwargs["name"] = _calc_worker_name(existing_worker_names)