refactor(web): use shared session factory for JWT decoding (#41007)

api/controllers/web/wraps.py

@@ -6,11 +6,11 @@ from typing import Concatenate
 from flask import request
 from flask_restx import Resource
 from sqlalchemy import select
-from sqlalchemy.orm import sessionmaker
 from werkzeug.exceptions import BadRequest, NotFound, Unauthorized
 
 from constants import HEADER_NAME_APP_CODE
 from controllers.web.error import WebAppAuthAccessDeniedError, WebAppAuthRequiredError
+from core.db.session_factory import session_factory
 from core.logging.context import set_identity_context
 from extensions.ext_database import db
 from libs.passport import PassportService
@@ -54,7 +54,7 @@ def decode_jwt_token(app_code: str | None = None, user_id: str | None = None) ->
         decoded = PassportService().verify(tk)
         app_code = decoded.get("app_code")
         app_id = decoded.get("app_id")
-        with sessionmaker(db.engine, expire_on_commit=False).begin() as session:
+        with session_factory.create_session() as session:
             app_model = session.scalar(select(App).where(App.id == app_id))
             site = session.scalar(select(Site).where(Site.code == app_code))
             if not app_model:

api/tests/unit_tests/controllers/web/test_wraps.py

@@ -1,6 +1,10 @@
+from types import SimpleNamespace
 from unittest import mock
+from uuid import uuid4
 
 import pytest
+from flask import Flask
+from sqlalchemy.orm import Session
 from werkzeug.exceptions import Unauthorized
 
 from core.logging.context import clear_request_context, get_identity_context
@@ -48,3 +52,57 @@ def test_validate_jwt_token_does_not_set_identity_when_authentication_fails() ->
         protected_view()
 
     assert get_identity_context() == ("", "", "")
+
+
+def test_decode_jwt_token_uses_shared_session_factory(sqlite_session: Session) -> None:
+    from controllers.web import wraps
+    from models.enums import EndUserType
+    from models.model import AppMode, CustomizeTokenStrategy, Site
+
+    tenant_id = str(uuid4())
+    app_model = App(
+        tenant_id=tenant_id,
+        mode=AppMode.CHAT.value,
+        name="test-app",
+        enable_site=True,
+        enable_api=True,
+    )
+    sqlite_session.add(app_model)
+    sqlite_session.commit()
+
+    site = Site(
+        app_id=app_model.id,
+        title="test-site",
+        default_language="en-US",
+        customize_token_strategy=CustomizeTokenStrategy.NOT_ALLOW,
+        code="app-code",
+    )
+    end_user = EndUser(
+        tenant_id=tenant_id,
+        app_id=app_model.id,
+        type=EndUserType.BROWSER,
+        session_id="session-id",
+    )
+    sqlite_session.add_all((site, end_user))
+    sqlite_session.commit()
+
+    with (
+        mock.patch.object(wraps, "extract_webapp_passport", return_value="jwt-token"),
+        mock.patch.object(wraps, "PassportService") as mock_passport_service,
+        mock.patch.object(
+            wraps,
+            "FeatureService",
+            get_system_features=mock.Mock(return_value=SimpleNamespace(webapp_auth=SimpleNamespace(enabled=False))),
+        ),
+    ):
+        mock_passport_service.return_value.verify.return_value = {
+            "app_code": "app-code",
+            "app_id": app_model.id,
+            "end_user_id": end_user.id,
+        }
+
+        with Flask(__name__).test_request_context("/", headers={"X-App-Code": "app-code"}):
+            result_app, result_end_user = wraps.decode_jwt_token()
+
+    assert result_app.id == app_model.id
+    assert result_end_user.id == end_user.id