refactor: pass db session into service calls (#37403) (#38016)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

api/controllers/console/billing/billing.py

@@ -16,6 +16,7 @@ from controllers.console.wraps import (
     with_current_user,
 )
 from enums.cloud_plan import CloudPlan
+from extensions.ext_database import db
 from libs.login import login_required
 from models import Account
 from services.billing_service import BillingService
@@ -50,7 +51,7 @@ class Subscription(Resource):
     @with_current_tenant_id
     def get(self, current_tenant_id: str, current_user: Account):
         args = SubscriptionQuery.model_validate(request.args.to_dict(flat=True))
-        BillingService.is_tenant_owner_or_admin(current_user)
+        BillingService.is_tenant_owner_or_admin(db.session, current_user)
         return BillingService.get_subscription(args.plan, args.interval, current_user.email, current_tenant_id)
 
 
@@ -64,7 +65,7 @@ class Invoices(Resource):
     @with_current_user
     @with_current_tenant_id
     def get(self, current_tenant_id: str, current_user: Account):
-        BillingService.is_tenant_owner_or_admin(current_user)
+        BillingService.is_tenant_owner_or_admin(db.session, current_user)
         return BillingService.get_invoices(current_user.email, current_tenant_id)
 
 

api/controllers/console/workspace/model_providers.py

@@ -18,6 +18,7 @@ from controllers.console.wraps import (
     with_current_tenant_id,
     with_current_user,
 )
+from extensions.ext_database import db
 from fields.base import ResponseModel
 from graphon.model_runtime.entities.model_entities import ModelType
 from graphon.model_runtime.errors.validate import CredentialsValidateFailedError
@@ -352,7 +353,7 @@ class ModelProviderPaymentCheckoutUrlApi(Resource):
     def get(self, current_tenant_id: str, current_user: Account, provider: str):
         if provider != "anthropic":
             raise ValueError(f"provider name {provider} is invalid")
-        BillingService.is_tenant_owner_or_admin(current_user)
+        BillingService.is_tenant_owner_or_admin(db.session, current_user)
         data = BillingService.get_model_provider_payment_link(
             provider_name=provider,
             tenant_id=current_tenant_id,

api/services/billing_service.py

@@ -7,12 +7,12 @@ from typing import Any, Literal, NotRequired, TypedDict
 import httpx
 from pydantic import TypeAdapter
 from sqlalchemy import select
+from sqlalchemy.orm import Session, scoped_session
 from tenacity import retry, retry_if_exception_type, stop_before_delay, wait_fixed
 from werkzeug.exceptions import InternalServerError
 
 from core.helper.http_client_pooling import get_pooled_http_client
 from enums.cloud_plan import CloudPlan
-from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from libs.helper import RateLimiter
 from models import Account, TenantAccountJoin, TenantAccountRole
@@ -363,10 +363,10 @@ class BillingService:
         return response.json()
 
     @staticmethod
-    def is_tenant_owner_or_admin(current_user: Account):
+    def is_tenant_owner_or_admin(session: Session | scoped_session, current_user: Account):
         tenant_id = current_user.current_tenant_id
 
-        join: TenantAccountJoin | None = db.session.scalar(
+        join: TenantAccountJoin | None = session.scalar(
             select(TenantAccountJoin)
             .where(TenantAccountJoin.tenant_id == tenant_id, TenantAccountJoin.account_id == current_user.id)
             .limit(1)

api/services/clear_free_plan_tenant_expired_logs.py

@@ -125,9 +125,9 @@ class ClearFreePlanTenantExpiredLogs:
             )
 
     @classmethod
-    def process_tenant(cls, flask_app: Flask, tenant_id: str, days: int, batch: int):
+    def process_tenant(cls, flask_app: Flask, tenant_id: str, days: int, batch: int, session: Session):
         with flask_app.app_context():
-            apps = db.session.scalars(select(App).where(App.tenant_id == tenant_id)).all()
+            apps = session.scalars(select(App).where(App.tenant_id == tenant_id)).all()
             app_ids = [app.id for app in apps]
             while True:
                 with sessionmaker(bind=db.engine, autoflush=False).begin() as session:
@@ -375,7 +375,8 @@ class ClearFreePlanTenantExpiredLogs:
                     or BillingService.get_info(tenant_id)["subscription"]["plan"] == CloudPlan.SANDBOX
                 ):
                     # only process sandbox tenant
-                    cls.process_tenant(flask_app, tenant_id, days, batch)
+                    with sessionmaker(db.engine).begin() as session:
+                        cls.process_tenant(flask_app, tenant_id, days, batch, session)
             except Exception:
                 logger.exception("Failed to process tenant %s", tenant_id)
             finally:

api/services/credential_permission_service.py

@@ -1,9 +1,8 @@
 from collections.abc import Sequence
 
 from sqlalchemy import or_, select
-from sqlalchemy.orm import InstrumentedAttribute
+from sqlalchemy.orm import InstrumentedAttribute, Session, scoped_session
 
-from extensions.ext_database import db
 from models.account import Account
 from models.credential_permission import CredentialPermission
 from models.enums import PermissionEnum
@@ -17,9 +16,11 @@ class CredentialPermissionService:
     """
 
     @classmethod
-    def get_partial_member_list(cls, credential_id: str, credential_type: str) -> Sequence[str]:
+    def get_partial_member_list(
+        cls, session: Session | scoped_session, credential_id: str, credential_type: str
+    ) -> Sequence[str]:
         """Return account_ids that have partial-member access to a credential."""
-        return db.session.scalars(
+        return session.scalars(
             select(CredentialPermission.account_id).where(
                 CredentialPermission.credential_id == credential_id,
                 CredentialPermission.credential_type == credential_type,

api/services/tools/builtin_tools_manage_service.py

@@ -427,7 +427,7 @@ class BuiltinToolManageService:
                 if vis_str == "partial_members":
                     credential_entity.partial_member_list = list(
                         CredentialPermissionService.get_partial_member_list(
-                            provider.id, CredPermType.BUILTIN_TOOL_PROVIDER
+                            db.session, provider.id, CredPermType.BUILTIN_TOOL_PROVIDER
                         )
                     )
                 if provider.id in borrowed_ids:

api/tests/test_containers_integration_tests/services/test_billing_service.py

@@ -417,7 +417,7 @@ class TestBillingServiceIsTenantOwnerOrAdmin:
         account, _ = self._create_account_with_tenant_role(db_session_with_containers, TenantAccountRole.EDITOR)
 
         with pytest.raises(ValueError, match="Only team owner or team admin can perform this action"):
-            BillingService.is_tenant_owner_or_admin(account)
+            BillingService.is_tenant_owner_or_admin(db_session_with_containers, account)
 
     def test_is_tenant_owner_or_admin_dataset_operator_raises_error(self, db_session_with_containers: Session) -> None:
         """is_tenant_owner_or_admin raises ValueError for DATASET_OPERATOR role."""
@@ -426,4 +426,4 @@ class TestBillingServiceIsTenantOwnerOrAdmin:
         )
 
         with pytest.raises(ValueError, match="Only team owner or team admin can perform this action"):
-            BillingService.is_tenant_owner_or_admin(account)
+            BillingService.is_tenant_owner_or_admin(db_session_with_containers, account)

api/tests/unit_tests/services/test_billing_service.py

@@ -1031,8 +1031,7 @@ class TestBillingServiceAccountManagement:
     @pytest.fixture
     def mock_db_session(self):
         """Mock database session."""
-        with patch("services.billing_service.db.session") as mock_session:
+        return MagicMock()
-            yield mock_session
 
     def test_delete_account(self, mock_send_request):
         """Test account deletion."""
@@ -1116,7 +1115,8 @@ class TestBillingServiceAccountManagement:
         mock_db_session.scalar.return_value = mock_join
 
         # Act - should not raise exception
-        BillingService.is_tenant_owner_or_admin(current_user)
+        BillingService.is_tenant_owner_or_admin(mock_db_session, current_user)
+        mock_db_session.scalar.assert_called_once()
 
     def test_is_tenant_owner_or_admin_admin(self, mock_db_session):
         """Test tenant owner/admin check for admin role."""
@@ -1131,7 +1131,8 @@ class TestBillingServiceAccountManagement:
         mock_db_session.scalar.return_value = mock_join
 
         # Act - should not raise exception
-        BillingService.is_tenant_owner_or_admin(current_user)
+        BillingService.is_tenant_owner_or_admin(mock_db_session, current_user)
+        mock_db_session.scalar.assert_called_once()
 
     def test_is_tenant_owner_or_admin_normal_user_raises_error(self, mock_db_session):
         """Test tenant owner/admin check raises error for normal user."""
@@ -1147,8 +1148,9 @@ class TestBillingServiceAccountManagement:
 
         # Act & Assert
         with pytest.raises(ValueError) as exc_info:
-            BillingService.is_tenant_owner_or_admin(current_user)
+            BillingService.is_tenant_owner_or_admin(mock_db_session, current_user)
         assert "Only team owner or team admin can perform this action" in str(exc_info.value)
+        mock_db_session.scalar.assert_called_once()
 
     def test_is_tenant_owner_or_admin_no_join_raises_error(self, mock_db_session):
         """Test tenant owner/admin check raises error when join not found."""
@@ -1161,8 +1163,9 @@ class TestBillingServiceAccountManagement:
 
         # Act & Assert
         with pytest.raises(ValueError) as exc_info:
-            BillingService.is_tenant_owner_or_admin(current_user)
+            BillingService.is_tenant_owner_or_admin(mock_db_session, current_user)
         assert "Tenant account join not found" in str(exc_info.value)
+        mock_db_session.scalar.assert_called_once()
 
 
 class TestBillingServiceCacheManagement:

api/tests/unit_tests/services/test_clear_free_plan_tenant_expired_logs.py

@@ -196,9 +196,7 @@ class _ImmediateExecutor:
 
 def _session_wrapper_for_no_autoflush(session: Mock) -> Mock:
     """
-    ClearFreePlanTenantExpiredLogs.process_tenant uses:
+    Return an object with a no_autoflush context manager for legacy tests that need Session-like wrappers.
-      with Session(db.engine).no_autoflush as session:
-    so Session(db.engine) must return an object with a no_autoflush context manager.
     """
     cm = MagicMock()
     cm.__enter__.return_value = session
@@ -224,7 +222,7 @@ def _sessionmaker_wrapper_for_begin(session: Mock) -> Mock:
 
 
 def _session_wrapper_for_direct(session: Mock) -> Mock:
-    """ClearFreePlanTenantExpiredLogs.process uses: with Session(db.engine) as session: (for old code paths)"""
+    """Return an object usable as a direct context manager for legacy Session-like test paths."""
     wrapper = MagicMock()
     wrapper.__enter__.return_value = session
     wrapper.__exit__.return_value = None
@@ -234,17 +232,13 @@ def _session_wrapper_for_direct(session: Mock) -> Mock:
 def test_process_tenant_processes_all_batches(monkeypatch: pytest.MonkeyPatch) -> None:
     flask_app = service_module.Flask("test-app")
 
+    app_session = MagicMock()
+    app_session.scalars.return_value.all.return_value = [SimpleNamespace(id="app-1"), SimpleNamespace(id="app-2")]
+
     monkeypatch.setattr(
         service_module,
         "db",
-        SimpleNamespace(
+        SimpleNamespace(engine=object()),
-            engine=object(),
-            session=SimpleNamespace(
-                scalars=lambda _stmt: SimpleNamespace(
-                    all=lambda: [SimpleNamespace(id="app-1"), SimpleNamespace(id="app-2")]
-                )
-            ),
-        ),
     )
 
     mock_storage = MagicMock()
@@ -326,9 +320,10 @@ def test_process_tenant_processes_all_batches(monkeypatch: pytest.MonkeyPatch) -
         lambda _sm: run_repo,
     )
 
-    ClearFreePlanTenantExpiredLogs.process_tenant(flask_app, "tenant-1", days=7, batch=10)
+    ClearFreePlanTenantExpiredLogs.process_tenant(flask_app, "tenant-1", days=7, batch=10, session=app_session)
 
     # messages backup, conversations backup, node executions backup, runs backup, workflow app logs backup
+    app_session.scalars.assert_called_once()
     assert mock_storage.save.call_count >= 5
     clear_related.assert_called()
 
@@ -389,6 +384,7 @@ def test_process_with_tenant_ids_filters_by_plan_and_logs_errors(
 
     # Only sandbox tenant should attempt processing, and its failure should be swallowed + logged.
     assert process_tenant_mock.call_count == 1
+    assert process_tenant_mock.call_args.args[4] is count_session
     assert "Failed to process tenant t_sandbox" in caplog.messages
     assert "Failed to process tenant t_fail" in caplog.messages
 
@@ -429,7 +425,14 @@ def test_process_without_tenant_ids_batches_and_scales_interval(monkeypatch: pyt
     batch_session.scalar.side_effect = [200, 200, 200, 50]
     batch_session.execute.return_value = rows
 
-    sessions = [_sessionmaker_wrapper_for_begin(total_session), _sessionmaker_wrapper_for_begin(batch_session)]
+    tenant_session_a = MagicMock()
+    tenant_session_b = MagicMock()
+    sessions = [
+        _sessionmaker_wrapper_for_begin(total_session),
+        _sessionmaker_wrapper_for_begin(batch_session),
+        _sessionmaker_wrapper_for_begin(tenant_session_a),
+        _sessionmaker_wrapper_for_begin(tenant_session_b),
+    ]
     monkeypatch.setattr(service_module, "sessionmaker", lambda _engine: sessions.pop(0))
 
     process_tenant_mock = MagicMock()
@@ -500,7 +503,12 @@ def test_process_without_tenant_ids_all_intervals_too_many_uses_min_interval(mon
     batch_session.scalar.side_effect = [200, 200, 200, 200, 200]
     batch_session.execute.return_value = rows
 
-    sessions = [_sessionmaker_wrapper_for_begin(total_session), _sessionmaker_wrapper_for_begin(batch_session)]
+    tenant_session = MagicMock()
+    sessions = [
+        _sessionmaker_wrapper_for_begin(total_session),
+        _sessionmaker_wrapper_for_begin(batch_session),
+        _sessionmaker_wrapper_for_begin(tenant_session),
+    ]
     monkeypatch.setattr(service_module, "sessionmaker", lambda _engine: sessions.pop(0))
 
     process_tenant_mock = MagicMock()
@@ -515,13 +523,13 @@ def test_process_without_tenant_ids_all_intervals_too_many_uses_min_interval(mon
 def test_process_tenant_repo_loops_break_on_empty_second_batch(monkeypatch: pytest.MonkeyPatch) -> None:
     flask_app = service_module.Flask("test-app")
 
+    app_session = MagicMock()
+    app_session.scalars.return_value.all.return_value = [SimpleNamespace(id="app-1")]
+
     monkeypatch.setattr(
         service_module,
         "db",
-        SimpleNamespace(
+        SimpleNamespace(engine=object()),
-            engine=object(),
-            session=SimpleNamespace(scalars=lambda _stmt: SimpleNamespace(all=lambda: [SimpleNamespace(id="app-1")])),
-        ),
     )
     mock_storage = MagicMock()
     monkeypatch.setattr(service_module, "storage", mock_storage)
@@ -585,7 +593,8 @@ def test_process_tenant_repo_loops_break_on_empty_second_batch(monkeypatch: pyte
         lambda _sm: run_repo,
     )
 
-    ClearFreePlanTenantExpiredLogs.process_tenant(flask_app, "tenant-1", days=7, batch=2)
+    ClearFreePlanTenantExpiredLogs.process_tenant(flask_app, "tenant-1", days=7, batch=2, session=app_session)
 
+    app_session.scalars.assert_called_once()
     assert node_repo.get_expired_executions_batch.call_count == 2
     assert run_repo.get_expired_runs_batch.call_count == 2

api/tests/unit_tests/services/test_credential_permission_service.py

@@ -5,7 +5,7 @@ and admin bypass behavior.
 """
 
 from types import SimpleNamespace
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock
 from uuid import uuid4
 
 import pytest
@@ -37,20 +37,22 @@ def credential_id():
 
 class TestGetPartialMemberList:
     def test_returns_empty_when_no_permissions(self, credential_id):
-        with patch("services.credential_permission_service.db") as mock_db:
+        session = MagicMock()
-            mock_db.session.scalars.return_value.all.return_value = []
+        session.scalars.return_value.all.return_value = []
-            result = CredentialPermissionService.get_partial_member_list(
+        result = CredentialPermissionService.get_partial_member_list(
-                credential_id, CredentialType.TRIGGER_SUBSCRIPTION
+            session, credential_id, CredentialType.TRIGGER_SUBSCRIPTION
-            )
+        )
-            assert result == []
+        assert result == []
+        session.scalars.assert_called_once()
 
     def test_returns_account_ids(self, credential_id, user_id, other_user_id):
-        with patch("services.credential_permission_service.db") as mock_db:
+        session = MagicMock()
-            mock_db.session.scalars.return_value.all.return_value = [user_id, other_user_id]
+        session.scalars.return_value.all.return_value = [user_id, other_user_id]
-            result = CredentialPermissionService.get_partial_member_list(
+        result = CredentialPermissionService.get_partial_member_list(
-                credential_id, CredentialType.TRIGGER_SUBSCRIPTION
+            session, credential_id, CredentialType.TRIGGER_SUBSCRIPTION
-            )
+        )
-            assert set(result) == {user_id, other_user_id}
+        assert set(result) == {user_id, other_user_id}
+        session.scalars.assert_called_once()
 
 
 class TestApplyVisibilityFilter: