diff --git a/api/controllers/console/billing/billing.py b/api/controllers/console/billing/billing.py index 1c2d129e6c3..cd719966e6a 100644 --- a/api/controllers/console/billing/billing.py +++ b/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) diff --git a/api/controllers/console/workspace/model_providers.py b/api/controllers/console/workspace/model_providers.py index 3ce7211703e..779399f9055 100644 --- a/api/controllers/console/workspace/model_providers.py +++ b/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, diff --git a/api/services/billing_service.py b/api/services/billing_service.py index 5c59a5f8dd1..ec391e51676 100644 --- a/api/services/billing_service.py +++ b/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) diff --git a/api/services/clear_free_plan_tenant_expired_logs.py b/api/services/clear_free_plan_tenant_expired_logs.py index 3fb340d3a75..963938982d0 100644 --- a/api/services/clear_free_plan_tenant_expired_logs.py +++ b/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: diff --git a/api/services/credential_permission_service.py b/api/services/credential_permission_service.py index 21806e0178a..2b1082d132b 100644 --- a/api/services/credential_permission_service.py +++ b/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, diff --git a/api/services/tools/builtin_tools_manage_service.py b/api/services/tools/builtin_tools_manage_service.py index 22086317c02..e49ab8398f1 100644 --- a/api/services/tools/builtin_tools_manage_service.py +++ b/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: diff --git a/api/tests/test_containers_integration_tests/services/test_billing_service.py b/api/tests/test_containers_integration_tests/services/test_billing_service.py index 4893126d7f2..a3a4a0e6edd 100644 --- a/api/tests/test_containers_integration_tests/services/test_billing_service.py +++ b/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) diff --git a/api/tests/unit_tests/services/test_billing_service.py b/api/tests/unit_tests/services/test_billing_service.py index f244b69407f..e5610545aa9 100644 --- a/api/tests/unit_tests/services/test_billing_service.py +++ b/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: - yield mock_session + return MagicMock() 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: diff --git a/api/tests/unit_tests/services/test_clear_free_plan_tenant_expired_logs.py b/api/tests/unit_tests/services/test_clear_free_plan_tenant_expired_logs.py index 5d0a93e239d..9be5af2b046 100644 --- a/api/tests/unit_tests/services/test_clear_free_plan_tenant_expired_logs.py +++ b/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: - with Session(db.engine).no_autoflush as session: - so Session(db.engine) must return an object with a no_autoflush context manager. + Return an object with a no_autoflush context manager for legacy tests that need Session-like wrappers. """ 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( - engine=object(), - session=SimpleNamespace( - scalars=lambda _stmt: SimpleNamespace( - all=lambda: [SimpleNamespace(id="app-1"), SimpleNamespace(id="app-2")] - ) - ), - ), + SimpleNamespace(engine=object()), ) 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( - engine=object(), - session=SimpleNamespace(scalars=lambda _stmt: SimpleNamespace(all=lambda: [SimpleNamespace(id="app-1")])), - ), + SimpleNamespace(engine=object()), ) 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 diff --git a/api/tests/unit_tests/services/test_credential_permission_service.py b/api/tests/unit_tests/services/test_credential_permission_service.py index 687d6880178..e467e9c8c5e 100644 --- a/api/tests/unit_tests/services/test_credential_permission_service.py +++ b/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: - mock_db.session.scalars.return_value.all.return_value = [] - result = CredentialPermissionService.get_partial_member_list( - credential_id, CredentialType.TRIGGER_SUBSCRIPTION - ) - assert result == [] + session = MagicMock() + session.scalars.return_value.all.return_value = [] + result = CredentialPermissionService.get_partial_member_list( + session, credential_id, CredentialType.TRIGGER_SUBSCRIPTION + ) + 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: - mock_db.session.scalars.return_value.all.return_value = [user_id, other_user_id] - result = CredentialPermissionService.get_partial_member_list( - credential_id, CredentialType.TRIGGER_SUBSCRIPTION - ) - assert set(result) == {user_id, other_user_id} + session = MagicMock() + session.scalars.return_value.all.return_value = [user_id, other_user_id] + result = CredentialPermissionService.get_partial_member_list( + session, credential_id, CredentialType.TRIGGER_SUBSCRIPTION + ) + assert set(result) == {user_id, other_user_id} + session.scalars.assert_called_once() class TestApplyVisibilityFilter: