opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-07-21 02:28:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
Mr_xie 2026-06-26 15:53:27 +08:00 committed by GitHub
parent eb4ec93cea
commit 3f2ef24755
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 73 additions and 55 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
api/controllers/console/billing/billing.py
View File

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

3
api/controllers/console/workspace/model_providers.py
View File

@ -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,

6
api/services/billing_service.py
View File

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

7
api/services/clear_free_plan_tenant_expired_logs.py
View File

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

9
api/services/credential_permission_service.py
View File

@ -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,

2
api/services/tools/builtin_tools_manage_service.py
View File

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

4
api/tests/test_containers_integration_tests/services/test_billing_service.py
View File

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

15
api/tests/unit_tests/services/test_billing_service.py
View File

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

49
api/tests/unit_tests/services/test_clear_free_plan_tenant_expired_logs.py
View File

@ -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

28
api/tests/unit_tests/services/test_credential_permission_service.py
View File

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