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/tools/tool_label_manager.py

@@ -1,3 +1,5 @@
+from sqlalchemy import select
+
 from core.tools.__base.tool_provider import ToolProviderController
 from core.tools.builtin_tool.provider import BuiltinToolProviderController
 from core.tools.custom_tool.provider import ApiToolProviderController
@@ -54,17 +56,13 @@ class ToolLabelManager:
             return controller.tool_labels
         else:
             raise ValueError("Unsupported tool type")
-
-        labels = (
-            db.session.query(ToolLabelBinding.label_name)
-            .where(
-                ToolLabelBinding.tool_id == provider_id,
-                ToolLabelBinding.tool_type == controller.provider_type.value,
-            )
-            .all()
+        stmt = select(ToolLabelBinding.label_name).where(
+            ToolLabelBinding.tool_id == provider_id,
+            ToolLabelBinding.tool_type == controller.provider_type.value,
         )
+        labels = db.session.scalars(stmt).all()
 
-        return [label.label_name for label in labels]
+        return list(labels)
 
     @classmethod
     def get_tools_labels(cls, tool_providers: list[ToolProviderController]) -> dict[str, list[str]]:

api/core/tools/tool_manager.py

@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Literal, Optional, Union, cast
 
 import sqlalchemy as sa
 from pydantic import TypeAdapter
+from sqlalchemy import select
 from sqlalchemy.orm import Session
 from yarl import URL
 
@@ -198,14 +199,11 @@ class ToolManager:
                 # get specific credentials
                 if is_valid_uuid(credential_id):
                     try:
-                        builtin_provider = (
-                            db.session.query(BuiltinToolProvider)
-                            .where(
-                                BuiltinToolProvider.tenant_id == tenant_id,
-                                BuiltinToolProvider.id == credential_id,
-                            )
-                            .first()
+                        builtin_provider_stmt = select(BuiltinToolProvider).where(
+                            BuiltinToolProvider.tenant_id == tenant_id,
+                            BuiltinToolProvider.id == credential_id,
                         )
+                        builtin_provider = db.session.scalar(builtin_provider_stmt)
                     except Exception as e:
                         builtin_provider = None
                         logger.info("Error getting builtin provider %s:%s", credential_id, e, exc_info=True)
@@ -317,11 +315,10 @@ class ToolManager:
                 ),
             )
         elif provider_type == ToolProviderType.WORKFLOW:
-            workflow_provider = (
-                db.session.query(WorkflowToolProvider)
-                .where(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id)
-                .first()
+            workflow_provider_stmt = select(WorkflowToolProvider).where(
+                WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id
             )
+            workflow_provider = db.session.scalar(workflow_provider_stmt)
 
             if workflow_provider is None:
                 raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found")

api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py

@@ -3,6 +3,7 @@ from typing import Any
 
 from flask import Flask, current_app
 from pydantic import BaseModel, Field
+from sqlalchemy import select
 
 from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler
 from core.model_manager import ModelManager
@@ -85,17 +86,14 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool):
 
         document_context_list = []
         index_node_ids = [document.metadata["doc_id"] for document in all_documents if document.metadata]
-        segments = (
-            db.session.query(DocumentSegment)
-            .where(
-                DocumentSegment.dataset_id.in_(self.dataset_ids),
-                DocumentSegment.completed_at.isnot(None),
-                DocumentSegment.status == "completed",
-                DocumentSegment.enabled == True,
-                DocumentSegment.index_node_id.in_(index_node_ids),
-            )
-            .all()
+        document_segment_stmt = select(DocumentSegment).where(
+            DocumentSegment.dataset_id.in_(self.dataset_ids),
+            DocumentSegment.completed_at.isnot(None),
+            DocumentSegment.status == "completed",
+            DocumentSegment.enabled == True,
+            DocumentSegment.index_node_id.in_(index_node_ids),
         )
+        segments = db.session.scalars(document_segment_stmt).all()
 
         if segments:
             index_node_id_to_position = {id: position for position, id in enumerate(index_node_ids)}
@@ -112,15 +110,12 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool):
                 resource_number = 1
                 for segment in sorted_segments:
                     dataset = db.session.query(Dataset).filter_by(id=segment.dataset_id).first()
-                    document = (
-                        db.session.query(Document)
-                        .where(
-                            Document.id == segment.document_id,
-                            Document.enabled == True,
-                            Document.archived == False,
-                        )
-                        .first()
+                    document_stmt = select(Document).where(
+                        Document.id == segment.document_id,
+                        Document.enabled == True,
+                        Document.archived == False,
                     )
+                    document = db.session.scalar(document_stmt)
                     if dataset and document:
                         source = RetrievalSourceMetadata(
                             position=resource_number,
@@ -162,9 +157,8 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool):
         hit_callbacks: list[DatasetIndexToolCallbackHandler],
     ):
         with flask_app.app_context():
-            dataset = (
-                db.session.query(Dataset).where(Dataset.tenant_id == self.tenant_id, Dataset.id == dataset_id).first()
-            )
+            stmt = select(Dataset).where(Dataset.tenant_id == self.tenant_id, Dataset.id == dataset_id)
+            dataset = db.session.scalar(stmt)
 
             if not dataset:
                 return []

api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py

@@ -1,6 +1,7 @@
 from typing import Any, Optional, cast
 
 from pydantic import BaseModel, Field
+from sqlalchemy import select
 
 from core.app.app_config.entities import DatasetRetrieveConfigEntity, ModelConfig
 from core.rag.datasource.retrieval_service import RetrievalService
@@ -56,9 +57,8 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool):
         )
 
     def _run(self, query: str) -> str:
-        dataset = (
-            db.session.query(Dataset).where(Dataset.tenant_id == self.tenant_id, Dataset.id == self.dataset_id).first()
-        )
+        dataset_stmt = select(Dataset).where(Dataset.tenant_id == self.tenant_id, Dataset.id == self.dataset_id)
+        dataset = db.session.scalar(dataset_stmt)
 
         if not dataset:
             return ""
@@ -188,15 +188,12 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool):
                         for record in records:
                             segment = record.segment
                             dataset = db.session.query(Dataset).filter_by(id=segment.dataset_id).first()
-                            document = (
-                                db.session.query(DatasetDocument)  # type: ignore
-                                .where(
-                                    DatasetDocument.id == segment.document_id,
-                                    DatasetDocument.enabled == True,
-                                    DatasetDocument.archived == False,
-                                )
-                                .first()
+                            dataset_document_stmt = select(DatasetDocument).where(
+                                DatasetDocument.id == segment.document_id,
+                                DatasetDocument.enabled == True,
+                                DatasetDocument.archived == False,
                             )
+                            document = db.session.scalar(dataset_document_stmt)  # type: ignore
                             if dataset and document:
                                 source = RetrievalSourceMetadata(
                                     dataset_id=dataset.id,

api/core/tools/workflow_as_tool/tool.py

@@ -3,6 +3,8 @@ import logging
 from collections.abc import Generator
 from typing import Any, Optional
 
+from sqlalchemy import select
+
 from core.file import FILE_MODEL_IDENTITY, File, FileTransferMethod
 from core.tools.__base.tool import Tool
 from core.tools.__base.tool_runtime import ToolRuntime
@@ -136,7 +138,8 @@ class WorkflowTool(Tool):
                 .first()
             )
         else:
-            workflow = db.session.query(Workflow).where(Workflow.app_id == app_id, Workflow.version == version).first()
+            stmt = select(Workflow).where(Workflow.app_id == app_id, Workflow.version == version)
+            workflow = db.session.scalar(stmt)
 
         if not workflow:
             raise ValueError("workflow not found or not published")
@@ -147,7 +150,8 @@ class WorkflowTool(Tool):
         """
         get the app by app id
         """
-        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")