From 68bd29eda24d83f1fd4ce95727667a89e62565af Mon Sep 17 00:00:00 2001
From: Renzo <170978465+RenzoMXD@users.noreply.github.com>
Date: Mon, 6 Apr 2026 23:23:14 -0500
Subject: [PATCH] refactor: migrate session.query to select API in sync task
 and services (#34619)

---
 api/services/attachment_service.py               |  6 +++---
 api/services/plugin/plugin_parameter_service.py  | 13 +++++++------
 api/tasks/sync_website_document_indexing_task.py | 10 ++++++----
 3 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/api/services/attachment_service.py b/api/services/attachment_service.py
index 2bd5627d5eb..54e664e9443 100644
--- a/api/services/attachment_service.py
+++ b/api/services/attachment_service.py
@@ -1,6 +1,6 @@
 import base64
 
-from sqlalchemy import Engine
+from sqlalchemy import Engine, select
 from sqlalchemy.orm import sessionmaker
 from werkzeug.exceptions import NotFound
 
@@ -22,8 +22,8 @@ class AttachmentService:
             raise AssertionError("must be a sessionmaker or an Engine.")
 
     def get_file_base64(self, file_id: str) -> str:
-        upload_file = (
-            self._session_maker(expire_on_commit=False).query(UploadFile).where(UploadFile.id == file_id).first()
+        upload_file = self._session_maker(expire_on_commit=False).scalar(
+            select(UploadFile).where(UploadFile.id == file_id).limit(1)
         )
         if not upload_file:
             raise NotFound("File not found")
diff --git a/api/services/plugin/plugin_parameter_service.py b/api/services/plugin/plugin_parameter_service.py
index 40565c56ed9..786c09b44e1 100644
--- a/api/services/plugin/plugin_parameter_service.py
+++ b/api/services/plugin/plugin_parameter_service.py
@@ -1,6 +1,7 @@
 from collections.abc import Mapping, Sequence
 from typing import Any, Literal
 
+from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from core.plugin.entities.parameters import PluginParameterOption
@@ -56,24 +57,24 @@ class PluginParameterService:
                     # fetch credentials from db
                     with Session(db.engine) as session:
                         if credential_id:
-                            db_record = (
-                                session.query(BuiltinToolProvider)
+                            db_record = session.scalar(
+                                select(BuiltinToolProvider)
                                 .where(
                                     BuiltinToolProvider.tenant_id == tenant_id,
                                     BuiltinToolProvider.provider == provider,
                                     BuiltinToolProvider.id == credential_id,
                                 )
-                                .first()
+                                .limit(1)
                             )
                         else:
-                            db_record = (
-                                session.query(BuiltinToolProvider)
+                            db_record = session.scalar(
+                                select(BuiltinToolProvider)
                                 .where(
                                     BuiltinToolProvider.tenant_id == tenant_id,
                                     BuiltinToolProvider.provider == provider,
                                 )
                                 .order_by(BuiltinToolProvider.is_default.desc(), BuiltinToolProvider.created_at.asc())
-                                .first()
+                                .limit(1)
                             )
 
                     if db_record is None:
diff --git a/api/tasks/sync_website_document_indexing_task.py b/api/tasks/sync_website_document_indexing_task.py
index aa6bce958ba..ab21f63f7ea 100644
--- a/api/tasks/sync_website_document_indexing_task.py
+++ b/api/tasks/sync_website_document_indexing_task.py
@@ -29,7 +29,7 @@ def sync_website_document_indexing_task(dataset_id: str, document_id: str):
     start_at = time.perf_counter()
 
     with session_factory.create_session() as session:
-        dataset = session.query(Dataset).where(Dataset.id == dataset_id).first()
+        dataset = session.scalar(select(Dataset).where(Dataset.id == dataset_id).limit(1))
         if dataset is None:
             raise ValueError("Dataset not found")
 
@@ -45,8 +45,8 @@ def sync_website_document_indexing_task(dataset_id: str, document_id: str):
                         "your subscription."
                     )
         except Exception as e:
-            document = (
-                session.query(Document).where(Document.id == document_id, Document.dataset_id == dataset_id).first()
+            document = session.scalar(
+                select(Document).where(Document.id == document_id, Document.dataset_id == dataset_id).limit(1)
             )
             if document:
                 document.indexing_status = IndexingStatus.ERROR
@@ -58,7 +58,9 @@ def sync_website_document_indexing_task(dataset_id: str, document_id: str):
             return
 
         logger.info(click.style(f"Start sync website document: {document_id}", fg="green"))
-        document = session.query(Document).where(Document.id == document_id, Document.dataset_id == dataset_id).first()
+        document = session.scalar(
+            select(Document).where(Document.id == document_id, Document.dataset_id == dataset_id).limit(1)
+        )
         if not document:
             logger.info(click.style(f"Document not found: {document_id}", fg="yellow"))
             return