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

refactor(api): use sessionmaker in trigger provider service & dataset… (#34774)

Browse Source
This commit is contained in:
carlos4s 2026-04-08 18:13:38 -05:00 committed by GitHub
parent 3a4756449a
commit 4c70bfa8b8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 22 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
api/core/rag/retrieval/dataset_retrieval.py
View File

@ -15,7 +15,7 @@ from graphon.model_runtime.entities.message_entities import PromptMessage, Promp
from graphon.model_runtime.entities.model_entities import ModelFeature, ModelType from graphon.model_runtime.entities.model_entities import ModelFeature, ModelType
from graphon.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel from graphon.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
from sqlalchemy import and_, func, literal, or_, select from sqlalchemy import and_, func, literal, or_, select
from sqlalchemy.orm import Session from sqlalchemy.orm import sessionmaker
from core.app.app_config.entities import ( from core.app.app_config.entities import (
DatasetEntity, DatasetEntity,
@ -884,7 +884,7 @@ class DatasetRetrieval:
self._send_trace_task(message_id, documents, timer) self._send_trace_task(message_id, documents, timer)
return return
with Session(db.engine) as session: with sessionmaker(bind=db.engine).begin() as session:
# Collect all document_ids and batch fetch DatasetDocuments # Collect all document_ids and batch fetch DatasetDocuments
document_ids = { document_ids = {
doc.metadata["document_id"] doc.metadata["document_id"]
@ -975,7 +975,6 @@ class DatasetRetrieval:
{DocumentSegment.hit_count: DocumentSegment.hit_count + 1}, {DocumentSegment.hit_count: DocumentSegment.hit_count + 1},
synchronize_session=False, synchronize_session=False,
) )
session.commit()
self._send_trace_task(message_id, documents, timer) self._send_trace_task(message_id, documents, timer)

22
api/services/trigger/trigger_provider_service.py
View File

@ -6,7 +6,7 @@ from collections.abc import Mapping
from typing import Any from typing import Any
from sqlalchemy import desc, func from sqlalchemy import desc, func
from sqlalchemy.orm import Session from sqlalchemy.orm import Session, sessionmaker
from configs import dify_config from configs import dify_config
from constants import HIDDEN_VALUE, UNKNOWN_VALUE from constants import HIDDEN_VALUE, UNKNOWN_VALUE
@ -146,7 +146,7 @@ class TriggerProviderService:
""" """
try: try:
provider_controller = TriggerManager.get_trigger_provider(tenant_id, provider_id) provider_controller = TriggerManager.get_trigger_provider(tenant_id, provider_id)
with Session(db.engine, expire_on_commit=False) as session: with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session:
# Use distributed lock to prevent race conditions # Use distributed lock to prevent race conditions
lock_key = f"trigger_provider_create_lock:{tenant_id}_{provider_id}" lock_key = f"trigger_provider_create_lock:{tenant_id}_{provider_id}"
with redis_client.lock(lock_key, timeout=20): with redis_client.lock(lock_key, timeout=20):
@ -205,7 +205,6 @@ class TriggerProviderService:
subscription.id = subscription_id or str(uuid.uuid4()) subscription.id = subscription_id or str(uuid.uuid4())
session.add(subscription) session.add(subscription)
session.commit()
return { return {
"result": "success", "result": "success",
@ -241,7 +240,7 @@ class TriggerProviderService:
:param expires_at: Optional new expiration timestamp :param expires_at: Optional new expiration timestamp
:return: Success response with updated subscription info :return: Success response with updated subscription info
""" """
with Session(db.engine, expire_on_commit=False) as session: with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session:
# Use distributed lock to prevent race conditions on the same subscription # Use distributed lock to prevent race conditions on the same subscription
lock_key = f"trigger_subscription_update_lock:{tenant_id}_{subscription_id}" lock_key = f"trigger_subscription_update_lock:{tenant_id}_{subscription_id}"
with redis_client.lock(lock_key, timeout=20): with redis_client.lock(lock_key, timeout=20):
@ -302,8 +301,6 @@ class TriggerProviderService:
if expires_at is not None: if expires_at is not None:
subscription.expires_at = expires_at subscription.expires_at = expires_at
session.commit()
# Clear subscription cache # Clear subscription cache
delete_cache_for_subscription( delete_cache_for_subscription(
tenant_id=tenant_id, tenant_id=tenant_id,
@ -404,7 +401,7 @@ class TriggerProviderService:
:param subscription_id: Subscription instance ID :param subscription_id: Subscription instance ID
:return: New token info :return: New token info
""" """
with Session(db.engine) as session: with sessionmaker(bind=db.engine).begin() as session:
subscription = session.query(TriggerSubscription).filter_by(tenant_id=tenant_id, id=subscription_id).first() subscription = session.query(TriggerSubscription).filter_by(tenant_id=tenant_id, id=subscription_id).first()
if not subscription: if not subscription:
@ -448,7 +445,6 @@ class TriggerProviderService:
# Update credentials # Update credentials
subscription.credentials = dict(encrypter.encrypt(dict(refreshed_credentials.credentials))) subscription.credentials = dict(encrypter.encrypt(dict(refreshed_credentials.credentials)))
subscription.credential_expires_at = refreshed_credentials.expires_at subscription.credential_expires_at = refreshed_credentials.expires_at
session.commit()
# Clear cache # Clear cache
cache.delete() cache.delete()
@ -478,7 +474,7 @@ class TriggerProviderService:
""" """
now_ts: int = int(now if now is not None else _time.time()) now_ts: int = int(now if now is not None else _time.time())
with Session(db.engine) as session: with sessionmaker(bind=db.engine).begin() as session:
subscription: TriggerSubscription | None = ( subscription: TriggerSubscription | None = (
session.query(TriggerSubscription).filter_by(tenant_id=tenant_id, id=subscription_id).first() session.query(TriggerSubscription).filter_by(tenant_id=tenant_id, id=subscription_id).first()
) )
@ -531,7 +527,6 @@ class TriggerProviderService:
# Persist refreshed properties and expires_at # Persist refreshed properties and expires_at
subscription.properties = dict(properties_encrypter.encrypt(dict(refreshed.properties))) subscription.properties = dict(properties_encrypter.encrypt(dict(refreshed.properties)))
subscription.expires_at = int(refreshed.expires_at) subscription.expires_at = int(refreshed.expires_at)
session.commit()
properties_cache.delete() properties_cache.delete()
logger.info( logger.info(
@ -639,7 +634,7 @@ class TriggerProviderService:
tenant_id=tenant_id, provider_id=provider_id tenant_id=tenant_id, provider_id=provider_id
) )
with Session(db.engine) as session: with sessionmaker(bind=db.engine).begin() as session:
# Find existing custom client params # Find existing custom client params
custom_client = ( custom_client = (
session.query(TriggerOAuthTenantClient) session.query(TriggerOAuthTenantClient)
@ -683,8 +678,6 @@ class TriggerProviderService:
if enabled is not None: if enabled is not None:
custom_client.enabled = enabled custom_client.enabled = enabled
session.commit()
return {"result": "success"} return {"result": "success"}
@classmethod @classmethod
@ -733,13 +726,12 @@ class TriggerProviderService:
:param provider_id: Provider identifier :param provider_id: Provider identifier
:return: Success response :return: Success response
""" """
with Session(db.engine) as session: with sessionmaker(bind=db.engine).begin() as session:
session.query(TriggerOAuthTenantClient).filter_by( session.query(TriggerOAuthTenantClient).filter_by(
tenant_id=tenant_id, tenant_id=tenant_id,
provider=provider_id.provider_name, provider=provider_id.provider_name,
plugin_id=provider_id.plugin_id, plugin_id=provider_id.plugin_id,
).delete() ).delete()
session.commit()
return {"result": "success"} return {"result": "success"}

6
api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval.py
View File

@ -4909,15 +4909,17 @@ class TestInternalHooksCoverage:
session_ctx.__enter__.return_value = session session_ctx.__enter__.return_value = session
session_ctx.__exit__.return_value = False session_ctx.__exit__.return_value = False
sessionmaker_ctx = MagicMock()
sessionmaker_ctx.begin.return_value = session_ctx
with ( with (
patch("core.rag.retrieval.dataset_retrieval.db", SimpleNamespace(engine=Mock())), patch("core.rag.retrieval.dataset_retrieval.db", SimpleNamespace(engine=Mock())),
patch("core.rag.retrieval.dataset_retrieval.Session", return_value=session_ctx), patch("core.rag.retrieval.dataset_retrieval.sessionmaker", return_value=sessionmaker_ctx),
patch.object(retrieval, "_send_trace_task") as mock_trace, patch.object(retrieval, "_send_trace_task") as mock_trace,
): ):
retrieval._on_retrieval_end(flask_app=app, documents=docs, message_id="m1", timer={"cost": 1}) retrieval._on_retrieval_end(flask_app=app, documents=docs, message_id="m1", timer={"cost": 1})
query.update.assert_called_once() query.update.assert_called_once()
session.commit.assert_called_once()
mock_trace.assert_called_once() mock_trace.assert_called_once()
def test_retriever_variants(self, retrieval: DatasetRetrieval) -> None: def test_retriever_variants(self, retrieval: DatasetRetrieval) -> None:

16
api/tests/unit_tests/services/test_trigger_provider_service.py
View File

@ -63,6 +63,12 @@ def mock_session(mocker: MockerFixture) -> MagicMock:
mock_session_cm.__enter__.return_value = mock_session_instance mock_session_cm.__enter__.return_value = mock_session_instance
mock_session_cm.__exit__.return_value = False mock_session_cm.__exit__.return_value = False
mocker.patch("services.trigger.trigger_provider_service.Session", return_value=mock_session_cm) mocker.patch("services.trigger.trigger_provider_service.Session", return_value=mock_session_cm)
mock_begin_cm = MagicMock()
mock_begin_cm.__enter__.return_value = mock_session_instance
mock_begin_cm.__exit__.return_value = False
mock_sessionmaker_instance = MagicMock()
mock_sessionmaker_instance.begin.return_value = mock_begin_cm
mocker.patch("services.trigger.trigger_provider_service.sessionmaker", return_value=mock_sessionmaker_instance)
return mock_session_instance return mock_session_instance
@ -212,7 +218,6 @@ def test_add_trigger_subscription_should_create_subscription_successfully_for_ap
# Assert # Assert
assert result["result"] == "success" assert result["result"] == "success"
mock_session.add.assert_called_once() mock_session.add.assert_called_once()
mock_session.commit.assert_called_once()
def test_add_trigger_subscription_should_store_empty_credentials_for_unauthorized_type( def test_add_trigger_subscription_should_store_empty_credentials_for_unauthorized_type(
@ -406,7 +411,7 @@ def test_update_trigger_subscription_should_update_fields_and_clear_cache(
assert subscription.credentials == {"api_key": "new-key"} assert subscription.credentials == {"api_key": "new-key"}
assert subscription.credential_expires_at == 100 assert subscription.credential_expires_at == 100
assert subscription.expires_at == 200 assert subscription.expires_at == 200
mock_session.commit.assert_called_once()
mock_delete_cache.assert_called_once() mock_delete_cache.assert_called_once()
@ -593,7 +598,7 @@ def test_refresh_oauth_token_should_refresh_and_persist_new_credentials(
assert result == {"result": "success", "expires_at": 12345} assert result == {"result": "success", "expires_at": 12345}
assert subscription.credentials == {"access_token": "new"} assert subscription.credentials == {"access_token": "new"}
assert subscription.credential_expires_at == 12345 assert subscription.credential_expires_at == 12345
mock_session.commit.assert_called_once()
cache.delete.assert_called_once() cache.delete.assert_called_once()
@ -664,7 +669,7 @@ def test_refresh_subscription_should_refresh_and_persist_properties(
assert result == {"result": "success", "expires_at": 999} assert result == {"result": "success", "expires_at": 999}
assert subscription.properties == {"p": "new-enc"} assert subscription.properties == {"p": "new-enc"}
assert subscription.expires_at == 999 assert subscription.expires_at == 999
mock_session.commit.assert_called_once()
prop_cache.delete.assert_called_once() prop_cache.delete.assert_called_once()
@ -838,7 +843,6 @@ def test_save_custom_oauth_client_params_should_create_record_and_clear_params_w
assert fake_model.encrypted_oauth_params == "{}" assert fake_model.encrypted_oauth_params == "{}"
assert fake_model.enabled is True assert fake_model.enabled is True
mock_session.add.assert_called_once_with(fake_model) mock_session.add.assert_called_once_with(fake_model)
mock_session.commit.assert_called_once()
def test_save_custom_oauth_client_params_should_merge_hidden_values_and_delete_cache( def test_save_custom_oauth_client_params_should_merge_hidden_values_and_delete_cache(
@ -870,7 +874,6 @@ def test_save_custom_oauth_client_params_should_merge_hidden_values_and_delete_c
assert result == {"result": "success"} assert result == {"result": "success"}
assert json.loads(custom_client.encrypted_oauth_params) == {"client_id": "new-id"} assert json.loads(custom_client.encrypted_oauth_params) == {"client_id": "new-id"}
cache.delete.assert_called_once() cache.delete.assert_called_once()
mock_session.commit.assert_called_once()
def test_get_custom_oauth_client_params_should_return_empty_when_record_missing( def test_get_custom_oauth_client_params_should_return_empty_when_record_missing(
@ -921,7 +924,6 @@ def test_delete_custom_oauth_client_params_should_delete_record_and_commit(
# Assert # Assert
assert result == {"result": "success"} assert result == {"result": "success"}
mock_session.commit.assert_called_once()
@pytest.mark.parametrize("exists", [True, False]) @pytest.mark.parametrize("exists", [True, False])