ast-grep --pattern 'session.query($NAME).where($COND).first()' -r 'session.scalars(select($NAME).where($COND).limit(1)).first()' -l py --update-all

Author: asukaminato0721

api/controllers/console/workspace/__init__.py

@@ -3,6 +3,7 @@
 from typing import ParamSpec, TypeVar
 
 from flask_login import current_user
+from sqlalchemy import select
 from sqlalchemy.orm import Session
 from werkzeug.exceptions import Forbidden
 
@@ -24,13 +25,9 @@ def decorated(*args: P.args, **kwargs: P.kwargs):
             tenant_id = user.current_tenant_id
 
             with Session(db.engine) as session:
-                permission = (
-                    session.query(TenantPluginPermission)
-                    .where(
-                        TenantPluginPermission.tenant_id == tenant_id,
-                    )
-                    .first()
-                )
+                permission = session.scalars(
+                    select(TenantPluginPermission).where(TenantPluginPermission.tenant_id == tenant_id).limit(1)
+                ).first()
 
                 if not permission:
                     # no permission set, allow access for everyone

api/controllers/mcp/mcp.py

@@ -3,6 +3,7 @@
 from flask import Response
 from flask_restx import Resource, reqparse
 from pydantic import ValidationError
+from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from controllers.console.app.mcp_server import AppMCPServerStatus
@@ -89,11 +90,13 @@ def post(self, server_code: str):
 
     def _get_mcp_server_and_app(self, server_code: str, session: Session) -> tuple[AppMCPServer, App]:
         """Get and validate MCP server and app in one query session"""
-        mcp_server = session.query(AppMCPServer).where(AppMCPServer.server_code == server_code).first()
+        mcp_server = session.scalars(
+            select(AppMCPServer).where(AppMCPServer.server_code == server_code).limit(1)
+        ).first()
         if not mcp_server:
             raise MCPRequestError(mcp_types.INVALID_REQUEST, "Server Not Found")
 
-        app = session.query(App).where(App.id == mcp_server.app_id).first()
+        app = session.scalars(select(App).where(App.id == mcp_server.app_id).limit(1)).first()
         if not app:
             raise MCPRequestError(mcp_types.INVALID_REQUEST, "App Not Found")
 

api/core/app/task_pipeline/easy_ui_based_generate_task_pipeline.py

@@ -473,9 +473,9 @@ def _agent_thought_to_stream_response(self, event: QueueAgentThoughtEvent) -> Op
         :return:
         """
         with Session(db.engine, expire_on_commit=False) as session:
-            agent_thought: Optional[MessageAgentThought] = (
-                session.query(MessageAgentThought).where(MessageAgentThought.id == event.agent_thought_id).first()
-            )
+            agent_thought: Optional[MessageAgentThought] = session.scalars(
+                select(MessageAgentThought).where(MessageAgentThought.id == event.agent_thought_id).limit(1)
+            ).first()
 
         if agent_thought:
             return AgentThoughtStreamResponse(

api/core/rag/datasource/retrieval_service.py

@@ -147,7 +147,7 @@ def external_retrieve(
     @classmethod
     def _get_dataset(cls, dataset_id: str) -> Optional[Dataset]:
         with Session(db.engine) as session:
-            return session.query(Dataset).where(Dataset.id == dataset_id).first()
+            return session.scalars(select(Dataset).where(Dataset.id == dataset_id).limit(1)).first()
 
     @classmethod
     def keyword_search(

api/core/tools/tool_file_manager.py

@@ -10,6 +10,7 @@
 from uuid import uuid4
 
 import httpx
+from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from configs import dify_config
@@ -158,13 +159,7 @@ def get_file_binary(self, id: str) -> Union[tuple[bytes, str], None]:
         :return: the binary of the file, mime type
         """
         with Session(self._engine, expire_on_commit=False) as session:
-            tool_file: ToolFile | None = (
-                session.query(ToolFile)
-                .where(
-                    ToolFile.id == id,
-                )
-                .first()
-            )
+            tool_file: ToolFile | None = session.scalars(select(ToolFile).where(ToolFile.id == id).limit(1)).first()
 
         if not tool_file:
             return None
@@ -182,13 +177,9 @@ def get_file_binary_by_message_file_id(self, id: str) -> Union[tuple[bytes, str]
         :return: the binary of the file, mime type
         """
         with Session(self._engine, expire_on_commit=False) as session:
-            message_file: MessageFile | None = (
-                session.query(MessageFile)
-                .where(
-                    MessageFile.id == id,
-                )
-                .first()
-            )
+            message_file: MessageFile | None = session.scalars(
+                select(MessageFile).where(MessageFile.id == id).limit(1)
+            ).first()
 
             # Check if message_file is not None
             if message_file is not None:
@@ -202,13 +193,9 @@ def get_file_binary_by_message_file_id(self, id: str) -> Union[tuple[bytes, str]
             else:
                 tool_file_id = None
 
-            tool_file: ToolFile | None = (
-                session.query(ToolFile)
-                .where(
-                    ToolFile.id == tool_file_id,
-                )
-                .first()
-            )
+            tool_file: ToolFile | None = session.scalars(
+                select(ToolFile).where(ToolFile.id == tool_file_id).limit(1)
+            ).first()
 
         if not tool_file:
             return None
@@ -226,13 +213,9 @@ def get_file_generator_by_tool_file_id(self, tool_file_id: str) -> tuple[Optiona
         :return: the binary of the file, mime type
         """
         with Session(self._engine, expire_on_commit=False) as session:
-            tool_file: ToolFile | None = (
-                session.query(ToolFile)
-                .where(
-                    ToolFile.id == tool_file_id,
-                )
-                .first()
-            )
+            tool_file: ToolFile | None = session.scalars(
+                select(ToolFile).where(ToolFile.id == tool_file_id).limit(1)
+            ).first()
 
         if not tool_file:
             return None, None

api/services/plugin/plugin_auto_upgrade_service.py

@@ -1,3 +1,4 @@
+from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from extensions.ext_database import db
@@ -8,11 +9,11 @@ class PluginAutoUpgradeService:
     @staticmethod
     def get_strategy(tenant_id: str) -> TenantPluginAutoUpgradeStrategy | None:
         with Session(db.engine) as session:
-            return (
-                session.query(TenantPluginAutoUpgradeStrategy)
+            return session.scalars(
+                select(TenantPluginAutoUpgradeStrategy)
                 .where(TenantPluginAutoUpgradeStrategy.tenant_id == tenant_id)
-                .first()
-            )
+                .limit(1)
+            ).first()
 
     @staticmethod
     def change_strategy(
@@ -24,11 +25,11 @@ def change_strategy(
         include_plugins: list[str],
     ) -> bool:
         with Session(db.engine) as session:
-            exist_strategy = (
-                session.query(TenantPluginAutoUpgradeStrategy)
+            exist_strategy = session.scalars(
+                select(TenantPluginAutoUpgradeStrategy)
                 .where(TenantPluginAutoUpgradeStrategy.tenant_id == tenant_id)
-                .first()
-            )
+                .limit(1)
+            ).first()
             if not exist_strategy:
                 strategy = TenantPluginAutoUpgradeStrategy(
                     tenant_id=tenant_id,
@@ -52,11 +53,11 @@ def change_strategy(
     @staticmethod
     def exclude_plugin(tenant_id: str, plugin_id: str) -> bool:
         with Session(db.engine) as session:
-            exist_strategy = (
-                session.query(TenantPluginAutoUpgradeStrategy)
+            exist_strategy = session.scalars(
+                select(TenantPluginAutoUpgradeStrategy)
                 .where(TenantPluginAutoUpgradeStrategy.tenant_id == tenant_id)
-                .first()
-            )
+                .limit(1)
+            ).first()
             if not exist_strategy:
                 # create for this tenant
                 PluginAutoUpgradeService.change_strategy(

api/services/plugin/plugin_permission_service.py

@@ -1,3 +1,4 @@
+from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from extensions.ext_database import db
@@ -8,7 +9,9 @@ class PluginPermissionService:
     @staticmethod
     def get_permission(tenant_id: str) -> TenantPluginPermission | None:
         with Session(db.engine) as session:
-            return session.query(TenantPluginPermission).where(TenantPluginPermission.tenant_id == tenant_id).first()
+            return session.scalars(
+                select(TenantPluginPermission).where(TenantPluginPermission.tenant_id == tenant_id).limit(1)
+            ).first()
 
     @staticmethod
     def change_permission(
@@ -17,9 +20,9 @@ def change_permission(
         debug_permission: TenantPluginPermission.DebugPermission,
     ):
         with Session(db.engine) as session:
-            permission = (
-                session.query(TenantPluginPermission).where(TenantPluginPermission.tenant_id == tenant_id).first()
-            )
+            permission = session.scalars(
+                select(TenantPluginPermission).where(TenantPluginPermission.tenant_id == tenant_id).limit(1)
+            ).first()
             if not permission:
                 permission = TenantPluginPermission(
                     tenant_id=tenant_id, install_permission=install_permission, debug_permission=debug_permission

api/tests/unit_tests/models/test_types_enum_text.py

@@ -5,7 +5,7 @@
 import pytest
 import sqlalchemy as sa
 from sqlalchemy import exc as sa_exc
-from sqlalchemy import insert
+from sqlalchemy import insert, select
 from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column
 from sqlalchemy.sql.sqltypes import VARCHAR
 
@@ -114,12 +114,14 @@ def test_insert_and_select(self):
             session.commit()
 
         with Session(engine) as session:
-            user = session.query(_User).where(_User.id == admin_user_id).first()
+            user = session.scalars(select(_User).where(_User.id == admin_user_id).limit(1)).first()
+            assert user is not None
             assert user.user_type == _UserType.admin
             assert user.user_type_nullable is None
 
         with Session(engine) as session:
-            user = session.query(_User).where(_User.id == normal_user_id).first()
+            user = session.scalars(select(_User).where(_User.id == normal_user_id).limit(1)).first()
+            assert user is not None
             assert user.user_type == _UserType.normal
             assert user.user_type_nullable == _UserType.normal
 
@@ -188,4 +190,4 @@ def test_select_invalid_values(self):
 
         with pytest.raises(ValueError) as exc:
             with Session(engine) as session:
-                _user = session.query(_User).where(_User.id == 1).first()
+                _user = session.scalars(select(_User).where(_User.id == 1).limit(1)).first()