update sql in batch (langgenius#24801)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: -LAN- <laipz8200@outlook.com>

Author: 3 people

api/commands.py

@@ -212,7 +212,9 @@ def migrate_annotation_vector_database():
                 if not dataset_collection_binding:
                     click.echo(f"App annotation collection binding not found: {app.id}")
                     continue
-                annotations = db.session.query(MessageAnnotation).where(MessageAnnotation.app_id == app.id).all()
+                annotations = db.session.scalars(
+                    select(MessageAnnotation).where(MessageAnnotation.app_id == app.id)
+                ).all()
                 dataset = Dataset(
                     id=app.id,
                     tenant_id=app.tenant_id,
@@ -367,29 +369,25 @@ def migrate_knowledge_vector_database():
                     )
                     raise e
 
-                dataset_documents = (
-                    db.session.query(DatasetDocument)
-                    .where(
+                dataset_documents = db.session.scalars(
+                    select(DatasetDocument).where(
                         DatasetDocument.dataset_id == dataset.id,
                         DatasetDocument.indexing_status == "completed",
                         DatasetDocument.enabled == True,
                         DatasetDocument.archived == False,
                     )
-                    .all()
-                )
+                ).all()
 
                 documents = []
                 segments_count = 0
                 for dataset_document in dataset_documents:
-                    segments = (
-                        db.session.query(DocumentSegment)
-                        .where(
+                    segments = db.session.scalars(
+                        select(DocumentSegment).where(
                             DocumentSegment.document_id == dataset_document.id,
                             DocumentSegment.status == "completed",
                             DocumentSegment.enabled == True,
                         )
-                        .all()
-                    )
+                    ).all()
 
                     for segment in segments:
                         document = Document(

api/controllers/console/apikey.py

@@ -60,11 +60,11 @@ def get(self, resource_id):
         assert self.resource_id_field is not None, "resource_id_field must be set"
         resource_id = str(resource_id)
         _get_resource(resource_id, current_user.current_tenant_id, self.resource_model)
-        keys = (
-            db.session.query(ApiToken)
-            .where(ApiToken.type == self.resource_type, getattr(ApiToken, self.resource_id_field) == resource_id)
-            .all()
-        )
+        keys = db.session.scalars(
+            select(ApiToken).where(
+                ApiToken.type == self.resource_type, getattr(ApiToken, self.resource_id_field) == resource_id
+            )
+        ).all()
         return {"items": keys}
 
     @marshal_with(api_key_fields)

api/controllers/console/datasets/data_source.py

@@ -29,14 +29,12 @@ class DataSourceApi(Resource):
     @marshal_with(integrate_list_fields)
     def get(self):
         # get workspace data source integrates
-        data_source_integrates = (
-            db.session.query(DataSourceOauthBinding)
-            .where(
+        data_source_integrates = db.session.scalars(
+            select(DataSourceOauthBinding).where(
                 DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
                 DataSourceOauthBinding.disabled == False,
             )
-            .all()
-        )
+        ).all()
 
         base_url = request.url_root.rstrip("/")
         data_source_oauth_base_path = "/console/api/oauth/data-source"

api/controllers/console/datasets/datasets.py

@@ -2,6 +2,7 @@
 from flask import request
 from flask_login import current_user
 from flask_restx import Resource, marshal, marshal_with, reqparse
+from sqlalchemy import select
 from werkzeug.exceptions import Forbidden, NotFound
 
 import services
@@ -411,11 +412,11 @@ def post(self):
         extract_settings = []
         if args["info_list"]["data_source_type"] == "upload_file":
             file_ids = args["info_list"]["file_info_list"]["file_ids"]
-            file_details = (
-                db.session.query(UploadFile)
-                .where(UploadFile.tenant_id == current_user.current_tenant_id, UploadFile.id.in_(file_ids))
-                .all()
-            )
+            file_details = db.session.scalars(
+                select(UploadFile).where(
+                    UploadFile.tenant_id == current_user.current_tenant_id, UploadFile.id.in_(file_ids)
+                )
+            ).all()
 
             if file_details is None:
                 raise NotFound("File not found.")
@@ -518,11 +519,11 @@ class DatasetIndexingStatusApi(Resource):
     @account_initialization_required
     def get(self, dataset_id):
         dataset_id = str(dataset_id)
-        documents = (
-            db.session.query(Document)
-            .where(Document.dataset_id == dataset_id, Document.tenant_id == current_user.current_tenant_id)
-            .all()
-        )
+        documents = db.session.scalars(
+            select(Document).where(
+                Document.dataset_id == dataset_id, Document.tenant_id == current_user.current_tenant_id
+            )
+        ).all()
         documents_status = []
         for document in documents:
             completed_segments = (
@@ -569,11 +570,11 @@ class DatasetApiKeyApi(Resource):
     @account_initialization_required
     @marshal_with(api_key_list)
     def get(self):
-        keys = (
-            db.session.query(ApiToken)
-            .where(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id)
-            .all()
-        )
+        keys = db.session.scalars(
+            select(ApiToken).where(
+                ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id
+            )
+        ).all()
         return {"items": keys}
 
     @setup_required

api/controllers/console/datasets/datasets_document.py

@@ -1,5 +1,6 @@
 import logging
 from argparse import ArgumentTypeError
+from collections.abc import Sequence
 from typing import Literal, cast
 
 from flask import request
@@ -79,7 +80,7 @@ def get_document(self, dataset_id: str, document_id: str) -> Document:
 
         return document
 
-    def get_batch_documents(self, dataset_id: str, batch: str) -> list[Document]:
+    def get_batch_documents(self, dataset_id: str, batch: str) -> Sequence[Document]:
         dataset = DatasetService.get_dataset(dataset_id)
         if not dataset:
             raise NotFound("Dataset not found.")

api/controllers/console/explore/installed_app.py

@@ -3,7 +3,7 @@
 
 from flask import request
 from flask_restx import Resource, inputs, marshal_with, reqparse
-from sqlalchemy import and_
+from sqlalchemy import and_, select
 from werkzeug.exceptions import BadRequest, Forbidden, NotFound
 
 from controllers.console import api
@@ -33,13 +33,15 @@ def get(self):
         current_tenant_id = current_user.current_tenant_id
 
         if app_id:
-            installed_apps = (
-                db.session.query(InstalledApp)
-                .where(and_(InstalledApp.tenant_id == current_tenant_id, InstalledApp.app_id == app_id))
-                .all()
-            )
+            installed_apps = db.session.scalars(
+                select(InstalledApp).where(
+                    and_(InstalledApp.tenant_id == current_tenant_id, InstalledApp.app_id == app_id)
+                )
+            ).all()
         else:
-            installed_apps = db.session.query(InstalledApp).where(InstalledApp.tenant_id == current_tenant_id).all()
+            installed_apps = db.session.scalars(
+                select(InstalledApp).where(InstalledApp.tenant_id == current_tenant_id)
+            ).all()
 
         if current_user.current_tenant is None:
             raise ValueError("current_user.current_tenant must not be None")

api/controllers/console/workspace/account.py

@@ -248,7 +248,9 @@ def get(self):
             raise ValueError("Invalid user account")
         account = current_user
 
-        account_integrates = db.session.query(AccountIntegrate).where(AccountIntegrate.account_id == account.id).all()
+        account_integrates = db.session.scalars(
+            select(AccountIntegrate).where(AccountIntegrate.account_id == account.id)
+        ).all()
 
         base_url = request.url_root.rstrip("/")
         oauth_base_path = "/console/api/oauth/login"

api/core/memory/token_buffer_memory.py

@@ -32,11 +32,16 @@ def __init__(
         self.model_instance = model_instance
 
     def _build_prompt_message_with_files(
-        self, message_files: list[MessageFile], text_content: str, message: Message, app_record, is_user_message: bool
+        self,
+        message_files: Sequence[MessageFile],
+        text_content: str,
+        message: Message,
+        app_record,
+        is_user_message: bool,
     ) -> PromptMessage:
         """
         Build prompt message with files.
-        :param message_files: list of MessageFile objects
+        :param message_files: Sequence of MessageFile objects
         :param text_content: text content of the message
         :param message: Message object
         :param app_record: app record
@@ -128,14 +133,12 @@ def get_history_prompt_messages(
         prompt_messages: list[PromptMessage] = []
         for message in messages:
             # Process user message with files
-            user_files = (
-                db.session.query(MessageFile)
-                .where(
+            user_files = db.session.scalars(
+                select(MessageFile).where(
                     MessageFile.message_id == message.id,
                     (MessageFile.belongs_to == "user") | (MessageFile.belongs_to.is_(None)),
                 )
-                .all()
-            )
+            ).all()
 
             if user_files:
                 user_prompt_message = self._build_prompt_message_with_files(
@@ -150,11 +153,9 @@ def get_history_prompt_messages(
                 prompt_messages.append(UserPromptMessage(content=message.query))
 
             # Process assistant message with files
-            assistant_files = (
-                db.session.query(MessageFile)
-                .where(MessageFile.message_id == message.id, MessageFile.belongs_to == "assistant")
-                .all()
-            )
+            assistant_files = db.session.scalars(
+                select(MessageFile).where(MessageFile.message_id == message.id, MessageFile.belongs_to == "assistant")
+            ).all()
 
             if assistant_files:
                 assistant_prompt_message = self._build_prompt_message_with_files(

api/core/ops/arize_phoenix_trace/arize_phoenix_trace.py

@@ -15,6 +15,7 @@
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.sdk.trace.id_generator import RandomIdGenerator
 from opentelemetry.trace import SpanContext, TraceFlags, TraceState
+from sqlalchemy import select
 
 from core.ops.base_trace_instance import BaseTraceInstance
 from core.ops.entities.config_entity import ArizeConfig, PhoenixConfig
@@ -699,8 +700,8 @@ def get_project_url(self):
 
     def _get_workflow_nodes(self, workflow_run_id: str):
         """Helper method to get workflow nodes"""
-        workflow_nodes = (
-            db.session.query(
+        workflow_nodes = db.session.scalars(
+            select(
                 WorkflowNodeExecutionModel.id,
                 WorkflowNodeExecutionModel.tenant_id,
                 WorkflowNodeExecutionModel.app_id,
@@ -713,10 +714,8 @@ def _get_workflow_nodes(self, workflow_run_id: str):
                 WorkflowNodeExecutionModel.elapsed_time,
                 WorkflowNodeExecutionModel.process_data,
                 WorkflowNodeExecutionModel.execution_metadata,
-            )
-            .where(WorkflowNodeExecutionModel.workflow_run_id == workflow_run_id)
-            .all()
-        )
+            ).where(WorkflowNodeExecutionModel.workflow_run_id == workflow_run_id)
+        ).all()
         return workflow_nodes
 
     def _construct_llm_attributes(self, prompts: dict | list | str | None) -> dict[str, str]:

api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_service.py

@@ -1,5 +1,6 @@
 import time
 import uuid
+from collections.abc import Sequence
 
 import requests
 from requests.auth import HTTPDigestAuth
@@ -139,7 +140,7 @@ def change_tidb_serverless_root_password(
 
     @staticmethod
     def batch_update_tidb_serverless_cluster_status(
-        tidb_serverless_list: list[TidbAuthBinding],
+        tidb_serverless_list: Sequence[TidbAuthBinding],
         project_id: str,
         api_url: str,
         iam_url: str,