From 1033b114c73f5d64897ca7c885ab9861f7f6d5a7 Mon Sep 17 00:00:00 2001
From: Jens Langhammer <jens@goauthentik.io>
Date: Sat, 13 Jan 2024 20:54:42 +0100
Subject: [PATCH] prefill in app startup

Signed-off-by: Jens Langhammer <jens@goauthentik.io>
---
 authentik/blueprints/apps.py        | 10 ++++-----
 authentik/events/apps.py            | 11 ++++++++++
 authentik/events/monitored_tasks.py | 34 ++++++++++++++---------------
 3 files changed, 32 insertions(+), 23 deletions(-)

diff --git a/authentik/blueprints/apps.py b/authentik/blueprints/apps.py
index aba14d552..ddc31e314 100644
--- a/authentik/blueprints/apps.py
+++ b/authentik/blueprints/apps.py
@@ -11,11 +11,11 @@ from structlog.stdlib import BoundLogger, get_logger
 class ManagedAppConfig(AppConfig):
     """Basic reconciliation logic for apps"""
 
-    _logger: BoundLogger
+    logger: BoundLogger
 
     def __init__(self, app_name: str, *args, **kwargs) -> None:
         super().__init__(app_name, *args, **kwargs)
-        self._logger = get_logger().bind(app_name=app_name)
+        self.logger = get_logger().bind(app_name=app_name)
 
     def ready(self) -> None:
         self.reconcile()
@@ -36,11 +36,11 @@ class ManagedAppConfig(AppConfig):
                 continue
             name = meth_name.replace(prefix, "")
             try:
-                self._logger.debug("Starting reconciler", name=name)
+                self.logger.debug("Starting reconciler", name=name)
                 meth()
-                self._logger.debug("Successfully reconciled", name=name)
+                self.logger.debug("Successfully reconciled", name=name)
             except (DatabaseError, ProgrammingError, InternalError) as exc:
-                self._logger.warning("Failed to run reconcile", name=name, exc=exc)
+                self.logger.warning("Failed to run reconcile", name=name, exc=exc)
 
 
 class AuthentikBlueprintsConfig(ManagedAppConfig):
diff --git a/authentik/events/apps.py b/authentik/events/apps.py
index 3d939d568..cd79b54bd 100644
--- a/authentik/events/apps.py
+++ b/authentik/events/apps.py
@@ -43,3 +43,14 @@ class AuthentikEventsConfig(ManagedAppConfig):
                 replacement_env=replace_env,
                 message=msg,
             ).save()
+
+    def reconcile_prefill_tasks(self):
+        """Prefill tasks"""
+        from authentik.events.models import SystemTask
+        from authentik.events.monitored_tasks import _prefill_tasks
+
+        for task in _prefill_tasks:
+            if SystemTask.objects.filter(name=task.name).exists():
+                continue
+            task.save()
+            self.logger.debug("prefilled task", task_name=task.name)
diff --git a/authentik/events/monitored_tasks.py b/authentik/events/monitored_tasks.py
index 3ec5ab3a2..43eb63719 100644
--- a/authentik/events/monitored_tasks.py
+++ b/authentik/events/monitored_tasks.py
@@ -4,12 +4,12 @@ from timeit import default_timer
 from typing import Any, Optional
 
 from celery import Task
-from django.db import DatabaseError, InternalError, ProgrammingError
 from django.utils.timezone import now
 from django.utils.translation import gettext_lazy as _
 from structlog.stdlib import get_logger
 
 from authentik.events.models import Event, EventAction, SystemTask, TaskStatus
+from authentik.events.utils import sanitize_item
 from authentik.lib.utils.errors import exception_to_string
 
 LOGGER = get_logger()
@@ -78,7 +78,7 @@ class MonitoredTask(Task):
                 "task_call_args": args,
                 "task_call_kwargs": kwargs,
                 "status": self._status,
-                "messages": self._messages,
+                "messages": sanitize_item(self._messages),
                 "expires": now() + timedelta(hours=self.result_timeout_hours),
                 "expiring": True,
             },
@@ -104,7 +104,7 @@ class MonitoredTask(Task):
                 "task_call_args": args,
                 "task_call_kwargs": kwargs,
                 "status": self._status,
-                "messages": self._messages,
+                "messages": sanitize_item(self._messages),
                 "expires": now() + timedelta(hours=self.result_timeout_hours),
                 "expiring": True,
             },
@@ -120,20 +120,18 @@ class MonitoredTask(Task):
 
 def prefill_task(func):
     """Ensure a task's details are always in cache, so it can always be triggered via API"""
-    try:
-        status = SystemTask.objects.filter(name=func.__name__).first()
-    except (DatabaseError, InternalError, ProgrammingError):
-        return func
-    if status:
-        return func
-    SystemTask.objects.create(
-        name=func.__name__,
-        description=func.__doc__,
-        status=TaskStatus.UNKNOWN,
-        messages=[_("Task has not been run yet.")],
-        task_call_module=func.__module__,
-        task_call_func=func.__name__,
-        expiring=False,
+    _prefill_tasks.append(
+        SystemTask(
+            name=func.__name__,
+            description=func.__doc__,
+            status=TaskStatus.UNKNOWN,
+            messages=sanitize_item([_("Task has not been run yet.")]),
+            task_call_module=func.__module__,
+            task_call_func=func.__name__,
+            expiring=False,
+        )
     )
-    LOGGER.debug("prefilled task", task_name=func.__name__)
     return func
+
+
+_prefill_tasks = []