Move task registration to the CeleryQueue class

Signed-off-by: Roberto Rosario <roberto.rosario.gonzalez@gmail.com>

mayan/apps/task_manager/classes.py

@@ -1,13 +1,23 @@
 from __future__ import absolute_import, unicode_literals
 
 from datetime import timedelta
+from importlib import import_module
+import logging
+
+from kombu import Exchange, Queue
 
 from celery.five import monotonic
 from celery.task.control import inspect
 
+from django.apps import apps
+
 from django.utils.encoding import force_text, python_2_unicode_compatible
 from django.utils.timezone import now
 
+from mayan.celery import app as celery_app
+
+logger = logging.getLogger(__name__)
+
 
 @python_2_unicode_compatible
 class TaskType(object):
@@ -21,9 +31,11 @@ class TaskType(object):
     def get(cls, name):
         return cls._registry[name]
 
-    def __init__(self, name, label):
-        self.name = name
+    def __init__(self, dotted_path, label, name=None, schedule=None):
+        self.name = name or dotted_path.split('.')[-1]
         self.label = label
+        self.dotted_path = dotted_path
+        self.schedule = schedule
         self.__class__._registry[name] = self
 
     def __str__(self):
@@ -52,6 +64,20 @@ class CeleryQueue(object):
     _registry = {}
     _inspect_instance = inspect()
 
+    @staticmethod
+    def initialize():
+        for app in apps.get_app_configs():
+            try:
+                import_module('{}.queues'.format(app.name))
+            except ImportError as exception:
+                if force_text(exception) not in ('No module named queues', 'No module named \'{}.queues\''.format(app.name)):
+                    logger.error(
+                        'Error importing %s queues.py file; %s', app.name,
+                        exception
+                    )
+
+        CeleryQueue.update_celery()
+
     @classmethod
     def all(cls):
         return sorted(
@@ -62,17 +88,32 @@ class CeleryQueue(object):
     def get(cls, queue_name):
         return cls._registry[queue_name]
 
-    def __init__(self, name, label, is_default_queue=False, transient=False):
+    @classmethod
+    def update_celery(cls):
+        for instance in cls.all():
+            instance._update_celery()
+
+    def __init__(self, name, label, default_queue=False, transient=False):
         self.name = name
         self.label = label
-        self.is_default_queue = is_default_queue
-        self.is_transient = transient
+        self.default_queue = default_queue
+        self.transient = transient
         self.task_types = []
         self.__class__._registry[name] = self
 
     def __str__(self):
         return force_text(self.label)
 
+    def _process_task_dictionary(self, task_dictionary):
+        result = []
+        for worker, tasks in task_dictionary.items():
+            for task in tasks:
+                if task['delivery_info']['routing_key'] == self.name:
+                    task_type = TaskType.get(name=task['name'])
+                    result.append(Task(task_type=task_type, kwargs=task))
+
+        return result
+
     def add_task_type(self, *args, **kwargs):
         self.task_types.append(TaskType(*args, **kwargs))
 
@@ -91,12 +132,35 @@ class CeleryQueue(object):
             task_dictionary=self.__class__._inspect_instance.scheduled()
         )
 
-    def _process_task_dictionary(self, task_dictionary):
-        result = []
-        for worker, tasks in task_dictionary.items():
-            for task in tasks:
-                if task['delivery_info']['routing_key'] == self.name:
-                    task_type = TaskType.get(name=task['name'])
-                    result.append(Task(task_type=task_type, kwargs=task))
+    def _update_celery(self):
+        kwargs = {
+            'name': self.name, 'exchange': Exchange(self.name),
+            'routing_key': self.name
+        }
 
-        return result
+        if self.transient:
+            kwargs['delivery_mode'] = 1
+
+        celery_app.conf.CELERY_QUEUES.append(Queue(**kwargs))
+
+        if self.default_queue:
+            celery_app.conf.CELERY_DEFAULT_QUEUE = self.name
+
+        for task_type in self.task_types:
+            celery_app.conf.CELERY_ROUTES.update(
+                {
+                    task_type.dotted_path: {
+                        'queue': self.name
+                    },
+                }
+            )
+
+            if task_type.schedule:
+                celery_app.conf.CELERYBEAT_SCHEDULE.update(
+                    {
+                        task_type.name: {
+                            'task': task_type.dotted_path,
+                            'schedule': task_type.schedule
+                        },
+                    }
+                )