From 8a22cc06c9e9a20fd9e83daa220a5e55733177b2 Mon Sep 17 00:00:00 2001 From: Renzo <170978465+RenzoMXD@users.noreply.github.com> Date: Thu, 19 Mar 2026 07:48:29 +0100 Subject: [PATCH] refactor: migrate db.session.query to select in infra layer (#33694) Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- api/commands/plugin.py | 73 +++++++++++-------- api/commands/storage.py | 22 ++++-- api/commands/system.py | 15 ++-- api/commands/vector.py | 45 ++++++------ .../event_handlers/create_document_index.py | 7 +- ...aset_join_when_app_model_config_updated.py | 8 +- ...oin_when_app_published_workflow_updated.py | 8 +- api/extensions/ext_login.py | 20 ++--- api/factories/file_factory.py | 6 +- api/schedule/check_upgradable_plugin_task.py | 9 +-- api/schedule/clean_embedding_cache_task.py | 10 +-- api/schedule/clean_unused_datasets_task.py | 10 +-- api/schedule/create_tidb_serverless_task.py | 3 +- .../mail_clean_document_notify_task.py | 12 +-- 14 files changed, 131 insertions(+), 117 deletions(-) diff --git a/api/commands/plugin.py b/api/commands/plugin.py index 2dfbd73b3a..c34391025a 100644 --- a/api/commands/plugin.py +++ b/api/commands/plugin.py @@ -1,9 +1,11 @@ import json import logging -from typing import Any +from typing import Any, cast import click from pydantic import TypeAdapter +from sqlalchemy import delete, select +from sqlalchemy.engine import CursorResult from configs import dify_config from core.helper import encrypter @@ -48,14 +50,15 @@ def setup_system_tool_oauth_client(provider, client_params): click.echo(click.style(f"Error parsing client params: {str(e)}", fg="red")) return - deleted_count = ( - db.session.query(ToolOAuthSystemClient) - .filter_by( - provider=provider_name, - plugin_id=plugin_id, - ) - .delete() - ) + deleted_count = cast( + CursorResult, + db.session.execute( + delete(ToolOAuthSystemClient).where( + ToolOAuthSystemClient.provider == provider_name, + ToolOAuthSystemClient.plugin_id == plugin_id, + ) + ), + ).rowcount if deleted_count > 0: click.echo(click.style(f"Deleted {deleted_count} existing oauth client params.", fg="yellow")) @@ -97,14 +100,15 @@ def setup_system_trigger_oauth_client(provider, client_params): click.echo(click.style(f"Error parsing client params: {str(e)}", fg="red")) return - deleted_count = ( - db.session.query(TriggerOAuthSystemClient) - .filter_by( - provider=provider_name, - plugin_id=plugin_id, - ) - .delete() - ) + deleted_count = cast( + CursorResult, + db.session.execute( + delete(TriggerOAuthSystemClient).where( + TriggerOAuthSystemClient.provider == provider_name, + TriggerOAuthSystemClient.plugin_id == plugin_id, + ) + ), + ).rowcount if deleted_count > 0: click.echo(click.style(f"Deleted {deleted_count} existing oauth client params.", fg="yellow")) @@ -139,14 +143,15 @@ def setup_datasource_oauth_client(provider, client_params): return click.echo(click.style(f"Ready to delete existing oauth client params: {provider_name}", fg="yellow")) - deleted_count = ( - db.session.query(DatasourceOauthParamConfig) - .filter_by( - provider=provider_name, - plugin_id=plugin_id, - ) - .delete() - ) + deleted_count = cast( + CursorResult, + db.session.execute( + delete(DatasourceOauthParamConfig).where( + DatasourceOauthParamConfig.provider == provider_name, + DatasourceOauthParamConfig.plugin_id == plugin_id, + ) + ), + ).rowcount if deleted_count > 0: click.echo(click.style(f"Deleted {deleted_count} existing oauth client params.", fg="yellow")) @@ -192,7 +197,9 @@ def transform_datasource_credentials(environment: str): # deal notion credentials deal_notion_count = 0 - notion_credentials = db.session.query(DataSourceOauthBinding).filter_by(provider="notion").all() + notion_credentials = db.session.scalars( + select(DataSourceOauthBinding).where(DataSourceOauthBinding.provider == "notion") + ).all() if notion_credentials: notion_credentials_tenant_mapping: dict[str, list[DataSourceOauthBinding]] = {} for notion_credential in notion_credentials: @@ -201,7 +208,7 @@ def transform_datasource_credentials(environment: str): notion_credentials_tenant_mapping[tenant_id] = [] notion_credentials_tenant_mapping[tenant_id].append(notion_credential) for tenant_id, notion_tenant_credentials in notion_credentials_tenant_mapping.items(): - tenant = db.session.query(Tenant).filter_by(id=tenant_id).first() + tenant = db.session.scalar(select(Tenant).where(Tenant.id == tenant_id)) if not tenant: continue try: @@ -250,7 +257,9 @@ def transform_datasource_credentials(environment: str): db.session.commit() # deal firecrawl credentials deal_firecrawl_count = 0 - firecrawl_credentials = db.session.query(DataSourceApiKeyAuthBinding).filter_by(provider="firecrawl").all() + firecrawl_credentials = db.session.scalars( + select(DataSourceApiKeyAuthBinding).where(DataSourceApiKeyAuthBinding.provider == "firecrawl") + ).all() if firecrawl_credentials: firecrawl_credentials_tenant_mapping: dict[str, list[DataSourceApiKeyAuthBinding]] = {} for firecrawl_credential in firecrawl_credentials: @@ -259,7 +268,7 @@ def transform_datasource_credentials(environment: str): firecrawl_credentials_tenant_mapping[tenant_id] = [] firecrawl_credentials_tenant_mapping[tenant_id].append(firecrawl_credential) for tenant_id, firecrawl_tenant_credentials in firecrawl_credentials_tenant_mapping.items(): - tenant = db.session.query(Tenant).filter_by(id=tenant_id).first() + tenant = db.session.scalar(select(Tenant).where(Tenant.id == tenant_id)) if not tenant: continue try: @@ -312,7 +321,9 @@ def transform_datasource_credentials(environment: str): db.session.commit() # deal jina credentials deal_jina_count = 0 - jina_credentials = db.session.query(DataSourceApiKeyAuthBinding).filter_by(provider="jinareader").all() + jina_credentials = db.session.scalars( + select(DataSourceApiKeyAuthBinding).where(DataSourceApiKeyAuthBinding.provider == "jinareader") + ).all() if jina_credentials: jina_credentials_tenant_mapping: dict[str, list[DataSourceApiKeyAuthBinding]] = {} for jina_credential in jina_credentials: @@ -321,7 +332,7 @@ def transform_datasource_credentials(environment: str): jina_credentials_tenant_mapping[tenant_id] = [] jina_credentials_tenant_mapping[tenant_id].append(jina_credential) for tenant_id, jina_tenant_credentials in jina_credentials_tenant_mapping.items(): - tenant = db.session.query(Tenant).filter_by(id=tenant_id).first() + tenant = db.session.scalar(select(Tenant).where(Tenant.id == tenant_id)) if not tenant: continue try: diff --git a/api/commands/storage.py b/api/commands/storage.py index fa890a855a..f23b17680a 100644 --- a/api/commands/storage.py +++ b/api/commands/storage.py @@ -1,7 +1,10 @@ import json +from typing import cast import click import sqlalchemy as sa +from sqlalchemy import update +from sqlalchemy.engine import CursorResult from configs import dify_config from extensions.ext_database import db @@ -740,14 +743,17 @@ def migrate_oss( else: try: source_storage_type = StorageType.LOCAL if is_source_local else StorageType.OPENDAL - updated = ( - db.session.query(UploadFile) - .where( - UploadFile.storage_type == source_storage_type, - UploadFile.key.in_(copied_upload_file_keys), - ) - .update({UploadFile.storage_type: dify_config.STORAGE_TYPE}, synchronize_session=False) - ) + updated = cast( + CursorResult, + db.session.execute( + update(UploadFile) + .where( + UploadFile.storage_type == source_storage_type, + UploadFile.key.in_(copied_upload_file_keys), + ) + .values(storage_type=dify_config.STORAGE_TYPE) + ), + ).rowcount db.session.commit() click.echo(click.style(f"Updated storage_type for {updated} upload_files records.", fg="green")) except Exception as e: diff --git a/api/commands/system.py b/api/commands/system.py index 604f0e34d0..39b2e991ed 100644 --- a/api/commands/system.py +++ b/api/commands/system.py @@ -2,6 +2,7 @@ import logging import click import sqlalchemy as sa +from sqlalchemy import delete, select, update from sqlalchemy.orm import sessionmaker from configs import dify_config @@ -41,7 +42,7 @@ def reset_encrypt_key_pair(): click.echo(click.style("This command is only for SELF_HOSTED installations.", fg="red")) return with sessionmaker(db.engine, expire_on_commit=False).begin() as session: - tenants = session.query(Tenant).all() + tenants = session.scalars(select(Tenant)).all() for tenant in tenants: if not tenant: click.echo(click.style("No workspaces found. Run /install first.", fg="red")) @@ -49,8 +50,8 @@ def reset_encrypt_key_pair(): tenant.encrypt_public_key = generate_key_pair(tenant.id) - session.query(Provider).where(Provider.provider_type == "custom", Provider.tenant_id == tenant.id).delete() - session.query(ProviderModel).where(ProviderModel.tenant_id == tenant.id).delete() + session.execute(delete(Provider).where(Provider.provider_type == "custom", Provider.tenant_id == tenant.id)) + session.execute(delete(ProviderModel).where(ProviderModel.tenant_id == tenant.id)) click.echo( click.style( @@ -93,7 +94,7 @@ def convert_to_agent_apps(): app_id = str(i.id) if app_id not in proceeded_app_ids: proceeded_app_ids.append(app_id) - app = db.session.query(App).where(App.id == app_id).first() + app = db.session.scalar(select(App).where(App.id == app_id)) if app is not None: apps.append(app) @@ -108,8 +109,8 @@ def convert_to_agent_apps(): db.session.commit() # update conversation mode to agent - db.session.query(Conversation).where(Conversation.app_id == app.id).update( - {Conversation.mode: AppMode.AGENT_CHAT} + db.session.execute( + update(Conversation).where(Conversation.app_id == app.id).values(mode=AppMode.AGENT_CHAT) ) db.session.commit() @@ -177,7 +178,7 @@ where sites.id is null limit 1000""" continue try: - app = db.session.query(App).where(App.id == app_id).first() + app = db.session.scalar(select(App).where(App.id == app_id)) if not app: logger.info("App %s not found", app_id) continue diff --git a/api/commands/vector.py b/api/commands/vector.py index 52ce26c26d..4cf11c9ad1 100644 --- a/api/commands/vector.py +++ b/api/commands/vector.py @@ -41,14 +41,13 @@ def migrate_annotation_vector_database(): # get apps info per_page = 50 with sessionmaker(db.engine, expire_on_commit=False).begin() as session: - apps = ( - session.query(App) + apps = session.scalars( + select(App) .where(App.status == "normal") .order_by(App.created_at.desc()) .limit(per_page) .offset((page - 1) * per_page) - .all() - ) + ).all() if not apps: break except SQLAlchemyError: @@ -63,8 +62,8 @@ def migrate_annotation_vector_database(): try: click.echo(f"Creating app annotation index: {app.id}") with sessionmaker(db.engine, expire_on_commit=False).begin() as session: - app_annotation_setting = ( - session.query(AppAnnotationSetting).where(AppAnnotationSetting.app_id == app.id).first() + app_annotation_setting = session.scalar( + select(AppAnnotationSetting).where(AppAnnotationSetting.app_id == app.id).limit(1) ) if not app_annotation_setting: @@ -72,10 +71,10 @@ def migrate_annotation_vector_database(): click.echo(f"App annotation setting disabled: {app.id}") continue # get dataset_collection_binding info - dataset_collection_binding = ( - session.query(DatasetCollectionBinding) - .where(DatasetCollectionBinding.id == app_annotation_setting.collection_binding_id) - .first() + dataset_collection_binding = session.scalar( + select(DatasetCollectionBinding).where( + DatasetCollectionBinding.id == app_annotation_setting.collection_binding_id + ) ) if not dataset_collection_binding: click.echo(f"App annotation collection binding not found: {app.id}") @@ -205,11 +204,11 @@ def migrate_knowledge_vector_database(): collection_name = Dataset.gen_collection_name_by_id(dataset_id) elif vector_type == VectorType.QDRANT: if dataset.collection_binding_id: - dataset_collection_binding = ( - db.session.query(DatasetCollectionBinding) - .where(DatasetCollectionBinding.id == dataset.collection_binding_id) - .one_or_none() - ) + dataset_collection_binding = db.session.execute( + select(DatasetCollectionBinding).where( + DatasetCollectionBinding.id == dataset.collection_binding_id + ) + ).scalar_one_or_none() if dataset_collection_binding: collection_name = dataset_collection_binding.collection_name else: @@ -334,7 +333,7 @@ def add_qdrant_index(field: str): create_count = 0 try: - bindings = db.session.query(DatasetCollectionBinding).all() + bindings = db.session.scalars(select(DatasetCollectionBinding)).all() if not bindings: click.echo(click.style("No dataset collection bindings found.", fg="red")) return @@ -421,10 +420,10 @@ def old_metadata_migration(): if field.value == key: break else: - dataset_metadata = ( - db.session.query(DatasetMetadata) + dataset_metadata = db.session.scalar( + select(DatasetMetadata) .where(DatasetMetadata.dataset_id == document.dataset_id, DatasetMetadata.name == key) - .first() + .limit(1) ) if not dataset_metadata: dataset_metadata = DatasetMetadata( @@ -436,7 +435,7 @@ def old_metadata_migration(): ) db.session.add(dataset_metadata) db.session.flush() - dataset_metadata_binding = DatasetMetadataBinding( + dataset_metadata_binding: DatasetMetadataBinding | None = DatasetMetadataBinding( tenant_id=document.tenant_id, dataset_id=document.dataset_id, metadata_id=dataset_metadata.id, @@ -445,14 +444,14 @@ def old_metadata_migration(): ) db.session.add(dataset_metadata_binding) else: - dataset_metadata_binding = ( - db.session.query(DatasetMetadataBinding) # type: ignore + dataset_metadata_binding = db.session.scalar( + select(DatasetMetadataBinding) .where( DatasetMetadataBinding.dataset_id == document.dataset_id, DatasetMetadataBinding.document_id == document.id, DatasetMetadataBinding.metadata_id == dataset_metadata.id, ) - .first() + .limit(1) ) if not dataset_metadata_binding: dataset_metadata_binding = DatasetMetadataBinding( diff --git a/api/events/event_handlers/create_document_index.py b/api/events/event_handlers/create_document_index.py index 76de5a0740..b7e7a6e60f 100644 --- a/api/events/event_handlers/create_document_index.py +++ b/api/events/event_handlers/create_document_index.py @@ -3,6 +3,7 @@ import logging import time import click +from sqlalchemy import select from werkzeug.exceptions import NotFound from core.indexing_runner import DocumentIsPausedError, IndexingRunner @@ -24,13 +25,11 @@ def handle(sender, **kwargs): for document_id in document_ids: logger.info(click.style(f"Start process document: {document_id}", fg="green")) - document = ( - db.session.query(Document) - .where( + document = db.session.scalar( + select(Document).where( Document.id == document_id, Document.dataset_id == dataset_id, ) - .first() ) if not document: diff --git a/api/events/event_handlers/update_app_dataset_join_when_app_model_config_updated.py b/api/events/event_handlers/update_app_dataset_join_when_app_model_config_updated.py index b70c2183d2..4709534ae6 100644 --- a/api/events/event_handlers/update_app_dataset_join_when_app_model_config_updated.py +++ b/api/events/event_handlers/update_app_dataset_join_when_app_model_config_updated.py @@ -1,6 +1,6 @@ from typing import Any, cast -from sqlalchemy import select +from sqlalchemy import delete, select from events.app_event import app_model_config_was_updated from extensions.ext_database import db @@ -31,9 +31,9 @@ def handle(sender, **kwargs): if removed_dataset_ids: for dataset_id in removed_dataset_ids: - db.session.query(AppDatasetJoin).where( - AppDatasetJoin.app_id == app.id, AppDatasetJoin.dataset_id == dataset_id - ).delete() + db.session.execute( + delete(AppDatasetJoin).where(AppDatasetJoin.app_id == app.id, AppDatasetJoin.dataset_id == dataset_id) + ) if added_dataset_ids: for dataset_id in added_dataset_ids: diff --git a/api/events/event_handlers/update_app_dataset_join_when_app_published_workflow_updated.py b/api/events/event_handlers/update_app_dataset_join_when_app_published_workflow_updated.py index 92bc9db075..20852b818e 100644 --- a/api/events/event_handlers/update_app_dataset_join_when_app_published_workflow_updated.py +++ b/api/events/event_handlers/update_app_dataset_join_when_app_published_workflow_updated.py @@ -1,6 +1,6 @@ from typing import cast -from sqlalchemy import select +from sqlalchemy import delete, select from core.workflow.nodes.knowledge_retrieval.entities import KnowledgeRetrievalNodeData from dify_graph.nodes import BuiltinNodeTypes @@ -31,9 +31,9 @@ def handle(sender, **kwargs): if removed_dataset_ids: for dataset_id in removed_dataset_ids: - db.session.query(AppDatasetJoin).where( - AppDatasetJoin.app_id == app.id, AppDatasetJoin.dataset_id == dataset_id - ).delete() + db.session.execute( + delete(AppDatasetJoin).where(AppDatasetJoin.app_id == app.id, AppDatasetJoin.dataset_id == dataset_id) + ) if added_dataset_ids: for dataset_id in added_dataset_ids: diff --git a/api/extensions/ext_login.py b/api/extensions/ext_login.py index 74299956c0..02e50a90fc 100644 --- a/api/extensions/ext_login.py +++ b/api/extensions/ext_login.py @@ -3,6 +3,7 @@ import json import flask_login from flask import Response, request from flask_login import user_loaded_from_request, user_logged_in +from sqlalchemy import select from werkzeug.exceptions import NotFound, Unauthorized from configs import dify_config @@ -34,16 +35,15 @@ def load_user_from_request(request_from_flask_login): if admin_api_key and admin_api_key == auth_token: workspace_id = request.headers.get("X-WORKSPACE-ID") if workspace_id: - tenant_account_join = ( - db.session.query(Tenant, TenantAccountJoin) + tenant_account_join = db.session.execute( + select(Tenant, TenantAccountJoin) .where(Tenant.id == workspace_id) .where(TenantAccountJoin.tenant_id == Tenant.id) .where(TenantAccountJoin.role == "owner") - .one_or_none() - ) + ).one_or_none() if tenant_account_join: tenant, ta = tenant_account_join - account = db.session.query(Account).filter_by(id=ta.account_id).first() + account = db.session.scalar(select(Account).where(Account.id == ta.account_id)) if account: account.current_tenant = tenant return account @@ -70,7 +70,7 @@ def load_user_from_request(request_from_flask_login): end_user_id = decoded.get("end_user_id") if not end_user_id: raise Unauthorized("Invalid Authorization token.") - end_user = db.session.query(EndUser).where(EndUser.id == end_user_id).first() + end_user = db.session.scalar(select(EndUser).where(EndUser.id == end_user_id)) if not end_user: raise NotFound("End user not found.") return end_user @@ -80,7 +80,7 @@ def load_user_from_request(request_from_flask_login): decoded = PassportService().verify(auth_token) end_user_id = decoded.get("end_user_id") if end_user_id: - end_user = db.session.query(EndUser).where(EndUser.id == end_user_id).first() + end_user = db.session.scalar(select(EndUser).where(EndUser.id == end_user_id)) if not end_user: raise NotFound("End user not found.") return end_user @@ -90,11 +90,11 @@ def load_user_from_request(request_from_flask_login): server_code = request.view_args.get("server_code") if request.view_args else None if not server_code: raise Unauthorized("Invalid Authorization token.") - app_mcp_server = db.session.query(AppMCPServer).where(AppMCPServer.server_code == server_code).first() + app_mcp_server = db.session.scalar(select(AppMCPServer).where(AppMCPServer.server_code == server_code).limit(1)) if not app_mcp_server: raise NotFound("App MCP server not found.") - end_user = ( - db.session.query(EndUser).where(EndUser.session_id == app_mcp_server.id, EndUser.type == "mcp").first() + end_user = db.session.scalar( + select(EndUser).where(EndUser.session_id == app_mcp_server.id, EndUser.type == "mcp").limit(1) ) if not end_user: raise NotFound("End user not found.") diff --git a/api/factories/file_factory.py b/api/factories/file_factory.py index ef55fe53c5..cb07ba58ae 100644 --- a/api/factories/file_factory.py +++ b/api/factories/file_factory.py @@ -424,13 +424,11 @@ def _build_from_datasource_file( datasource_file_id = mapping.get("datasource_file_id") if not datasource_file_id: raise ValueError(f"DatasourceFile {datasource_file_id} not found") - datasource_file = ( - db.session.query(UploadFile) - .where( + datasource_file = db.session.scalar( + select(UploadFile).where( UploadFile.id == datasource_file_id, UploadFile.tenant_id == tenant_id, ) - .first() ) if datasource_file is None: diff --git a/api/schedule/check_upgradable_plugin_task.py b/api/schedule/check_upgradable_plugin_task.py index 13d2f24ca0..cf223f6e9e 100644 --- a/api/schedule/check_upgradable_plugin_task.py +++ b/api/schedule/check_upgradable_plugin_task.py @@ -3,6 +3,7 @@ import math import time import click +from sqlalchemy import select import app from core.helper.marketplace import fetch_global_plugin_manifest @@ -28,17 +29,15 @@ def check_upgradable_plugin_task(): now_seconds_of_day = time.time() % 86400 - 30 # we assume the tz is UTC click.echo(click.style(f"Now seconds of day: {now_seconds_of_day}", fg="green")) - strategies = ( - db.session.query(TenantPluginAutoUpgradeStrategy) - .where( + strategies = db.session.scalars( + select(TenantPluginAutoUpgradeStrategy).where( TenantPluginAutoUpgradeStrategy.upgrade_time_of_day >= now_seconds_of_day, TenantPluginAutoUpgradeStrategy.upgrade_time_of_day < now_seconds_of_day + AUTO_UPGRADE_MINIMAL_CHECKING_INTERVAL, TenantPluginAutoUpgradeStrategy.strategy_setting != TenantPluginAutoUpgradeStrategy.StrategySetting.DISABLED, ) - .all() - ) + ).all() total_strategies = len(strategies) click.echo(click.style(f"Total strategies: {total_strategies}", fg="green")) diff --git a/api/schedule/clean_embedding_cache_task.py b/api/schedule/clean_embedding_cache_task.py index 2b74fb2dd0..04c954875f 100644 --- a/api/schedule/clean_embedding_cache_task.py +++ b/api/schedule/clean_embedding_cache_task.py @@ -2,7 +2,7 @@ import datetime import time import click -from sqlalchemy import text +from sqlalchemy import select, text from sqlalchemy.exc import SQLAlchemyError import app @@ -19,14 +19,12 @@ def clean_embedding_cache_task(): thirty_days_ago = datetime.datetime.now() - datetime.timedelta(days=clean_days) while True: try: - embedding_ids = ( - db.session.query(Embedding.id) + embedding_ids = db.session.scalars( + select(Embedding.id) .where(Embedding.created_at < thirty_days_ago) .order_by(Embedding.created_at.desc()) .limit(100) - .all() - ) - embedding_ids = [embedding_id[0] for embedding_id in embedding_ids] + ).all() except SQLAlchemyError: raise if embedding_ids: diff --git a/api/schedule/clean_unused_datasets_task.py b/api/schedule/clean_unused_datasets_task.py index d9fb6a24f1..0b0fc1b229 100644 --- a/api/schedule/clean_unused_datasets_task.py +++ b/api/schedule/clean_unused_datasets_task.py @@ -3,7 +3,7 @@ import time from typing import TypedDict import click -from sqlalchemy import func, select +from sqlalchemy import func, select, update from sqlalchemy.exc import SQLAlchemyError import app @@ -51,7 +51,7 @@ def clean_unused_datasets_task(): try: # Subquery for counting new documents document_subquery_new = ( - db.session.query(Document.dataset_id, func.count(Document.id).label("document_count")) + select(Document.dataset_id, func.count(Document.id).label("document_count")) .where( Document.indexing_status == "completed", Document.enabled == True, @@ -64,7 +64,7 @@ def clean_unused_datasets_task(): # Subquery for counting old documents document_subquery_old = ( - db.session.query(Document.dataset_id, func.count(Document.id).label("document_count")) + select(Document.dataset_id, func.count(Document.id).label("document_count")) .where( Document.indexing_status == "completed", Document.enabled == True, @@ -142,8 +142,8 @@ def clean_unused_datasets_task(): index_processor.clean(dataset, None) # Update document - db.session.query(Document).filter_by(dataset_id=dataset.id).update( - {Document.enabled: False} + db.session.execute( + update(Document).where(Document.dataset_id == dataset.id).values(enabled=False) ) db.session.commit() click.echo(click.style(f"Cleaned unused dataset {dataset.id} from db success!", fg="green")) diff --git a/api/schedule/create_tidb_serverless_task.py b/api/schedule/create_tidb_serverless_task.py index ed46c1c70a..8b9d973d6d 100644 --- a/api/schedule/create_tidb_serverless_task.py +++ b/api/schedule/create_tidb_serverless_task.py @@ -1,6 +1,7 @@ import time import click +from sqlalchemy import func, select import app from configs import dify_config @@ -20,7 +21,7 @@ def create_tidb_serverless_task(): try: # check the number of idle tidb serverless idle_tidb_serverless_number = ( - db.session.query(TidbAuthBinding).where(TidbAuthBinding.active == False).count() + db.session.scalar(select(func.count(TidbAuthBinding.id)).where(TidbAuthBinding.active == False)) or 0 ) if idle_tidb_serverless_number >= tidb_serverless_number: break diff --git a/api/schedule/mail_clean_document_notify_task.py b/api/schedule/mail_clean_document_notify_task.py index d738bf46fa..8479cdfb0c 100644 --- a/api/schedule/mail_clean_document_notify_task.py +++ b/api/schedule/mail_clean_document_notify_task.py @@ -49,16 +49,18 @@ def mail_clean_document_notify_task(): if plan != CloudPlan.SANDBOX: knowledge_details = [] # check tenant - tenant = db.session.query(Tenant).where(Tenant.id == tenant_id).first() + tenant = db.session.scalar(select(Tenant).where(Tenant.id == tenant_id)) if not tenant: continue # check current owner - current_owner_join = ( - db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, role="owner").first() + current_owner_join = db.session.scalar( + select(TenantAccountJoin) + .where(TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.role == "owner") + .limit(1) ) if not current_owner_join: continue - account = db.session.query(Account).where(Account.id == current_owner_join.account_id).first() + account = db.session.scalar(select(Account).where(Account.id == current_owner_join.account_id)) if not account: continue @@ -71,7 +73,7 @@ def mail_clean_document_notify_task(): ) for dataset_id, document_ids in dataset_auto_dataset_map.items(): - dataset = db.session.query(Dataset).where(Dataset.id == dataset_id).first() + dataset = db.session.scalar(select(Dataset).where(Dataset.id == dataset_id)) if dataset: document_count = len(document_ids) knowledge_details.append(rf"Knowledge base {dataset.name}: {document_count} documents")