Migrate SQLAlchemy from 1.x to 2.0 with automated and manual adjustments (#23224)

Co-authored-by: Yongtao Huang <99629139+hyongtao-db@users.noreply.github.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

api/core/agent/base_agent_runner.py

@@ -334,7 +334,8 @@ class BaseAgentRunner(AppRunner):
         """
         Save agent thought
         """
-        agent_thought = db.session.query(MessageAgentThought).where(MessageAgentThought.id == agent_thought_id).first()
+        stmt = select(MessageAgentThought).where(MessageAgentThought.id == agent_thought_id)
+        agent_thought = db.session.scalar(stmt)
         if not agent_thought:
             raise ValueError("agent thought not found")
 
@@ -492,7 +493,8 @@ class BaseAgentRunner(AppRunner):
         return result
 
     def organize_agent_user_prompt(self, message: Message) -> UserPromptMessage:
-        files = db.session.query(MessageFile).where(MessageFile.message_id == message.id).all()
+        stmt = select(MessageFile).where(MessageFile.message_id == message.id)
+        files = db.session.scalars(stmt).all()
         if not files:
             return UserPromptMessage(content=message.query)
         if message.app_model_config:

api/core/app/apps/advanced_chat/app_runner.py

@@ -74,6 +74,7 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
 
         with Session(db.engine, expire_on_commit=False) as session:
             app_record = session.scalar(select(App).where(App.id == app_config.app_id))
+
         if not app_record:
             raise ValueError("App not found")
 

api/core/app/apps/agent_chat/app_runner.py

@@ -1,6 +1,8 @@
 import logging
 from typing import cast
 
+from sqlalchemy import select
+
 from core.agent.cot_chat_agent_runner import CotChatAgentRunner
 from core.agent.cot_completion_agent_runner import CotCompletionAgentRunner
 from core.agent.entities import AgentEntity
@@ -44,8 +46,8 @@ class AgentChatAppRunner(AppRunner):
         """
         app_config = application_generate_entity.app_config
         app_config = cast(AgentChatAppConfig, app_config)
-
-        app_record = db.session.query(App).where(App.id == app_config.app_id).first()
+        app_stmt = select(App).where(App.id == app_config.app_id)
+        app_record = db.session.scalar(app_stmt)
         if not app_record:
             raise ValueError("App not found")
 
@@ -182,11 +184,12 @@ class AgentChatAppRunner(AppRunner):
 
         if {ModelFeature.MULTI_TOOL_CALL, ModelFeature.TOOL_CALL}.intersection(model_schema.features or []):
             agent_entity.strategy = AgentEntity.Strategy.FUNCTION_CALLING
-
-        conversation_result = db.session.query(Conversation).where(Conversation.id == conversation.id).first()
+        conversation_stmt = select(Conversation).where(Conversation.id == conversation.id)
+        conversation_result = db.session.scalar(conversation_stmt)
         if conversation_result is None:
             raise ValueError("Conversation not found")
-        message_result = db.session.query(Message).where(Message.id == message.id).first()
+        msg_stmt = select(Message).where(Message.id == message.id)
+        message_result = db.session.scalar(msg_stmt)
         if message_result is None:
             raise ValueError("Message not found")
         db.session.close()

api/core/app/apps/chat/app_runner.py

@@ -1,6 +1,8 @@
 import logging
 from typing import cast
 
+from sqlalchemy import select
+
 from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
 from core.app.apps.base_app_runner import AppRunner
 from core.app.apps.chat.app_config_manager import ChatAppConfig
@@ -42,8 +44,8 @@ class ChatAppRunner(AppRunner):
         """
         app_config = application_generate_entity.app_config
         app_config = cast(ChatAppConfig, app_config)
-
-        app_record = db.session.query(App).where(App.id == app_config.app_id).first()
+        stmt = select(App).where(App.id == app_config.app_id)
+        app_record = db.session.scalar(stmt)
         if not app_record:
             raise ValueError("App not found")
 

api/core/app/apps/completion/app_generator.py

@@ -6,6 +6,7 @@ from typing import Any, Literal, Union, overload
 
 from flask import Flask, copy_current_request_context, current_app
 from pydantic import ValidationError
+from sqlalchemy import select
 
 from configs import dify_config
 from core.app.app_config.easy_ui_based_app.model_config.converter import ModelConfigConverter
@@ -248,17 +249,14 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
         :param invoke_from: invoke from source
         :param stream: is stream
         """
-        message = (
-            db.session.query(Message)
-            .where(
-                Message.id == message_id,
-                Message.app_id == app_model.id,
-                Message.from_source == ("api" if isinstance(user, EndUser) else "console"),
-                Message.from_end_user_id == (user.id if isinstance(user, EndUser) else None),
-                Message.from_account_id == (user.id if isinstance(user, Account) else None),
-            )
-            .first()
+        stmt = select(Message).where(
+            Message.id == message_id,
+            Message.app_id == app_model.id,
+            Message.from_source == ("api" if isinstance(user, EndUser) else "console"),
+            Message.from_end_user_id == (user.id if isinstance(user, EndUser) else None),
+            Message.from_account_id == (user.id if isinstance(user, Account) else None),
         )
+        message = db.session.scalar(stmt)
 
         if not message:
             raise MessageNotExistsError()

api/core/app/apps/completion/app_runner.py

@@ -1,6 +1,8 @@
 import logging
 from typing import cast
 
+from sqlalchemy import select
+
 from core.app.apps.base_app_queue_manager import AppQueueManager
 from core.app.apps.base_app_runner import AppRunner
 from core.app.apps.completion.app_config_manager import CompletionAppConfig
@@ -35,8 +37,8 @@ class CompletionAppRunner(AppRunner):
         """
         app_config = application_generate_entity.app_config
         app_config = cast(CompletionAppConfig, app_config)
-
-        app_record = db.session.query(App).where(App.id == app_config.app_id).first()
+        stmt = select(App).where(App.id == app_config.app_id)
+        app_record = db.session.scalar(stmt)
         if not app_record:
             raise ValueError("App not found")
 

api/core/app/apps/message_based_app_generator.py

@@ -86,11 +86,10 @@ class MessageBasedAppGenerator(BaseAppGenerator):
 
     def _get_app_model_config(self, app_model: App, conversation: Optional[Conversation] = None) -> AppModelConfig:
         if conversation:
-            app_model_config = (
-                db.session.query(AppModelConfig)
-                .where(AppModelConfig.id == conversation.app_model_config_id, AppModelConfig.app_id == app_model.id)
-                .first()
+            stmt = select(AppModelConfig).where(
+                AppModelConfig.id == conversation.app_model_config_id, AppModelConfig.app_id == app_model.id
             )
+            app_model_config = db.session.scalar(stmt)
 
             if not app_model_config:
                 raise AppModelConfigBrokenError()

api/core/app/features/annotation_reply/annotation_reply.py

@@ -1,6 +1,8 @@
 import logging
 from typing import Optional
 
+from sqlalchemy import select
+
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.rag.datasource.vdb.vector_factory import Vector
 from extensions.ext_database import db
@@ -25,9 +27,8 @@ class AnnotationReplyFeature:
         :param invoke_from: invoke from
         :return:
         """
-        annotation_setting = (
-            db.session.query(AppAnnotationSetting).where(AppAnnotationSetting.app_id == app_record.id).first()
-        )
+        stmt = select(AppAnnotationSetting).where(AppAnnotationSetting.app_id == app_record.id)
+        annotation_setting = db.session.scalar(stmt)
 
         if not annotation_setting:
             return None

api/core/app/task_pipeline/message_cycle_manager.py

@@ -86,7 +86,8 @@ class MessageCycleManager:
     def _generate_conversation_name_worker(self, flask_app: Flask, conversation_id: str, query: str):
         with flask_app.app_context():
             # get conversation and message
-            conversation = db.session.query(Conversation).where(Conversation.id == conversation_id).first()
+            stmt = select(Conversation).where(Conversation.id == conversation_id)
+            conversation = db.session.scalar(stmt)
 
             if not conversation:
                 return

api/core/callback_handler/index_tool_callback_handler.py

@@ -1,6 +1,8 @@
 import logging
 from collections.abc import Sequence
 
+from sqlalchemy import select
+
 from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.app.entities.queue_entities import QueueRetrieverResourcesEvent
@@ -49,7 +51,8 @@ class DatasetIndexToolCallbackHandler:
         for document in documents:
             if document.metadata is not None:
                 document_id = document.metadata["document_id"]
-                dataset_document = db.session.query(DatasetDocument).where(DatasetDocument.id == document_id).first()
+                dataset_document_stmt = select(DatasetDocument).where(DatasetDocument.id == document_id)
+                dataset_document = db.session.scalar(dataset_document_stmt)
                 if not dataset_document:
                     _logger.warning(
                         "Expected DatasetDocument record to exist, but none was found, document_id=%s",
@@ -57,15 +60,12 @@ class DatasetIndexToolCallbackHandler:
                     )
                     continue
                 if dataset_document.doc_form == IndexType.PARENT_CHILD_INDEX:
-                    child_chunk = (
-                        db.session.query(ChildChunk)
-                        .where(
-                            ChildChunk.index_node_id == document.metadata["doc_id"],
-                            ChildChunk.dataset_id == dataset_document.dataset_id,
-                            ChildChunk.document_id == dataset_document.id,
-                        )
-                        .first()
+                    child_chunk_stmt = select(ChildChunk).where(
+                        ChildChunk.index_node_id == document.metadata["doc_id"],
+                        ChildChunk.dataset_id == dataset_document.dataset_id,
+                        ChildChunk.document_id == dataset_document.id,
                     )
+                    child_chunk = db.session.scalar(child_chunk_stmt)
                     if child_chunk:
                         segment = (
                             db.session.query(DocumentSegment)

api/core/external_data_tool/api/api.py

@@ -1,5 +1,7 @@
 from typing import Optional
 
+from sqlalchemy import select
+
 from core.extension.api_based_extension_requestor import APIBasedExtensionRequestor
 from core.external_data_tool.base import ExternalDataTool
 from core.helper import encrypter
@@ -28,13 +30,11 @@ class ApiExternalDataTool(ExternalDataTool):
         api_based_extension_id = config.get("api_based_extension_id")
         if not api_based_extension_id:
             raise ValueError("api_based_extension_id is required")
-
         # get api_based_extension
-        api_based_extension = (
-            db.session.query(APIBasedExtension)
-            .where(APIBasedExtension.tenant_id == tenant_id, APIBasedExtension.id == api_based_extension_id)
-            .first()
+        stmt = select(APIBasedExtension).where(
+            APIBasedExtension.tenant_id == tenant_id, APIBasedExtension.id == api_based_extension_id
         )
+        api_based_extension = db.session.scalar(stmt)
 
         if not api_based_extension:
             raise ValueError("api_based_extension_id is invalid")
@@ -52,13 +52,11 @@ class ApiExternalDataTool(ExternalDataTool):
             raise ValueError(f"config is required, config: {self.config}")
         api_based_extension_id = self.config.get("api_based_extension_id")
         assert api_based_extension_id is not None, "api_based_extension_id is required"
-
         # get api_based_extension
-        api_based_extension = (
-            db.session.query(APIBasedExtension)
-            .where(APIBasedExtension.tenant_id == self.tenant_id, APIBasedExtension.id == api_based_extension_id)
-            .first()
+        stmt = select(APIBasedExtension).where(
+            APIBasedExtension.tenant_id == self.tenant_id, APIBasedExtension.id == api_based_extension_id
         )
+        api_based_extension = db.session.scalar(stmt)
 
         if not api_based_extension:
             raise ValueError(

api/core/indexing_runner.py

@@ -8,6 +8,7 @@ import uuid
 from typing import Any, Optional, cast
 
 from flask import current_app
+from sqlalchemy import select
 from sqlalchemy.orm.exc import ObjectDeletedError
 
 from configs import dify_config
@@ -56,13 +57,11 @@ class IndexingRunner:
 
                 if not dataset:
                     raise ValueError("no dataset found")
-
                 # get the process rule
-                processing_rule = (
-                    db.session.query(DatasetProcessRule)
-                    .where(DatasetProcessRule.id == dataset_document.dataset_process_rule_id)
-                    .first()
+                stmt = select(DatasetProcessRule).where(
+                    DatasetProcessRule.id == dataset_document.dataset_process_rule_id
                 )
+                processing_rule = db.session.scalar(stmt)
                 if not processing_rule:
                     raise ValueError("no process rule found")
                 index_type = dataset_document.doc_form
@@ -123,11 +122,8 @@ class IndexingRunner:
                     db.session.query(ChildChunk).where(ChildChunk.segment_id == document_segment.id).delete()
             db.session.commit()
             # get the process rule
-            processing_rule = (
-                db.session.query(DatasetProcessRule)
-                .where(DatasetProcessRule.id == dataset_document.dataset_process_rule_id)
-                .first()
-            )
+            stmt = select(DatasetProcessRule).where(DatasetProcessRule.id == dataset_document.dataset_process_rule_id)
+            processing_rule = db.session.scalar(stmt)
             if not processing_rule:
                 raise ValueError("no process rule found")
 
@@ -208,7 +204,6 @@ class IndexingRunner:
                                     child_documents.append(child_document)
                                 document.children = child_documents
                         documents.append(document)
-
             # build index
             index_type = dataset_document.doc_form
             index_processor = IndexProcessorFactory(index_type).init_index_processor()
@@ -310,7 +305,8 @@ class IndexingRunner:
                 # delete image files and related db records
                 image_upload_file_ids = get_image_upload_file_ids(document.page_content)
                 for upload_file_id in image_upload_file_ids:
-                    image_file = db.session.query(UploadFile).where(UploadFile.id == upload_file_id).first()
+                    stmt = select(UploadFile).where(UploadFile.id == upload_file_id)
+                    image_file = db.session.scalar(stmt)
                     if image_file is None:
                         continue
                     try:
@@ -339,10 +335,8 @@ class IndexingRunner:
         if dataset_document.data_source_type == "upload_file":
             if not data_source_info or "upload_file_id" not in data_source_info:
                 raise ValueError("no upload file found")
-
-            file_detail = (
-                db.session.query(UploadFile).where(UploadFile.id == data_source_info["upload_file_id"]).one_or_none()
-            )
+            stmt = select(UploadFile).where(UploadFile.id == data_source_info["upload_file_id"])
+            file_detail = db.session.scalars(stmt).one_or_none()
 
             if file_detail:
                 extract_setting = ExtractSetting(

api/core/memory/token_buffer_memory.py

@@ -110,9 +110,9 @@ class TokenBufferMemory:
         else:
             message_limit = 500
 
-        stmt = stmt.limit(message_limit)
+        msg_limit_stmt = stmt.limit(message_limit)
 
-        messages = db.session.scalars(stmt).all()
+        messages = db.session.scalars(msg_limit_stmt).all()
 
         # instead of all messages from the conversation, we only need to extract messages
         # that belong to the thread of last message

api/core/moderation/api/api.py

@@ -1,6 +1,7 @@
 from typing import Optional
 
 from pydantic import BaseModel, Field
+from sqlalchemy import select
 
 from core.extension.api_based_extension_requestor import APIBasedExtensionPoint, APIBasedExtensionRequestor
 from core.helper.encrypter import decrypt_token
@@ -87,10 +88,9 @@ class ApiModeration(Moderation):
 
     @staticmethod
     def _get_api_based_extension(tenant_id: str, api_based_extension_id: str) -> Optional[APIBasedExtension]:
-        extension = (
-            db.session.query(APIBasedExtension)
-            .where(APIBasedExtension.tenant_id == tenant_id, APIBasedExtension.id == api_based_extension_id)
-            .first()
+        stmt = select(APIBasedExtension).where(
+            APIBasedExtension.tenant_id == tenant_id, APIBasedExtension.id == api_based_extension_id
         )
+        extension = db.session.scalar(stmt)
 
         return extension

api/core/ops/aliyun_trace/aliyun_trace.py

@@ -5,6 +5,7 @@ from typing import Optional
 from urllib.parse import urljoin
 
 from opentelemetry.trace import Link, Status, StatusCode
+from sqlalchemy import select
 from sqlalchemy.orm import Session, sessionmaker
 
 from core.ops.aliyun_trace.data_exporter.traceclient import (
@@ -263,15 +264,15 @@ class AliyunDataTrace(BaseTraceInstance):
             app_id = trace_info.metadata.get("app_id")
             if not app_id:
                 raise ValueError("No app_id found in trace_info metadata")
-
-            app = session.query(App).where(App.id == app_id).first()
+            app_stmt = select(App).where(App.id == app_id)
+            app = session.scalar(app_stmt)
             if not app:
                 raise ValueError(f"App with id {app_id} not found")
 
             if not app.created_by:
                 raise ValueError(f"App with id {app_id} has no creator (created_by is None)")
-
-            service_account = session.query(Account).where(Account.id == app.created_by).first()
+            account_stmt = select(Account).where(Account.id == app.created_by)
+            service_account = session.scalar(account_stmt)
             if not service_account:
                 raise ValueError(f"Creator account with id {app.created_by} not found for app {app_id}")
             current_tenant = (

api/core/ops/base_trace_instance.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 
+from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from core.ops.entities.config_entity import BaseTracingConfig
@@ -44,14 +45,15 @@ class BaseTraceInstance(ABC):
         """
         with Session(db.engine, expire_on_commit=False) as session:
             # Get the app to find its creator
-            app = session.query(App).where(App.id == app_id).first()
+            app_stmt = select(App).where(App.id == app_id)
+            app = session.scalar(app_stmt)
             if not app:
                 raise ValueError(f"App with id {app_id} not found")
 
             if not app.created_by:
                 raise ValueError(f"App with id {app_id} has no creator (created_by is None)")
-
-            service_account = session.query(Account).where(Account.id == app.created_by).first()
+            account_stmt = select(Account).where(Account.id == app.created_by)
+            service_account = session.scalar(account_stmt)
             if not service_account:
                 raise ValueError(f"Creator account with id {app.created_by} not found for app {app_id}")
 

api/core/ops/ops_trace_manager.py

@@ -226,9 +226,9 @@ class OpsTraceManager:
 
         if not trace_config_data:
             return None
-
         # decrypt_token
-        app = db.session.query(App).where(App.id == app_id).first()
+        stmt = select(App).where(App.id == app_id)
+        app = db.session.scalar(stmt)
         if not app:
             raise ValueError("App not found")
 
@@ -295,20 +295,19 @@ class OpsTraceManager:
     @classmethod
     def get_app_config_through_message_id(cls, message_id: str):
         app_model_config = None
-        message_data = db.session.query(Message).where(Message.id == message_id).first()
+        message_stmt = select(Message).where(Message.id == message_id)
+        message_data = db.session.scalar(message_stmt)
         if not message_data:
             return None
         conversation_id = message_data.conversation_id
-        conversation_data = db.session.query(Conversation).where(Conversation.id == conversation_id).first()
+        conversation_stmt = select(Conversation).where(Conversation.id == conversation_id)
+        conversation_data = db.session.scalar(conversation_stmt)
         if not conversation_data:
             return None
 
         if conversation_data.app_model_config_id:
-            app_model_config = (
-                db.session.query(AppModelConfig)
-                .where(AppModelConfig.id == conversation_data.app_model_config_id)
-                .first()
-            )
+            config_stmt = select(AppModelConfig).where(AppModelConfig.id == conversation_data.app_model_config_id)
+            app_model_config = db.session.scalar(config_stmt)
         elif conversation_data.app_model_config_id is None and conversation_data.override_model_configs:
             app_model_config = conversation_data.override_model_configs
 

api/core/plugin/backwards_invocation/app.py

@@ -1,6 +1,8 @@
 from collections.abc import Generator, Mapping
 from typing import Optional, Union
 
+from sqlalchemy import select
+
 from controllers.service_api.wraps import create_or_update_end_user_for_user_id
 from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict
 from core.app.apps.advanced_chat.app_generator import AdvancedChatAppGenerator
@@ -192,10 +194,11 @@ class PluginAppBackwardsInvocation(BaseBackwardsInvocation):
         """
         get the user by user id
         """
-
-        user = db.session.query(EndUser).where(EndUser.id == user_id).first()
+        stmt = select(EndUser).where(EndUser.id == user_id)
+        user = db.session.scalar(stmt)
         if not user:
-            user = db.session.query(Account).where(Account.id == user_id).first()
+            stmt = select(Account).where(Account.id == user_id)
+            user = db.session.scalar(stmt)
 
         if not user:
             raise ValueError("user not found")

api/core/provider_manager.py

@@ -276,15 +276,11 @@ class ProviderManager:
         :param model_type: model type
         :return:
         """
-        # Get the corresponding TenantDefaultModel record
-        default_model = (
-            db.session.query(TenantDefaultModel)
-            .where(
-                TenantDefaultModel.tenant_id == tenant_id,
-                TenantDefaultModel.model_type == model_type.to_origin_model_type(),
-            )
-            .first()
+        stmt = select(TenantDefaultModel).where(
+            TenantDefaultModel.tenant_id == tenant_id,
+            TenantDefaultModel.model_type == model_type.to_origin_model_type(),
         )
+        default_model = db.session.scalar(stmt)
 
         # If it does not exist, get the first available provider model from get_configurations
         # and update the TenantDefaultModel record
@@ -367,16 +363,11 @@ class ProviderManager:
         model_names = [model.model for model in available_models]
         if model not in model_names:
             raise ValueError(f"Model {model} does not exist.")
-
-        # Get the list of available models from get_configurations and check if it is LLM
-        default_model = (
-            db.session.query(TenantDefaultModel)
-            .where(
-                TenantDefaultModel.tenant_id == tenant_id,
-                TenantDefaultModel.model_type == model_type.to_origin_model_type(),
-            )
-            .first()
+        stmt = select(TenantDefaultModel).where(
+            TenantDefaultModel.tenant_id == tenant_id,
+            TenantDefaultModel.model_type == model_type.to_origin_model_type(),
         )
+        default_model = db.session.scalar(stmt)
 
         # create or update TenantDefaultModel record
         if default_model:
@@ -598,16 +589,13 @@ class ProviderManager:
                             provider_name_to_provider_records_dict[provider_name].append(new_provider_record)
                         except IntegrityError:
                             db.session.rollback()
-                            existed_provider_record = (
-                                db.session.query(Provider)
-                                .where(
-                                    Provider.tenant_id == tenant_id,
-                                    Provider.provider_name == ModelProviderID(provider_name).provider_name,
-                                    Provider.provider_type == ProviderType.SYSTEM.value,
-                                    Provider.quota_type == ProviderQuotaType.TRIAL.value,
-                                )
-                                .first()
+                            stmt = select(Provider).where(
+                                Provider.tenant_id == tenant_id,
+                                Provider.provider_name == ModelProviderID(provider_name).provider_name,
+                                Provider.provider_type == ProviderType.SYSTEM.value,
+                                Provider.quota_type == ProviderQuotaType.TRIAL.value,
                             )
+                            existed_provider_record = db.session.scalar(stmt)
                             if not existed_provider_record:
                                 continue
 

api/core/rag/datasource/keyword/jieba/jieba.py

@@ -3,6 +3,7 @@ from typing import Any, Optional
 
 import orjson
 from pydantic import BaseModel
+from sqlalchemy import select
 
 from configs import dify_config
 from core.rag.datasource.keyword.jieba.jieba_keyword_table_handler import JiebaKeywordTableHandler
@@ -211,11 +212,10 @@ class Jieba(BaseKeyword):
         return sorted_chunk_indices[:k]
 
     def _update_segment_keywords(self, dataset_id: str, node_id: str, keywords: list[str]):
-        document_segment = (
-            db.session.query(DocumentSegment)
-            .where(DocumentSegment.dataset_id == dataset_id, DocumentSegment.index_node_id == node_id)
-            .first()
+        stmt = select(DocumentSegment).where(
+            DocumentSegment.dataset_id == dataset_id, DocumentSegment.index_node_id == node_id
         )
+        document_segment = db.session.scalar(stmt)
         if document_segment:
             document_segment.keywords = keywords
             db.session.add(document_segment)