tasks results

Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>

packages/django-dramatiq-postgres/django_dramatiq_postgres/apps.py

@@ -2,6 +2,7 @@ import dramatiq
 from django.apps import AppConfig
 from django.core.exceptions import ImproperlyConfigured
 from django.utils.module_loading import import_string
+from dramatiq.results.middleware import Results
 
 from django_dramatiq_postgres.conf import Conf
 
@@ -22,20 +23,21 @@ class DjangoDramatiqPostgres(AppConfig):
         encoder: dramatiq.encoder.Encoder = import_string(Conf().encoder_class)()
         dramatiq.set_encoder(encoder)
 
-        broker_args = Conf().broker_args
-        broker_kwargs = {
-            **Conf().broker_kwargs,
-            "middleware": [],
-        }
         broker: dramatiq.broker.Broker = import_string(Conf().broker_class)(
-            *broker_args,
-            **broker_kwargs,
+            *Conf().broker_args,
+            **Conf().broker_kwargs,
+            middleware=[],
         )
 
         for middleware_class, middleware_kwargs in Conf().middlewares:
             middleware: dramatiq.middleware.middleware.Middleware = import_string(middleware_class)(
                 **middleware_kwargs,
             )
+            if isinstance(middleware, Results):
+                middleware.backend = import_string(Conf().result_backend)(
+                    *Conf().result_backend_args,
+                    **Conf().result_backend_kwargs,
+                )
             broker.add_middleware(middleware)
 
         dramatiq.set_broker(broker)

packages/django-dramatiq-postgres/django_dramatiq_postgres/conf.py

@@ -66,6 +66,18 @@ class Conf:
         # 30 days
         return self.conf.get("task_expiration", 60 * 60 * 24 * 30)
 
+    @property
+    def result_backend(self) -> str:
+        return self.conf.get("result_backend", "django_dramatiq_postgres.results.PostgresBackend")
+
+    @property
+    def result_backend_args(self) -> tuple[Any]:
+        return self.conf.get("result_backend_args", ())
+
+    @property
+    def result_backend_kwargs(self) -> dict[str, Any]:
+        return self.conf.get("result_backend_kwargs", {})
+
     @property
     def autodiscovery(self) -> dict[str, Any]:
         autodiscovery = {

packages/django-dramatiq-postgres/django_dramatiq_postgres/results.py

@@ -0,0 +1,43 @@
+from django.db import DEFAULT_DB_ALIAS
+from django.db.models import QuerySet
+from django.utils import timezone
+from django.utils.functional import cached_property
+from django.utils.module_loading import import_string
+from dramatiq.message import Message
+from dramatiq.results.backend import Missing, MResult, Result, ResultBackend
+
+from django_dramatiq_postgres.conf import Conf
+from django_dramatiq_postgres.models import TaskBase
+
+
+class PostgresBackend(ResultBackend):
+    def __init__(self, *args, db_alias: str = DEFAULT_DB_ALIAS, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.db_alias = db_alias
+
+    @cached_property
+    def model(self) -> type[TaskBase]:
+        return import_string(Conf().task_model)
+
+    @property
+    def query_set(self) -> QuerySet:
+        return self.model.objects.using(self.db_alias)
+
+    def build_message_key(self, message: Message) -> str:
+        return str(message.message_id)
+
+    def _get(self, message_key: str) -> MResult:
+        message = self.query_set.filter(message_id=message_key).first()
+        if message is None:
+            return Missing
+        data = message.result
+        if data is None:
+            return Missing
+        return self.encoder.decode(data)
+
+    def _store(self, message_key: str, result: Result, ttl: int) -> None:
+        self.query_set.filter(message_id=message_key).update(
+            mtime=timezone.now(),
+            result=self.encoder.encode(result),
+            result_expiry=timezone.now() + timezone.timedelta(milliseconds=ttl),
+        )