wip

Author: cunla

scheduler/management/commands/export.py

@@ -5,8 +5,7 @@
 from django.core.management.base import BaseCommand
 
 from scheduler.models import Task
-from scheduler.rq_classes import MODEL_NAMES
-from scheduler.tools import OLD_MODEL_NAMES
+from scheduler.tools import MODEL_NAMES
 
 
 class Command(BaseCommand):
@@ -45,12 +44,7 @@ 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()
-        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:
+        for model_name in 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 OLD_MODEL_NAMES
+from scheduler.tools import MODEL_NAMES
 
 
 def job_model_str(model_str: str) -> str:
@@ -21,6 +21,8 @@ def job_model_str(model_str: str) -> str:
 
 def get_task_type(model_str: str) -> TaskType:
     model_str = job_model_str(model_str)
+    if TaskType(model_str):
+        return TaskType(model_str)
     if model_str == "CronTask":
         return TaskType.CRON
     elif model_str == "RepeatableTask":
@@ -133,7 +135,7 @@ def handle(self, *args, **options):
             jobs = yaml.load(file, yaml.SafeLoader)
 
         if options.get("reset"):
-            for model_name in OLD_MODEL_NAMES:
+            for model_name in MODEL_NAMES:
                 model = apps.get_model(app_label="scheduler", model_name=model_name)
                 model.objects.all().delete()
 

scheduler/models/task.py

@@ -48,13 +48,13 @@ def failure_callback(job, connection, result, *args, **kwargs):
 
 
 def success_callback(job, connection, result, *args, **kwargs):
-    model_name = job.meta.get("task_type", None)
-    if model_name is None:
+    task_type = job.meta.get("task_type", None)
+    if task_type is None:
         return
 
     task = Task.objects.filter(job_id=job.id).first()
     if task is None:
-        model = apps.get_model(app_label="scheduler", model_name=model_name)
+        model = apps.get_model(app_label="scheduler", model_name=task_type)
         task = model.objects.filter(job_id=job.id).first()
     if task is None:
         logger.warn(f"Could not find task for job {job.id}")
@@ -267,7 +267,7 @@ def ready_for_schedule(self) -> bool:
         if not self.enabled:
             logger.debug(f"Task {str(self)} disabled, enable task before scheduling")
             return False
-        if self.task_type == TaskType.REPEATABLE and self._schedule_time() < timezone.now():
+        if self.task_type in {TaskType.REPEATABLE, TaskType.ONCE} and self._schedule_time() < timezone.now():
             return False
         return True
 

scheduler/rq_classes.py

@@ -24,8 +24,8 @@
 from scheduler import settings
 from scheduler.broker_types import PipelineType, ConnectionType
 
-OLD_MODEL_NAMES = ["ScheduledTask", "RepeatableTask", "CronTask"]
-MODEL_NAMES = ["OnceTask", "RepeatableTask", "CronTask"]
+MODEL_NAMES = ["ScheduledTask", "RepeatableTask", "CronTask", "Task"]
+TASK_TYPES = ["OnceTaskType", "RepeatableTaskType", "CronTaskType"]
 
 rq_job_decorator = job
 ExecutionStatus = JobStatus
@@ -264,7 +264,7 @@ def __init__(self, *args, **kwargs) -> None:
 
     @staticmethod
     def reschedule_all_jobs():
-        for model_name in OLD_MODEL_NAMES:
+        for model_name in 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_mgmt_commands/test_export.py

@@ -7,16 +7,18 @@
 from django.core.management import call_command
 from django.test import TestCase
 
+from scheduler.tests import test_settings  # noqa
 from scheduler.tests.testtools import task_factory
 from scheduler.tools import TaskType
-from scheduler.tests import test_settings # noqa
 
 
 class ExportTest(TestCase):
     def setUp(self) -> None:
+        super().setUp()
         self.tmpfile = tempfile.NamedTemporaryFile()
 
     def tearDown(self) -> None:
+        super().tearDown()
         os.remove(self.tmpfile.name)
 
     def test_export__should_export_job(self):

scheduler/tools.py

@@ -1,6 +1,6 @@
 import importlib
 import os
-from typing import List, Any, Callable, Optional, Union
+from typing import List, Any, Callable, Optional
 
 import croniter
 from django.apps import apps
@@ -10,14 +10,14 @@
 from django.utils.translation import gettext_lazy as _
 
 from scheduler.queues import get_queues, logger, get_queue
-from scheduler.rq_classes import DjangoWorker, OLD_MODEL_NAMES, JobExecution, MODEL_NAMES
+from scheduler.rq_classes import DjangoWorker, JobExecution, TASK_TYPES, MODEL_NAMES
 from scheduler.settings import SCHEDULER_CONFIG, Broker
 
 
 class TaskType(models.TextChoices):
-    CRON = "CronTask", _("Cron Task")
-    REPEATABLE = "RepeatableTask", _("Repeatable Task")
-    ONCE = "OnceTask", _("Run once")
+    CRON = "CronTaskType", _("Cron Task")
+    REPEATABLE = "RepeatableTaskType", _("Repeatable Task")
+    ONCE = "OnceTaskType", _("Run once")
 
 
 def callable_func(callable_str: str) -> Callable:
@@ -39,31 +39,25 @@ def get_next_cron_time(cron_string: Optional[str]) -> Optional[timezone.datetime
     return next_itr
 
 
-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}")
-
+def get_scheduled_task(task_type_str: str, task_id: int) -> "BaseTask":  # noqa: F821
     # 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_type_str in TASK_TYPES:
+        try:
+            task_type = TaskType(task_type_str)
+            task = model.objects.filter(task_type=TaskType.ONCE, id=task_id).first()
+            if task is None:
+                raise ValueError(f"Job {task_type}:{task_id} does not exit")
+            return task
+        except ValueError:
+            raise ValueError(f"Invalid task type {task_type_str}")
+    elif task_type_str in MODEL_NAMES:
+        model = apps.get_model(app_label="scheduler", model_name=task_type_str)
+        task = model.objects.filter(id=task_id).first()
         if task is None:
-            raise ValueError(f"Job {task_model}:{task_id} does not exit")
+            raise ValueError(f"Job {task_type_str}:{task_id} does not exit")
         return task
-    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:
-        raise ValueError(f"Job {task_model}:{task_id} does not exit")
-    return task
+    raise ValueError(f"Job Model {task_type_str} does not exist, choices are {TASK_TYPES}")
 
 
 def run_task(task_model: str, task_id: int) -> Any: