From 7beebe030d82f31b9131ea912d246775c90b388c Mon Sep 17 00:00:00 2001
From: Jens Langhammer <jens.langhammer@beryju.org>
Date: Thu, 9 Jun 2022 20:09:51 +0200
Subject: [PATCH] lifecycle: run bootstrap tasks inline when using automated
 install

Signed-off-by: Jens Langhammer <jens.langhammer@beryju.org>
---
 .../core/management/commands/bootstrap_tasks.py  | 13 +++++++++++++
 authentik/root/celery.py                         | 16 +++++++++++-----
 lifecycle/ak                                     |  3 +++
 3 files changed, 27 insertions(+), 5 deletions(-)
 create mode 100644 authentik/core/management/commands/bootstrap_tasks.py

diff --git a/authentik/core/management/commands/bootstrap_tasks.py b/authentik/core/management/commands/bootstrap_tasks.py
new file mode 100644
index 000000000..74badb61e
--- /dev/null
+++ b/authentik/core/management/commands/bootstrap_tasks.py
@@ -0,0 +1,13 @@
+"""Run bootstrap tasks"""
+from django.core.management.base import BaseCommand
+
+from authentik.root.celery import _get_startup_tasks
+
+
+class Command(BaseCommand):  # pragma: no cover
+    """Run bootstrap tasks to ensure certain objects are created"""
+
+    def handle(self, **options):
+        tasks = _get_startup_tasks()
+        for task in tasks:
+            task()
diff --git a/authentik/root/celery.py b/authentik/root/celery.py
index a0c7641fd..2b3d16c85 100644
--- a/authentik/root/celery.py
+++ b/authentik/root/celery.py
@@ -1,6 +1,7 @@
 """authentik core celery"""
 import os
 from logging.config import dictConfig
+from typing import Callable
 
 from celery import Celery
 from celery.signals import (
@@ -76,23 +77,28 @@ def task_error_hook(task_id, exception: Exception, traceback, *args, **kwargs):
         Event.new(EventAction.SYSTEM_EXCEPTION, message=exception_to_string(exception)).save()
 
 
-@worker_ready.connect
-def worker_ready_hook(*args, **kwargs):
-    """Run certain tasks on worker start"""
+def _get_startup_tasks() -> list[Callable]:
+    """Get all tasks to be run on startup"""
     from authentik.admin.tasks import clear_update_notifications
     from authentik.managed.tasks import managed_reconcile
     from authentik.outposts.tasks import outpost_controller_all, outpost_local_connection
     from authentik.providers.proxy.tasks import proxy_set_defaults
 
-    tasks = [
+    return [
         clear_update_notifications,
         outpost_local_connection,
         outpost_controller_all,
         proxy_set_defaults,
         managed_reconcile,
     ]
+
+
+@worker_ready.connect
+def worker_ready_hook(*args, **kwargs):
+    """Run certain tasks on worker start"""
+
     LOGGER.info("Dispatching startup tasks...")
-    for task in tasks:
+    for task in _get_startup_tasks():
         try:
             task.delay()
         except ProgrammingError as exc:
diff --git a/lifecycle/ak b/lifecycle/ak
index b01ff2533..c304952ff 100755
--- a/lifecycle/ak
+++ b/lifecycle/ak
@@ -38,6 +38,9 @@ if [[ "$1" == "server" ]]; then
     wait_for_db
     echo "server" > $MODE_FILE
     python -m lifecycle.migrate
+    if [[ ! -z "${AUTHENTIK_BOOTSTRAP_PASSWORD}" || ! -z "${AUTHENTIK_BOOTSTRAP_TOKEN}" ]]; then
+        python -m manage bootstrap_tasks
+    fi
     /authentik-proxy
 elif [[ "$1" == "worker" ]]; then
     wait_for_db