diff --git a/api/controllers/console/app/workflow_trigger.py b/api/controllers/console/app/workflow_trigger.py index aa37d247383..e4a6afae1ee 100644 --- a/api/controllers/console/app/workflow_trigger.py +++ b/api/controllers/console/app/workflow_trigger.py @@ -66,13 +66,13 @@ class WebhookTriggerApi(Resource): with sessionmaker(db.engine).begin() as session: # Get webhook trigger for this app and node - webhook_trigger = ( - session.query(WorkflowWebhookTrigger) + webhook_trigger = session.scalar( + select(WorkflowWebhookTrigger) .where( WorkflowWebhookTrigger.app_id == app_model.id, WorkflowWebhookTrigger.node_id == node_id, ) - .first() + .limit(1) ) if not webhook_trigger: diff --git a/api/controllers/console/datasets/rag_pipeline/rag_pipeline.py b/api/controllers/console/datasets/rag_pipeline/rag_pipeline.py index 1758bad31da..4fe96902575 100644 --- a/api/controllers/console/datasets/rag_pipeline/rag_pipeline.py +++ b/api/controllers/console/datasets/rag_pipeline/rag_pipeline.py @@ -3,6 +3,7 @@ import logging from flask import request from flask_restx import Resource from pydantic import BaseModel, Field +from sqlalchemy import select from sqlalchemy.orm import sessionmaker from controllers.common.schema import register_schema_models @@ -86,8 +87,8 @@ class CustomizedPipelineTemplateApi(Resource): @enterprise_license_required def post(self, template_id: str): with sessionmaker(db.engine, expire_on_commit=False).begin() as session: - template = ( - session.query(PipelineCustomizedTemplate).where(PipelineCustomizedTemplate.id == template_id).first() + template = session.scalar( + select(PipelineCustomizedTemplate).where(PipelineCustomizedTemplate.id == template_id).limit(1) ) if not template: raise ValueError("Customized pipeline template not found.") diff --git a/api/controllers/console/workspace/__init__.py b/api/controllers/console/workspace/__init__.py index 971674cee2b..60f712e4768 100644 --- a/api/controllers/console/workspace/__init__.py +++ b/api/controllers/console/workspace/__init__.py @@ -1,6 +1,7 @@ from collections.abc import Callable from functools import wraps +from sqlalchemy import select from sqlalchemy.orm import sessionmaker from werkzeug.exceptions import Forbidden @@ -21,12 +22,12 @@ def plugin_permission_required( tenant_id = current_tenant_id with sessionmaker(db.engine).begin() as session: - permission = ( - session.query(TenantPluginPermission) + permission = session.scalar( + select(TenantPluginPermission) .where( TenantPluginPermission.tenant_id == tenant_id, ) - .first() + .limit(1) ) if not permission: diff --git a/api/controllers/mcp/mcp.py b/api/controllers/mcp/mcp.py index 3c59535a48f..d2ce0ea543d 100644 --- a/api/controllers/mcp/mcp.py +++ b/api/controllers/mcp/mcp.py @@ -4,6 +4,7 @@ from flask import Response from flask_restx import Resource from graphon.variables.input_entities import VariableEntity from pydantic import BaseModel, Field, ValidationError +from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker from controllers.common.schema import register_schema_model @@ -80,11 +81,11 @@ class MCPAppApi(Resource): 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.scalar(select(AppMCPServer).where(AppMCPServer.server_code == server_code).limit(1)) 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.scalar(select(App).where(App.id == mcp_server.app_id).limit(1)) if not app: raise MCPRequestError(mcp_types.INVALID_REQUEST, "App Not Found") @@ -190,12 +191,12 @@ class MCPAppApi(Resource): def _retrieve_end_user(self, tenant_id: str, mcp_server_id: str) -> EndUser | None: """Get end user - manages its own database session""" with sessionmaker(db.engine, expire_on_commit=False).begin() as session: - return ( - session.query(EndUser) + return session.scalar( + select(EndUser) .where(EndUser.tenant_id == tenant_id) .where(EndUser.session_id == mcp_server_id) .where(EndUser.type == "mcp") - .first() + .limit(1) ) def _create_end_user( diff --git a/api/tests/test_containers_integration_tests/controllers/console/app/test_app_apis.py b/api/tests/test_containers_integration_tests/controllers/console/app/test_app_apis.py index 0841217fcfe..c3a861c3e1c 100644 --- a/api/tests/test_containers_integration_tests/controllers/console/app/test_app_apis.py +++ b/api/tests/test_containers_integration_tests/controllers/console/app/test_app_apis.py @@ -555,7 +555,7 @@ class TestWorkflowTriggerEndpoints: trigger = MagicMock() session = MagicMock() - session.query.return_value.where.return_value.first.return_value = trigger + session.scalar.return_value = trigger class DummySessionCtx: def __enter__(self): diff --git a/api/tests/test_containers_integration_tests/controllers/mcp/test_mcp.py b/api/tests/test_containers_integration_tests/controllers/mcp/test_mcp.py index 90670a9db59..21b395a04c6 100644 --- a/api/tests/test_containers_integration_tests/controllers/mcp/test_mcp.py +++ b/api/tests/test_containers_integration_tests/controllers/mcp/test_mcp.py @@ -444,7 +444,7 @@ class TestMCPAppApi: ) session = MagicMock() - session.query().where().first.side_effect = [server, app] + session.scalar.side_effect = [server, app] result_server, result_app = api._get_mcp_server_and_app("server-1", session)