chore(api): cache the setup status to cut down DB access (#36966)

Co-authored-by: Asuka Minato <i@asukaminato.eu.org>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Byron.wang <byron@dify.ai>

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")
 

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()

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()

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):