Restructure management commands and Celery tasks

CHANGELOG.rst

@@ -19,11 +19,8 @@ Changed
 ~~~~~~~
 
 * Wrap all imports in complete revisions to make it possible to undo them completely and to track changes correctly.
-
-Fixed
-~~~~~
-
-* Management commands were run in Celery, not in the local console.
+* Management commands can run the import in the foreground or in the background.
+* The management commands were merged to one with an argument to call the subcommands.
 
 `2.0rc3`_ - 2021-09-30
 ----------------------

aleksis/apps/untis/commands.py

+from typing import Optional
+
+from django.db.models import QuerySet
+from django.utils.functional import classproperty
+
 from aleksis.apps.untis.util.mysql.importers.terms import (
     get_future_terms_for_date,
     get_terms_for_date,
@@ -6,35 +11,82 @@ from aleksis.apps.untis.util.mysql.importers.terms import (
 from .util.mysql.main import untis_import_mysql as _untis_import_mysql
 
 
-def import_mysql_current_term():
-    """Import UNTIS data from MySQL (current term)."""
-    terms = get_terms_for_date()
-    _untis_import_mysql(terms)
+class ImportCommand:
+    """A generic UNTIS import command."""
+
+    name = None
+
+    @classproperty
+    def task_name(cls) -> str:  # noqa
+        """Get the name for the related Celery task."""
+        return f"untis_import_mysql_{cls.name}"
+
+    @classmethod
+    def get_terms(cls) -> Optional[QuerySet]:
+        """Return which terms should be imported."""
+        return None
+
+    @classmethod
+    def run(cls, background: bool = False):
+        """Run the import command (foreground/background)."""
+        if background:
+            from .tasks import TASKS
+
+            task = TASKS[cls]
+            task.delay()
+        else:
+            _untis_import_mysql(cls.get_terms())
+
+
+class CurrentImportCommand(ImportCommand):
+    """Import data of current term from UNTIS."""
+
+    name = "current"
+
+    @classmethod
+    def get_terms(cls) -> Optional[QuerySet]:
+        return get_terms_for_date()
+
+
+class FutureImportCommand(ImportCommand):
+    """Import data of future terms from UNTIS."""
+
+    name = "future"
+
+    @classmethod
+    def get_terms(cls) -> Optional[QuerySet]:
+        return get_future_terms_for_date()
+
+
+class AllImportCommand(ImportCommand):
+    name = "all"
+
+
+class CurrentNextImportCommand(ImportCommand):
+    """Import data of the current and next term from UNTIS."""
 
+    name = "current_next"
 
-def import_mysql_future_terms():
-    """Import UNTIS data from MySQL (all future terms)."""
-    terms = get_future_terms_for_date()
-    _untis_import_mysql(terms)
+    @classmethod
+    def get_terms(cls) -> Optional[QuerySet]:
+        terms = get_terms_for_date()
+        future_terms = get_future_terms_for_date()
+        if future_terms.exists():
+            terms = terms.union(future_terms[0:1])
+        return terms
 
 
-def import_mysql_all_terms():
-    """Import UNTIS data from MySQL (all terms in DB)."""
-    _untis_import_mysql()
+class CurrentFutureImportCommand(ImportCommand):
+    """Import data of the current and future terms from UNTIS."""
 
+    name = "current_future"
 
-def import_mysql_current_next_term():
-    """Import UNTIS data from MySQL (current and next term)."""
-    terms = get_terms_for_date()
-    future_terms = get_future_terms_for_date()
-    if future_terms.exists():
-        terms = terms.union(future_terms[0:1])
-    _untis_import_mysql(terms)
+    @classmethod
+    def get_terms(cls) -> Optional[QuerySet]:
+        terms = get_terms_for_date()
+        future_terms = get_future_terms_for_date()
+        terms = terms.union(future_terms)
+        return terms
 
 
-def import_mysql_current_future_terms():
-    """Import UNTIS data from MySQL (current and future terms)."""
-    terms = get_terms_for_date()
-    future_terms = get_future_terms_for_date()
-    terms = terms.union(future_terms)
-    _untis_import_mysql(terms)
+COMMANDS_BY_NAME = {c.name: c for c in ImportCommand.__subclasses__()}

aleksis/apps/untis/management/commands/untis_import_mysql.py

 from django.core.management.base import BaseCommand
 
-from ...commands import import_mysql_current_term
+from ...commands import COMMANDS_BY_NAME
 
 
 class Command(BaseCommand):
+    def add_arguments(self, parser):
+        parser.add_argument(
+            "command", nargs="?", default="all", type=str, choices=list(COMMANDS_BY_NAME.keys())
+        )
+        parser.add_argument(
+            "--background", action="store_true", help="Run import job in background using Celery",
+        )
+
     def handle(self, *args, **options):
-        import_mysql_current_term()
+        command = COMMANDS_BY_NAME[options["command"]]
+        background = options["background"]
+        command.run(background=background)

aleksis/apps/untis/management/commands/untis_import_mysql_all.py

-from django.core.management.base import BaseCommand
-
-from ...commands import import_mysql_all_terms
-
-
-class Command(BaseCommand):
-    def handle(self, *args, **options):
-        import_mysql_all_terms()

aleksis/apps/untis/management/commands/untis_import_mysql_current_future.py

-from django.core.management.base import BaseCommand
-
-from ...commands import import_mysql_current_future_terms
-
-
-class Command(BaseCommand):
-    def handle(self, *args, **options):
-        import_mysql_current_future_terms()

aleksis/apps/untis/management/commands/untis_import_mysql_current_next.py

-from django.core.management.base import BaseCommand
-
-from ...commands import import_mysql_current_next_term
-
-
-class Command(BaseCommand):
-    def handle(self, *args, **options):
-        import_mysql_current_next_term()

aleksis/apps/untis/management/commands/untis_import_mysql_future.py

-from django.core.management.base import BaseCommand
-
-from ...commands import import_mysql_future_terms
-
-
-class Command(BaseCommand):
-    def handle(self, *args, **options):
-        import_mysql_future_terms()

aleksis/apps/untis/tasks.py

 from aleksis.core.celery import app
 
-from .commands import (
-    import_mysql_all_terms,
-    import_mysql_current_future_terms,
-    import_mysql_current_next_term,
-    import_mysql_current_term,
-    import_mysql_future_terms,
-)
+from .commands import ImportCommand
 
+TASKS = {}
+for import_command in ImportCommand.__subclasses__():
 
-@app.task
-def untis_import_mysql_current_term():
-    """Celery task for import of UNTIS data from MySQL (current term)."""
-    import_mysql_current_term()
+    @app.task(name=import_command.task_name)
+    def _task():
+        import_command.run()
 
-
-@app.task
-def untis_import_mysql_future_terms():
-    """Celery task for import of UNTIS data from MySQL (all future terms)."""
-    import_mysql_future_terms()
-
-
-@app.task
-def untis_import_mysql_all_terms():
-    """Celery task for import of UNTIS data from MySQL (all terms in DB)."""
-    import_mysql_all_terms()
-
-
-@app.task
-def untis_import_mysql_current_next_term():
-    """Celery task for import of UNTIS data from MySQL (current and next term)."""
-    import_mysql_current_next_term()
-
-
-@app.task
-def untis_import_mysql_current_future_terms():
-    """Celery task for import of UNTIS data from MySQL (current and future terms)."""
-    import_mysql_current_future_terms()
+    TASKS[import_command] = _task