Celery fixes

Lars Solberg · Lars Solberg · commit b5d5cf6867d7 · 2020-04-10T23:28:14.000+02:00
diff --git a/.gitignore b/.gitignore
@@ -12,3 +12,4 @@ api/data/opa/tests/*.log
 .coverage
 api/data/opa/tests/coverage.xml
 .ropeproject
+celerybeat-schedule
diff --git a/api/data/opa/core/plugin.py b/api/data/opa/core/plugin.py
@@ -429,7 +429,12 @@ def get_plugin_manager() -> PluginManager:
 
 
 def get_component(name: str):
+    try:
     return plugin_manager.optional_components[name]
+    except KeyError:
+        raise Exception(
+            f'Component is not defined (yet?). Defined components are: {list(plugin_manager.optional_components.keys())}'
+        )
 
 
 def get_instance(name: str):
diff --git a/api/data/opa/plugins/driver_celery.py b/api/data/opa/plugins/driver_celery.py
@@ -14,8 +14,7 @@ class celery_setup(Hook):
     name = 'driver.celery.setup'
     order = -1
 
-    def run(self, celery_app, task_candidates):
-        celery_app.autodiscover_tasks(task_candidates)
+    def run(self, celery_app):
         return celery_app
 
 
@@ -33,6 +32,7 @@ def connect(self):
             "tasks", backend=self.opts.BACKEND_URL, broker=self.opts.BROKER_URL,
         )
 
-        celery_app = call_hook('driver.celery.setup', celery_app=celery_app, task_candidates=self.pm.store['task_candidates'])
+        celery_app.autodiscover_tasks(self.pm.store['task_candidates'])
+        celery_app = call_hook('driver.celery.setup', celery_app=celery_app)
 
         self.instance = celery_app
diff --git a/api/root/etc/services.d/worker/run b/api/root/etc/services.d/worker/run
@@ -4,4 +4,4 @@
 
 cd /data
 
-exec celery worker -A opa.main ${CELERY_PARAMS}
+exec celery ${CELERY_COMMAND:-worker} -A opa.main ${CELERY_PARAMS}
diff --git a/examples/docker-compose/celery-task/docker-compose.yaml b/examples/docker-compose/celery-task/docker-compose.yaml
@@ -10,18 +10,39 @@ services:
     volumes:
       - ./plugins:/plugins
 
-  worker:
+  worker-math:
     image: opastack/api:latest
     environment:
       OPA_PLUGIN_PATHS: "/plugins"
+      OPA_LOGLEVEL: debug
       MODE: "worker"
+      CELERY_PARAMS: -E -Q math
     volumes:
       - ./plugins:/plugins
 
+  worker-counter:
+    image: opastack/api:latest
+    environment:
+      OPA_PLUGIN_PATHS: "/plugins"
+      MODE: "worker"
+      CELERY_PARAMS: -E -Q counter
+    volumes:
+      - ./plugins:/plugins
+
+  # ### untested
+  # worker-beat:
+  #   image: opastack/api:latest
+  #   environment:
+  #     OPA_PLUGIN_PATHS: "/plugins"
+  #     MODE: "worker"
+  #     CELERY_COMMAND: beat
+  #   volumes:
+  #     - ./plugins:/plugins
+
   redis:
     image: "redis:5"
 
   rabbitmq:
     image: "rabbitmq:3.8-management"
-    ports:
-      - "15672:15672" # Management port
+    # ports:
+    #   - "15672:15672" # Management port
diff --git a/examples/docker-compose/celery-task/plugins/celerydemo/__init__.py b/examples/docker-compose/celery-task/plugins/celerydemo/__init__.py
@@ -4,22 +4,31 @@
 class celery_config(Hook):
     name = 'driver.celery.setup'
 
-    def run(self, celery_app, task_candidates):
-        celery_app.conf.task_routes = {"worker.celery_worker.test_celery": "test-queue"}
+    def run(self, celery_app):
+        celery_app.conf.task_routes = {
+            'celerydemo.tasks.counter': 'counter',
+            'celerydemo.tasks.divider': 'math',
+        }
         celery_app.conf.update(task_track_started=True)
-        celery_app.autodiscover_tasks(task_candidates)
         return celery_app
 
 
 router = get_router()
 
 
-@router.get("/add/{num1}/{num2}")
-async def root(num1: int, num2: int):
-    from celerydemo.tasks import test_celery
+@router.get('/inc')
+def counter(num1: int, num2: int):
+    from celerydemo.tasks import counter
 
-    celery = get_instance('celery')
     walrus = get_instance('walrus')
-    test_celery.delay('abc')
-    count = str(walrus.get('celery'))
-    return {"message": "Word received", 'count': count}
+    current_count = str(walrus.get('celery'))
+    counter.delay()
+    return {'status': 'queued', 'current_count': current_count}
+
+
+@router.get('/div/{num1}/{num2}')
+def divider(num1: int, num2: int):
+    from celerydemo.tasks import divider
+
+    divider.delay(num1, num2)
+    return {'status': 'queued'}
diff --git a/examples/docker-compose/celery-task/plugins/celerydemo/tasks.py b/examples/docker-compose/celery-task/plugins/celerydemo/tasks.py
@@ -6,12 +6,21 @@
 
 celery = get_instance('celery')
 
+
 @celery.task
-def test_celery(word: str) -> str:
-    for i in range(1, 4):
+def counter() -> str:
+
+    for i in range(4):
         sleep(1)
-        current_task.update_state(state='PROGRESS',
-                                  meta={'process_percent': i*10})
+        current_task.update_state(state='PROGRESS', meta={'process_percent': i * 20})
+
+    status = get_instance('walrus').incr('celery')
+    return f'incremented to {status}'
+
+
+@celery.task
+def divider(num1, num2) -> str:
+    sleep(4)
 
-    get_instance('walrus').incr('celery')
-    return f"test task return {word}...."
+    result = num1 / num2
+    return f'The result was {result}'

Original file line number	Diff line number	Diff line change
`@@ -4,4 +4,4 @@`
`4`	`4`
`5`	`5`	`cd /data`
`6`	`6`
`7`		`-exec celery worker -A opa.main ${CELERY_PARAMS}`
	`7`	`+exec celery ${CELERY_COMMAND:-worker} -A opa.main ${CELERY_PARAMS}`