From 64aa1426815005c5fcdc3180d7b86439b450979a Mon Sep 17 00:00:00 2001 From: chariri Date: Tue, 7 Jul 2026 23:30:28 +0900 Subject: [PATCH] chore(api): cache the setup status to cut down DB access (#36966) Co-authored-by: Asuka Minato Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Byron.wang --- api/controllers/console/setup.py | 3 +- api/controllers/console/wraps.py | 62 ++++++++++++++++++- .../console/test_fastopenapi_setup.py | 2 + .../controllers/console/test_wraps.py | 40 ++++++++++++ 4 files changed, 103 insertions(+), 4 deletions(-) diff --git a/api/controllers/console/setup.py b/api/controllers/console/setup.py index 3b5c1bbe18f..2b99693a9ca 100644 --- a/api/controllers/console/setup.py +++ b/api/controllers/console/setup.py @@ -13,7 +13,7 @@ from services.account_service import RegisterService, TenantService from .error import AlreadySetupError, NotInitValidateError from .init_validate import get_init_validate_status -from .wraps import only_edition_self_hosted +from .wraps import mark_setup_completed, only_edition_self_hosted class SetupRequestPayload(BaseModel): @@ -96,6 +96,7 @@ def setup_system(payload: SetupRequestPayload) -> SetupResponse: language=payload.language, session=db.session, ) + mark_setup_completed() return SetupResponse(result="success") diff --git a/api/controllers/console/wraps.py b/api/controllers/console/wraps.py index 017793ffe0b..37d7239170c 100644 --- a/api/controllers/console/wraps.py +++ b/api/controllers/console/wraps.py @@ -4,7 +4,7 @@ import os import time from collections.abc import Callable from functools import wraps -from typing import Any, Concatenate, overload +from typing import Any, Concatenate, Protocol, cast, overload from flask import abort, request from pydantic import BaseModel, ValidationError @@ -46,6 +46,60 @@ ERROR_MSG_INVALID_ENCRYPTED_DATA = "Invalid encrypted data" ERROR_MSG_INVALID_ENCRYPTED_CODE = "Invalid encrypted code" +class OnceTrueCallable[**P](Protocol): + def __call__(self, *args: P.args, **kwargs: P.kwargs) -> bool: ... + + def mark_success(self) -> None: ... + + def reset_success(self) -> None: ... + + +def once_true[**P](func: Callable[P, bool]) -> OnceTrueCallable[P]: + """Wrap a predicate so only a strict True result is memoized.""" + has_success = False + + def mark_success() -> None: + nonlocal has_success + + has_success = True + + def reset_success() -> None: + nonlocal has_success + + has_success = False + + @wraps(func) + def wrapper(*args: P.args, **kwargs: P.kwargs) -> bool: + nonlocal has_success + + if has_success: + return True + + result = func(*args, **kwargs) + if result is True: + has_success = True + + return result + + wrapper.mark_success = mark_success # type: ignore[attr-defined] + wrapper.reset_success = reset_success # type: ignore[attr-defined] + return cast(OnceTrueCallable[P], wrapper) + + +def mark_setup_completed() -> None: + """Remember in this process that one-time self-hosted setup has completed.""" + _is_setup_completed.mark_success() + + +@once_true +def _is_setup_completed() -> bool: + """Check whether setup exists, caching only successful observations. + + Use `once_true` instead of `@cache` because a pre-setup False result must not be memoized. + """ + return db.session.scalar(select(DifySetup).limit(1)) is not None + + @overload def account_initialization_required[T, **P, R]( view: Callable[Concatenate[T, P], R], @@ -246,7 +300,9 @@ def setup_required[T, **P, R]( @overload -def setup_required[**P, R](view: Callable[P, R]) -> Callable[P, R]: ... +def setup_required[**P, R](view: Callable[P, R]) -> Callable[P, R]: + """Require self-hosted bootstrap setup before serving protected routes.""" + ... def setup_required[R](view: Callable[..., R]) -> Callable[..., R]: @@ -255,7 +311,7 @@ def setup_required[R](view: Callable[..., R]) -> Callable[..., R]: # The overloads keep Resource methods method-aware for pyrefly while # preserving support for plain functions used in tests and utilities. # check setup - if dify_config.EDITION == "SELF_HOSTED" and not db.session.scalar(select(DifySetup).limit(1)): + if dify_config.EDITION == "SELF_HOSTED" and not _is_setup_completed(): if os.environ.get("INIT_PASSWORD"): raise NotInitValidateError() raise NotSetupError() diff --git a/api/tests/unit_tests/controllers/console/test_fastopenapi_setup.py b/api/tests/unit_tests/controllers/console/test_fastopenapi_setup.py index 385539b6f30..2b385304d32 100644 --- a/api/tests/unit_tests/controllers/console/test_fastopenapi_setup.py +++ b/api/tests/unit_tests/controllers/console/test_fastopenapi_setup.py @@ -48,9 +48,11 @@ def test_console_setup_fastopenapi_post_success(app: Flask): patch("controllers.console.setup.TenantService.get_tenant_count", return_value=0), patch("controllers.console.setup.get_init_validate_status", return_value=True), patch("controllers.console.setup.RegisterService.setup"), + patch("controllers.console.setup.mark_setup_completed") as mark_setup_completed, ): client = app.test_client() response = client.post("/console/api/setup", json=payload) assert response.status_code == 201 assert response.get_json() == {"result": "success"} + mark_setup_completed.assert_called_once_with() diff --git a/api/tests/unit_tests/controllers/console/test_wraps.py b/api/tests/unit_tests/controllers/console/test_wraps.py index 618a1f52180..172125f4635 100644 --- a/api/tests/unit_tests/controllers/console/test_wraps.py +++ b/api/tests/unit_tests/controllers/console/test_wraps.py @@ -13,6 +13,7 @@ from controllers.console.workspace.error import AccountNotInitializedError from controllers.console.wraps import ( RBACPermission, RBACResourceScope, + _is_setup_completed, account_initialization_required, cloud_edition_billing_enabled, cloud_edition_billing_rate_limit_check, @@ -35,6 +36,12 @@ from models.account import AccountStatus, TenantAccountRole from services.feature_service import LicenseStatus +@pytest.fixture(autouse=True) +def reset_setup_required_cache(): + """Keep setup_required's process cache isolated across unit tests.""" + _is_setup_completed.reset_success() + + class MockUser(UserMixin): """Simple User class for testing.""" @@ -735,6 +742,39 @@ class TestSystemSetup: # Assert assert result == "admin_success" + @patch("controllers.console.wraps.db") + def test_should_cache_completed_setup(self, mock_db): + """Test that completed setup skips repeated DB reads in this process""" + mock_db.session.scalar.return_value = MagicMock() + + @setup_required + def admin_view(): + return "admin_success" + + with patch("controllers.console.wraps.dify_config.EDITION", "SELF_HOSTED"): + assert admin_view() == "admin_success" + assert admin_view() == "admin_success" + + assert mock_db.session.scalar.call_count == 1 + + @patch("controllers.console.wraps.db") + @patch("controllers.console.wraps.os.environ.get") + def test_should_not_cache_missing_setup(self, mock_environ_get, mock_db): + """Test that first-time bootstrap completion can be observed later in the same process""" + mock_db.session.scalar.side_effect = [None, MagicMock()] + mock_environ_get.return_value = None + + @setup_required + def admin_view(): + return "admin_success" + + with patch("controllers.console.wraps.dify_config.EDITION", "SELF_HOSTED"): + with pytest.raises(NotSetupError): + admin_view() + assert admin_view() == "admin_success" + + assert mock_db.session.scalar.call_count == 2 + @patch("controllers.console.wraps.db") @patch("controllers.console.wraps.os.environ.get") def test_should_raise_not_init_validate_error_with_init_password(self, mock_environ_get, mock_db: MagicMock):