From c755232f0a30a009a68bc8f6a316413e4781519a Mon Sep 17 00:00:00 2001 From: Marc 'risson' Schmitt Date: Wed, 17 Jun 2026 15:17:04 +0200 Subject: [PATCH] brands: select_related models accessed in the hot path (#23162) Co-authored-by: Ryan Pesek Co-authored-by: Jens Langhammer Signed-off-by: Marc 'risson' Schmitt Signed-off-by: Jens Langhammer --- authentik/brands/models.py | 14 ++++ authentik/brands/tests/__init__.py | 0 .../brands/{tests.py => tests/test_brands.py} | 0 .../brands/tests/test_current_queries.py | 79 +++++++++++++++++++ authentik/brands/utils.py | 5 +- 5 files changed, 96 insertions(+), 2 deletions(-) create mode 100644 authentik/brands/tests/__init__.py rename authentik/brands/{tests.py => tests/test_brands.py} (100%) create mode 100644 authentik/brands/tests/test_current_queries.py diff --git a/authentik/brands/models.py b/authentik/brands/models.py index 94c772767a..ab49ee96f5 100644 --- a/authentik/brands/models.py +++ b/authentik/brands/models.py @@ -18,6 +18,20 @@ from authentik.lib.models import SerializerModel LOGGER = get_logger() +# Brand FKs read on the request hot path. select_related pulls them into the +# same SELECT to avoid N+1 lazy loads; CurrentBrandSerializer alone reads 7. +_BRAND_RELATED_FK_FIELDS = ( + "flow_authentication", + "flow_invalidation", + "flow_recovery", + "flow_unenrollment", + "flow_user_settings", + "flow_device_code", + "flow_lockdown", + "default_application", +) + + class Brand(SerializerModel): """Single brand""" diff --git a/authentik/brands/tests/__init__.py b/authentik/brands/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/authentik/brands/tests.py b/authentik/brands/tests/test_brands.py similarity index 100% rename from authentik/brands/tests.py rename to authentik/brands/tests/test_brands.py diff --git a/authentik/brands/tests/test_current_queries.py b/authentik/brands/tests/test_current_queries.py new file mode 100644 index 0000000000..7de78cb062 --- /dev/null +++ b/authentik/brands/tests/test_current_queries.py @@ -0,0 +1,79 @@ +"""Test brands""" + +from django.http import HttpRequest +from django.test import TestCase + +from authentik.brands.models import Brand +from authentik.brands.utils import _BRAND_RELATED_FK_FIELDS, get_brand_for_request +from authentik.core.tests.utils import create_test_flow +from authentik.flows.models import FlowDesignation + + +class TestGetBrandForRequestSelectRelated(TestCase): + """``get_brand_for_request`` must hydrate the FK fields read on the + request hot path so later access doesn't trigger lazy loads.""" + + def setUp(self): + Brand.objects.all().delete() + self.flow_auth = create_test_flow(designation=FlowDesignation.AUTHENTICATION) + self.brand = Brand.objects.create( + domain="select-related-test.example.com", + flow_authentication=self.flow_auth, + ) + + def _make_request(self, host: str) -> HttpRequest: + request = HttpRequest() + request.META["HTTP_HOST"] = host + return request + + def test_brand_fks_are_loaded_in_single_query(self): + """Brand FK access after ``get_brand_for_request`` must not trigger + extra queries.""" + request = self._make_request("select-related-test.example.com") + with self.assertNumQueries(1): + brand = get_brand_for_request(request) + _ = brand.flow_authentication + _ = brand.flow_authentication.slug if brand.flow_authentication else None + _ = brand.flow_invalidation + _ = brand.flow_recovery + _ = brand.flow_unenrollment + _ = brand.flow_user_settings + _ = brand.flow_device_code + _ = brand.flow_lockdown + _ = brand.default_application + + def test_brand_related_fk_list_complete(self): + """``_BRAND_RELATED_FK_FIELDS`` covers every Flow/Application FK on + Brand — fails loud when a new FK is added but not registered here.""" + actual_fks = { + f.name + for f in Brand._meta.get_fields() + if f.many_to_one and f.related_model is not None + } + relevant_fks = { + name for name in actual_fks if name.startswith("flow_") or name == "default_application" + } + declared = set(_BRAND_RELATED_FK_FIELDS) + missing = relevant_fks - declared + self.assertFalse( + missing, + f"Brand FK fields {missing} aren't in _BRAND_RELATED_FK_FIELDS — " + "add them or the request hot path will incur extra queries.", + ) + + def test_brand_related_fks_all_exist_on_model(self): + """Every entry in ``_BRAND_RELATED_FK_FIELDS`` is a real FK on Brand. + ``select_related`` raises ``FieldError`` at first use if any entry + is stale, which would break every request.""" + actual_fks = { + f.name + for f in Brand._meta.get_fields() + if f.many_to_one and f.related_model is not None + } + declared = set(_BRAND_RELATED_FK_FIELDS) + extraneous = declared - actual_fks + self.assertFalse( + extraneous, + f"_BRAND_RELATED_FK_FIELDS contains {extraneous} which don't " + f"exist on Brand (actual FKs: {sorted(actual_fks)}).", + ) diff --git a/authentik/brands/utils.py b/authentik/brands/utils.py index 9b2b5be65c..b063a8ede8 100644 --- a/authentik/brands/utils.py +++ b/authentik/brands/utils.py @@ -9,7 +9,7 @@ from django.utils.html import _json_script_escapes from django.utils.safestring import mark_safe from authentik import authentik_full_version -from authentik.brands.models import Brand +from authentik.brands.models import _BRAND_RELATED_FK_FIELDS, Brand from authentik.lib.sentry import get_http_meta from authentik.tenants.models import Tenant @@ -21,7 +21,8 @@ def get_brand_for_request(request: HttpRequest) -> Brand: """Get brand object for current request""" brand = ( - Brand.objects.annotate( + Brand.objects.select_related(*_BRAND_RELATED_FK_FIELDS) + .annotate( host_domain=Value(request.get_host()), domain_length=Length("domain"), match_priority=Case(