From 341d866c00a9b9d1422e41e6b692fb028a70cdf7 Mon Sep 17 00:00:00 2001 From: Jens L Date: Fri, 23 Feb 2024 15:12:34 +0100 Subject: [PATCH] blueprints: use reconcile decorator instead of relying on function name prefix (#8483) Signed-off-by: Jens Langhammer --- authentik/blueprints/apps.py | 38 +++++++++++++++++++++--------- authentik/core/apps.py | 6 +++-- authentik/crypto/apps.py | 6 +++-- authentik/enterprise/audit/apps.py | 3 ++- authentik/events/apps.py | 9 ++++--- authentik/flows/apps.py | 3 ++- authentik/outposts/apps.py | 3 ++- authentik/sources/oauth/apps.py | 3 ++- authentik/tenants/apps.py | 3 ++- authentik/tenants/models.py | 2 +- 10 files changed, 52 insertions(+), 24 deletions(-) diff --git a/authentik/blueprints/apps.py b/authentik/blueprints/apps.py index 695bc68ce1..9a9fb53dfb 100644 --- a/authentik/blueprints/apps.py +++ b/authentik/blueprints/apps.py @@ -2,6 +2,7 @@ from importlib import import_module from inspect import ismethod +from typing import Callable from django.apps import AppConfig from django.db import DatabaseError, InternalError, ProgrammingError @@ -13,8 +14,8 @@ class ManagedAppConfig(AppConfig): logger: BoundLogger - RECONCILE_GLOBAL_PREFIX: str = "reconcile_global_" - RECONCILE_TENANT_PREFIX: str = "reconcile_tenant_" + RECONCILE_GLOBAL_CATEGORY: str = "global" + RECONCILE_TENANT_CATEGORY: str = "tenant" def __init__(self, app_name: str, *args, **kwargs) -> None: super().__init__(app_name, *args, **kwargs) @@ -22,8 +23,8 @@ class ManagedAppConfig(AppConfig): def ready(self) -> None: self.import_related() - self.reconcile_global() - self.reconcile_tenant() + self._reconcile_global() + self._reconcile_tenant() return super().ready() def import_related(self): @@ -51,7 +52,8 @@ class ManagedAppConfig(AppConfig): meth = getattr(self, meth_name) if not ismethod(meth): continue - if not meth_name.startswith(prefix): + category = getattr(meth, "_authentik_managed_reconcile", None) + if category != prefix: continue name = meth_name.replace(prefix, "") try: @@ -61,7 +63,19 @@ class ManagedAppConfig(AppConfig): except (DatabaseError, ProgrammingError, InternalError) as exc: self.logger.warning("Failed to run reconcile", name=name, exc=exc) - def reconcile_tenant(self) -> None: + @staticmethod + def reconcile_tenant(func: Callable): + """Mark a function to be called on startup (for each tenant)""" + setattr(func, "_authentik_managed_reconcile", ManagedAppConfig.RECONCILE_TENANT_CATEGORY) + return func + + @staticmethod + def reconcile_global(func: Callable): + """Mark a function to be called on startup (globally)""" + setattr(func, "_authentik_managed_reconcile", ManagedAppConfig.RECONCILE_GLOBAL_CATEGORY) + return func + + def _reconcile_tenant(self) -> None: """reconcile ourselves for tenanted methods""" from authentik.tenants.models import Tenant @@ -72,9 +86,9 @@ class ManagedAppConfig(AppConfig): return for tenant in tenants: with tenant: - self._reconcile(self.RECONCILE_TENANT_PREFIX) + self._reconcile(self.RECONCILE_TENANT_CATEGORY) - def reconcile_global(self) -> None: + def _reconcile_global(self) -> None: """ reconcile ourselves for global methods. Used for signals, tasks, etc. Database queries should not be made in here. @@ -82,7 +96,7 @@ class ManagedAppConfig(AppConfig): from django_tenants.utils import get_public_schema_name, schema_context with schema_context(get_public_schema_name()): - self._reconcile(self.RECONCILE_GLOBAL_PREFIX) + self._reconcile(self.RECONCILE_GLOBAL_CATEGORY) class AuthentikBlueprintsConfig(ManagedAppConfig): @@ -93,11 +107,13 @@ class AuthentikBlueprintsConfig(ManagedAppConfig): verbose_name = "authentik Blueprints" default = True - def reconcile_global_load_blueprints_v1_tasks(self): + @ManagedAppConfig.reconcile_global + def load_blueprints_v1_tasks(self): """Load v1 tasks""" self.import_module("authentik.blueprints.v1.tasks") - def reconcile_tenant_blueprints_discovery(self): + @ManagedAppConfig.reconcile_tenant + def blueprints_discovery(self): """Run blueprint discovery""" from authentik.blueprints.v1.tasks import blueprints_discovery, clear_failed_blueprints diff --git a/authentik/core/apps.py b/authentik/core/apps.py index 866405327c..6fff9cb89f 100644 --- a/authentik/core/apps.py +++ b/authentik/core/apps.py @@ -14,14 +14,16 @@ class AuthentikCoreConfig(ManagedAppConfig): mountpoint = "" default = True - def reconcile_global_debug_worker_hook(self): + @ManagedAppConfig.reconcile_global + def debug_worker_hook(self): """Dispatch startup tasks inline when debugging""" if settings.DEBUG: from authentik.root.celery import worker_ready_hook worker_ready_hook() - def reconcile_tenant_source_inbuilt(self): + @ManagedAppConfig.reconcile_tenant + def source_inbuilt(self): """Reconcile inbuilt source""" from authentik.core.models import Source diff --git a/authentik/crypto/apps.py b/authentik/crypto/apps.py index d80076657c..2336930cca 100644 --- a/authentik/crypto/apps.py +++ b/authentik/crypto/apps.py @@ -36,7 +36,8 @@ class AuthentikCryptoConfig(ManagedAppConfig): }, ) - def reconcile_tenant_managed_jwt_cert(self): + @ManagedAppConfig.reconcile_tenant + def managed_jwt_cert(self): """Ensure managed JWT certificate""" from authentik.crypto.models import CertificateKeyPair @@ -49,7 +50,8 @@ class AuthentikCryptoConfig(ManagedAppConfig): ): self._create_update_cert() - def reconcile_tenant_self_signed(self): + @ManagedAppConfig.reconcile_tenant + def self_signed(self): """Create self-signed keypair""" from authentik.crypto.builder import CertificateBuilder from authentik.crypto.models import CertificateKeyPair diff --git a/authentik/enterprise/audit/apps.py b/authentik/enterprise/audit/apps.py index 6779111bc3..181e266f20 100644 --- a/authentik/enterprise/audit/apps.py +++ b/authentik/enterprise/audit/apps.py @@ -13,7 +13,8 @@ class AuthentikEnterpriseAuditConfig(EnterpriseConfig): verbose_name = "authentik Enterprise.Audit" default = True - def reconcile_global_install_middleware(self): + @EnterpriseConfig.reconcile_global + def install_middleware(self): """Install enterprise audit middleware""" orig_import = "authentik.events.middleware.AuditMiddleware" new_import = "authentik.enterprise.audit.middleware.EnterpriseAuditMiddleware" diff --git a/authentik/events/apps.py b/authentik/events/apps.py index fdad4f9aa3..915b923afc 100644 --- a/authentik/events/apps.py +++ b/authentik/events/apps.py @@ -35,7 +35,8 @@ class AuthentikEventsConfig(ManagedAppConfig): verbose_name = "authentik Events" default = True - def reconcile_global_check_deprecations(self): + @ManagedAppConfig.reconcile_global + def check_deprecations(self): """Check for config deprecations""" from authentik.events.models import Event, EventAction @@ -56,7 +57,8 @@ class AuthentikEventsConfig(ManagedAppConfig): message=msg, ).save() - def reconcile_tenant_prefill_tasks(self): + @ManagedAppConfig.reconcile_tenant + def prefill_tasks(self): """Prefill tasks""" from authentik.events.models import SystemTask from authentik.events.system_tasks import _prefill_tasks @@ -67,7 +69,8 @@ class AuthentikEventsConfig(ManagedAppConfig): task.save() self.logger.debug("prefilled task", task_name=task.name) - def reconcile_tenant_run_scheduled_tasks(self): + @ManagedAppConfig.reconcile_tenant + def run_scheduled_tasks(self): """Run schedule tasks which are behind schedule (only applies to tasks of which we keep metrics)""" from authentik.events.models import TaskStatus diff --git a/authentik/flows/apps.py b/authentik/flows/apps.py index 976dec2e28..da36733f45 100644 --- a/authentik/flows/apps.py +++ b/authentik/flows/apps.py @@ -31,7 +31,8 @@ class AuthentikFlowsConfig(ManagedAppConfig): verbose_name = "authentik Flows" default = True - def reconcile_global_load_stages(self): + @ManagedAppConfig.reconcile_global + def load_stages(self): """Ensure all stages are loaded""" from authentik.flows.models import Stage diff --git a/authentik/outposts/apps.py b/authentik/outposts/apps.py index f07b8297c4..24798c269d 100644 --- a/authentik/outposts/apps.py +++ b/authentik/outposts/apps.py @@ -30,7 +30,8 @@ class AuthentikOutpostConfig(ManagedAppConfig): verbose_name = "authentik Outpost" default = True - def reconcile_tenant_embedded_outpost(self): + @ManagedAppConfig.reconcile_tenant + def embedded_outpost(self): """Ensure embedded outpost""" from authentik.outposts.models import ( DockerServiceConnection, diff --git a/authentik/sources/oauth/apps.py b/authentik/sources/oauth/apps.py index 5a50989f46..f2c110a0bf 100644 --- a/authentik/sources/oauth/apps.py +++ b/authentik/sources/oauth/apps.py @@ -32,7 +32,8 @@ class AuthentikSourceOAuthConfig(ManagedAppConfig): mountpoint = "source/oauth/" default = True - def reconcile_global_sources_loaded(self): + @ManagedAppConfig.reconcile_global + def load_source_types(self): """Load source_types from config file""" for source_type in AUTHENTIK_SOURCES_OAUTH_TYPES: try: diff --git a/authentik/tenants/apps.py b/authentik/tenants/apps.py index 3ae4f3014a..3b4c379de3 100644 --- a/authentik/tenants/apps.py +++ b/authentik/tenants/apps.py @@ -28,7 +28,8 @@ class AuthentikTenantsConfig(ManagedAppConfig): verbose_name = "authentik Tenants" default = True - def reconcile_global_default_tenant(self): + @ManagedAppConfig.reconcile_global + def default_tenant(self): """Make sure default tenant exists, especially after a migration""" post_migrate.connect(ensure_default_tenant) ensure_default_tenant() diff --git a/authentik/tenants/models.py b/authentik/tenants/models.py index 762c41c72e..c7a7876ea6 100644 --- a/authentik/tenants/models.py +++ b/authentik/tenants/models.py @@ -136,7 +136,7 @@ def tenant_needs_sync(sender, tenant, **kwargs): with tenant: for app in apps.get_app_configs(): if isinstance(app, ManagedAppConfig): - app._reconcile(ManagedAppConfig.RECONCILE_TENANT_PREFIX) + app._reconcile(ManagedAppConfig.RECONCILE_TENANT_CATEGORY) tenant.ready = True tenant.save()