Producer for running tasks on start (#90)

* Basics of an on-start producer

* Add test for on_start producer

* Update schema

* Update docs with new produer

Author: AlanCoding

dispatcher.yml

@@ -21,5 +21,8 @@ producers:
         schedule: 3
       'lambda: __import__("time").sleep(2)':
         schedule: 3
+  OnStartProducer:
+    task_list:
+      'lambda: print("This task runs on startup")': {}
 publish:
   default_broker: pg_notify

dispatcher/producers/__init__.py

@@ -1,5 +1,6 @@
 from .base import BaseProducer
 from .brokered import BrokeredProducer
+from .on_start import OnStartProducer
 from .scheduled import ScheduledProducer
 
-__all__ = ['BaseProducer', 'BrokeredProducer', 'ScheduledProducer']
+__all__ = ['BaseProducer', 'BrokeredProducer', 'ScheduledProducer', 'OnStartProducer']

dispatcher/producers/on_start.py

@@ -0,0 +1,31 @@
+import asyncio
+import logging
+from typing import Union
+
+from .base import BaseProducer
+
+logger = logging.getLogger(__name__)
+
+
+class OnStartProducer(BaseProducer):
+    def __init__(self, task_list: dict[str, dict[str, Union[int, str]]]):
+        self.task_list = task_list
+        super().__init__()
+
+    async def start_producing(self, dispatcher) -> None:
+        self.events.ready_event.set()
+
+        for task_name, options in self.task_list.items():
+            message = options.copy()
+            message['task'] = task_name
+            message['uuid'] = f'on-start-{self.produced_count}'
+
+            logger.debug(f"Produced on-start task: {task_name}")
+            self.produced_count += 1
+            await dispatcher.process_message(message)
+
+    def all_tasks(self) -> list[asyncio.Task]:
+        return []
+
+    async def shutdown(self) -> None:
+        pass

dispatcher/producers/scheduled.py

@@ -1,13 +1,14 @@
 import asyncio
 import logging
+from typing import Union
 
 from .base import BaseProducer
 
 logger = logging.getLogger(__name__)
 
 
 class ScheduledProducer(BaseProducer):
-    def __init__(self, task_schedule: dict[str, dict[str, int]]):
+    def __init__(self, task_schedule: dict[str, dict[str, Union[int, str]]]):
         self.task_schedule = task_schedule
         self.scheduled_tasks: list[asyncio.Task] = []
         super().__init__()

docs/config.md

@@ -116,8 +116,9 @@ For every listed broker, a `BrokeredProducer` is automatically
 created. That means that tasks may be produced from the messaging
 system that the dispatcher service is listening to.
 
-The other current use case is `ScheduledProducer`,
-which submits tasks every certain number of seconds.
+The others are:
+ - `ScheduledProducer` - submits tasks every certain number of seconds
+ - `OnStartProducer` - runs tasks once after starting
 
 #### Publish
 

schema.json

@@ -11,7 +11,10 @@
   },
   "producers": {
     "ScheduledProducer": {
-      "task_schedule": "dict[str, dict[str, int]]"
+      "task_schedule": "dict[str, dict[str, typing.Union[int, str]]]"
+    },
+    "OnStartProducer": {
+      "task_list": "dict[str, dict[str, typing.Union[int, str]]]"
     }
   },
   "service": {

tests/integration/test_producers.py

@@ -0,0 +1,39 @@
+import logging
+import asyncio
+
+import pytest
+
+from dispatcher.factories import from_settings
+from dispatcher.config import DispatcherSettings
+
+
+@pytest.mark.asyncio
+async def test_on_start_tasks(caplog):
+    try:
+        settings = DispatcherSettings({
+            'version': 2,
+            'service': {
+                'pool_kwargs': {'max_workers': 2}
+            },
+            'brokers': {},  # do not need them for this test
+            'producers': {
+                'OnStartProducer': {
+                    'task_list': {'lambda: return "confirmation_of_run"': {}}
+                }
+            }
+        })
+        dispatcher = from_settings(settings=settings)
+        assert dispatcher.pool.finished_count == 0
+
+        await dispatcher.connect_signals()
+        with caplog.at_level(logging.DEBUG):
+            await dispatcher.start_working()
+            await dispatcher.wait_for_producers_ready()
+            await asyncio.wait_for(dispatcher.pool.events.work_cleared.wait(), timeout=2)
+            await asyncio.sleep(0.02)  # still may be some time between clearing event and desired log
+
+        assert dispatcher.pool.finished_count == 1
+        assert 'result: confirmation_of_run' not in caplog.text
+    finally:
+        await dispatcher.shutdown()
+        await dispatcher.cancel_tasks()