opensource/dify
3
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-09-24 23:22:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(api): migrate tools, account, workflow and plugin services to SQLAlchemy 2.0 (#34966)

Browse Source 
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
wdeveloper16
2026-04-12 01:45:27 +00:00
committed by GitHub
co-authored by autofix-ci[bot]
parent 440602f52a
commit 0841b4c663
9 changed files with 111 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/services/account_service.py
+2 -2
View File
@@ -809,11 +809,11 @@ class AccountService:
rest of the system gradually normalizes new inputs.
"""
with session_factory.create_session() as session:
account = session.execute(select(Account).filter_by(email=email)).scalar_one_or_none()
account = session.execute(select(Account).where(Account.email == email)).scalar_one_or_none()
if account or email == email.lower():
return account
return session.execute(select(Account).filter_by(email=email.lower())).scalar_one_or_none()
return session.execute(select(Account).where(Account.email == email.lower())).scalar_one_or_none()
@classmethod
def get_email_code_login_data(cls, token: str) -> dict[str, Any] | None:
api/services/plugin/plugin_migration.py
+14 -10
View File
@@ -13,6 +13,7 @@ import sqlalchemy as sa
import tqdm
from flask import Flask, current_app
from pydantic import TypeAdapter
from sqlalchemy import func, select
from sqlalchemy.orm import Session
from core.agent.entities import AgentToolEntity
@@ -66,7 +67,7 @@ class PluginMigration:
current_time = started_at
with Session(db.engine) as session:
total_tenant_count = session.query(Tenant.id).count()
total_tenant_count = session.scalar(select(func.count(Tenant.id))) or 0
click.echo(click.style(f"Total tenant count: {total_tenant_count}", fg="white"))
@@ -123,9 +124,12 @@ class PluginMigration:
tenant_count = 0
for test_interval in test_intervals:
tenant_count = (
session.query(Tenant.id)
.where(Tenant.created_at.between(current_time, current_time + test_interval))
.count()
session.scalar(
select(func.count(Tenant.id)).where(
Tenant.created_at.between(current_time, current_time + test_interval)
)
)
or 0
)
if tenant_count <= 100:
interval = test_interval
@@ -147,8 +151,8 @@ class PluginMigration:
batch_end = min(current_time + interval, ended_at)
rs = (
session.query(Tenant.id)
rs = session.execute(
select(Tenant.id)
.where(Tenant.created_at.between(current_time, batch_end))
.order_by(Tenant.created_at)
)
@@ -235,7 +239,7 @@ class PluginMigration:
Extract tool tables.
"""
with Session(db.engine) as session:
rs = session.query(BuiltinToolProvider).where(BuiltinToolProvider.tenant_id == tenant_id).all()
rs = session.scalars(select(BuiltinToolProvider).where(BuiltinToolProvider.tenant_id == tenant_id)).all()
result = []
for row in rs:
result.append(ToolProviderID(row.provider).plugin_id)
@@ -249,7 +253,7 @@ class PluginMigration:
"""
with Session(db.engine) as session:
rs = session.query(Workflow).where(Workflow.tenant_id == tenant_id).all()
rs = session.scalars(select(Workflow).where(Workflow.tenant_id == tenant_id)).all()
result = []
for row in rs:
graph = row.graph_dict
@@ -272,7 +276,7 @@ class PluginMigration:
Extract app tables.
"""
with Session(db.engine) as session:
apps = session.query(App).where(App.tenant_id == tenant_id).all()
apps = session.scalars(select(App).where(App.tenant_id == tenant_id)).all()
if not apps:
return []
@@ -280,7 +284,7 @@ class PluginMigration:
app.app_model_config_id for app in apps if app.is_agent or app.mode == AppMode.AGENT_CHAT
]
rs = session.query(AppModelConfig).where(AppModelConfig.id.in_(agent_app_model_config_ids)).all()
rs = session.scalars(select(AppModelConfig).where(AppModelConfig.id.in_(agent_app_model_config_ids))).all()
result = []
for row in rs:
agent_config = row.agent_mode_dict
api/services/rag_pipeline/rag_pipeline_dsl_service.py
+20 -19
View File
@@ -283,7 +283,9 @@ class RagPipelineDslService:
):
raise ValueError("Chunk structure is not compatible with the published pipeline")
if not dataset:
datasets = self._session.query(Dataset).filter_by(tenant_id=account.current_tenant_id).all()
datasets = self._session.scalars(
select(Dataset).where(Dataset.tenant_id == account.current_tenant_id)
).all()
names = [dataset.name for dataset in datasets]
generate_name = generate_incremental_name(names, name)
dataset = Dataset(
@@ -303,8 +305,8 @@ class RagPipelineDslService:
chunk_structure=knowledge_configuration.chunk_structure,
)
if knowledge_configuration.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
dataset_collection_binding = (
self._session.query(DatasetCollectionBinding)
dataset_collection_binding = self._session.scalar(
select(DatasetCollectionBinding)
.where(
DatasetCollectionBinding.provider_name
== knowledge_configuration.embedding_model_provider,
@@ -312,7 +314,7 @@ class RagPipelineDslService:
DatasetCollectionBinding.type == CollectionBindingType.DATASET,
)
.order_by(DatasetCollectionBinding.created_at)
.first()
.limit(1)
)
if not dataset_collection_binding:
@@ -440,8 +442,8 @@ class RagPipelineDslService:
dataset.runtime_mode = DatasetRuntimeMode.RAG_PIPELINE
dataset.chunk_structure = knowledge_configuration.chunk_structure
if knowledge_configuration.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
dataset_collection_binding = (
self._session.query(DatasetCollectionBinding)
dataset_collection_binding = self._session.scalar(
select(DatasetCollectionBinding)
.where(
DatasetCollectionBinding.provider_name
== knowledge_configuration.embedding_model_provider,
@@ -449,7 +451,7 @@ class RagPipelineDslService:
DatasetCollectionBinding.type == CollectionBindingType.DATASET,
)
.order_by(DatasetCollectionBinding.created_at)
.first()
.limit(1)
)
if not dataset_collection_binding:
@@ -591,14 +593,14 @@ class RagPipelineDslService:
IMPORT_INFO_REDIS_EXPIRY,
CheckDependenciesPendingData(pipeline_id=pipeline.id, dependencies=dependencies).model_dump_json(),
)
workflow = (
self._session.query(Workflow)
workflow = self._session.scalar(
select(Workflow)
.where(
Workflow.tenant_id == pipeline.tenant_id,
Workflow.app_id == pipeline.id,
Workflow.version == "draft",
)
.first()
.limit(1)
)
# create draft workflow if not found
@@ -665,14 +667,12 @@ class RagPipelineDslService:
:param pipeline: Pipeline instance
"""
workflow = (
self._session.query(Workflow)
.where(
workflow = self._session.scalar(
select(Workflow).where(
Workflow.tenant_id == pipeline.tenant_id,
Workflow.app_id == pipeline.id,
Workflow.version == "draft",
)
.first()
)
if not workflow:
raise ValueError("Missing draft workflow configuration, please check.")
@@ -904,15 +904,16 @@ class RagPipelineDslService:
):
if rag_pipeline_dataset_create_entity.name:
# check if dataset name already exists
if (
self._session.query(Dataset)
.filter_by(name=rag_pipeline_dataset_create_entity.name, tenant_id=tenant_id)
.first()
if self._session.scalar(
select(Dataset).where(
Dataset.name == rag_pipeline_dataset_create_entity.name,
Dataset.tenant_id == tenant_id,
)
):
raise ValueError(f"Dataset with name {rag_pipeline_dataset_create_entity.name} already exists.")
else:
# generate a random name as Untitled 1 2 3 ...
datasets = self._session.query(Dataset).filter_by(tenant_id=tenant_id).all()
datasets = self._session.scalars(select(Dataset).where(Dataset.tenant_id == tenant_id)).all()
names = [dataset.name for dataset in datasets]
rag_pipeline_dataset_create_entity.name = generate_incremental_name(
names,
api/services/workflow_draft_variable_service.py
+48 -46
View File
@@ -19,7 +19,7 @@ from graphon.variables.segments import (
)
from graphon.variables.types import SegmentType
from graphon.variables.utils import dumps_with_segments
from sqlalchemy import Engine, orm, select
from sqlalchemy import Engine, delete, orm, select
from sqlalchemy.dialects.mysql import insert as mysql_insert
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.orm import Session, sessionmaker
@@ -222,11 +222,10 @@ class WorkflowDraftVariableService:
)
def get_variable(self, variable_id: str) -> WorkflowDraftVariable | None:
return (
self._session.query(WorkflowDraftVariable)
return self._session.scalar(
select(WorkflowDraftVariable)
.options(orm.selectinload(WorkflowDraftVariable.variable_file))
.where(WorkflowDraftVariable.id == variable_id)
.first()
)
def get_draft_variables_by_selectors(
@@ -254,20 +253,21 @@ class WorkflowDraftVariableService:
# Alternatively, a `SELECT` statement could be constructed for each selector and
# combined using `UNION` to fetch all rows.
# Benchmarking indicates that both approaches yield comparable performance.
query = (
self._session.query(WorkflowDraftVariable)
.options(
orm.selectinload(WorkflowDraftVariable.variable_file).selectinload(
WorkflowDraftVariableFile.upload_file
return list(
self._session.scalars(
select(WorkflowDraftVariable)
.options(
orm.selectinload(WorkflowDraftVariable.variable_file).selectinload(
WorkflowDraftVariableFile.upload_file
)
)
.where(
WorkflowDraftVariable.app_id == app_id,
WorkflowDraftVariable.user_id == user_id,
or_(*ors),
)
)
.where(
WorkflowDraftVariable.app_id == app_id,
WorkflowDraftVariable.user_id == user_id,
or_(*ors),
)
)
return query.all()
def list_variables_without_values(
self, app_id: str, page: int, limit: int, user_id: str
@@ -277,18 +277,21 @@ class WorkflowDraftVariableService:
WorkflowDraftVariable.user_id == user_id,
]
total = None
query = self._session.query(WorkflowDraftVariable).where(*criteria)
base_stmt = select(WorkflowDraftVariable).where(*criteria)
if page == 1:
total = query.count()
variables = (
# Do not load the `value` field
query.options(
orm.defer(WorkflowDraftVariable.value, raiseload=True),
from sqlalchemy import func as sa_func
total = self._session.scalar(select(sa_func.count()).select_from(base_stmt.subquery()))
variables = list(
self._session.scalars(
# Do not load the `value` field
base_stmt.options(
orm.defer(WorkflowDraftVariable.value, raiseload=True),
)
.order_by(WorkflowDraftVariable.created_at.desc())
.limit(limit)
.offset((page - 1) * limit)
)
.order_by(WorkflowDraftVariable.created_at.desc())
.limit(limit)
.offset((page - 1) * limit)
.all()
)
return WorkflowDraftVariableList(variables=variables, total=total)
@@ -299,11 +302,13 @@ class WorkflowDraftVariableService:
WorkflowDraftVariable.node_id == node_id,
WorkflowDraftVariable.user_id == user_id,
]
query = self._session.query(WorkflowDraftVariable).where(*criteria)
variables = (
query.options(orm.selectinload(WorkflowDraftVariable.variable_file))
.order_by(WorkflowDraftVariable.created_at.desc())
.all()
variables = list(
self._session.scalars(
select(WorkflowDraftVariable)
.options(orm.selectinload(WorkflowDraftVariable.variable_file))
.where(*criteria)
.order_by(WorkflowDraftVariable.created_at.desc())
)
)
return WorkflowDraftVariableList(variables=variables)
@@ -326,8 +331,8 @@ class WorkflowDraftVariableService:
return self._get_variable(app_id, node_id, name, user_id=user_id)
def _get_variable(self, app_id: str, node_id: str, name: str, user_id: str) -> WorkflowDraftVariable | None:
return (
self._session.query(WorkflowDraftVariable)
return self._session.scalar(
select(WorkflowDraftVariable)
.options(orm.selectinload(WorkflowDraftVariable.variable_file))
.where(
WorkflowDraftVariable.app_id == app_id,
@@ -335,7 +340,6 @@ class WorkflowDraftVariableService:
WorkflowDraftVariable.name == name,
WorkflowDraftVariable.user_id == user_id,
)
.first()
)
def update_variable(
@@ -488,20 +492,20 @@ class WorkflowDraftVariableService:
self._session.delete(variable)
def delete_user_workflow_variables(self, app_id: str, user_id: str):
(
self._session.query(WorkflowDraftVariable)
self._session.execute(
delete(WorkflowDraftVariable)
.where(
WorkflowDraftVariable.app_id == app_id,
WorkflowDraftVariable.user_id == user_id,
)
.delete(synchronize_session=False)
.execution_options(synchronize_session=False)
)
def delete_app_workflow_variables(self, app_id: str):
(
self._session.query(WorkflowDraftVariable)
self._session.execute(
delete(WorkflowDraftVariable)
.where(WorkflowDraftVariable.app_id == app_id)
.delete(synchronize_session=False)
.execution_options(synchronize_session=False)
)
def delete_workflow_draft_variable_file(self, deletions: list[DraftVarFileDeletion]):
@@ -540,14 +544,14 @@ class WorkflowDraftVariableService:
return self._delete_node_variables(app_id, node_id, user_id=user_id)
def _delete_node_variables(self, app_id: str, node_id: str, user_id: str):
(
self._session.query(WorkflowDraftVariable)
self._session.execute(
delete(WorkflowDraftVariable)
.where(
WorkflowDraftVariable.app_id == app_id,
WorkflowDraftVariable.node_id == node_id,
WorkflowDraftVariable.user_id == user_id,
)
.delete(synchronize_session=False)
.execution_options(synchronize_session=False)
)
def _get_conversation_id_from_draft_variable(self, app_id: str, user_id: str) -> str | None:
@@ -588,13 +592,11 @@ class WorkflowDraftVariableService:
conv_id = self._get_conversation_id_from_draft_variable(workflow.app_id, account_id)
if conv_id is not None:
conversation = (
self._session.query(Conversation)
.where(
conversation = self._session.scalar(
select(Conversation).where(
Conversation.id == conv_id,
Conversation.app_id == workflow.app_id,
)
.first()
)
# Only return the conversation ID if it exists and is valid (has a correspond conversation record in DB).
if conversation is not None:
api/services/workflow_service.py
+2 -4
View File
@@ -1512,14 +1512,12 @@ class WorkflowService:
# Don't use workflow.tool_published as it's not accurate for specific workflow versions
# Check if there's a tool provider using this specific workflow version
tool_provider = (
session.query(WorkflowToolProvider)
.where(
tool_provider = session.scalar(
select(WorkflowToolProvider).where(
WorkflowToolProvider.tenant_id == workflow.tenant_id,
WorkflowToolProvider.app_id == workflow.app_id,
WorkflowToolProvider.version == workflow.version,
)
.first()
)
if tool_provider: