diff --git a/api/controllers/console/app/workflow_draft_variable.py b/api/controllers/console/app/workflow_draft_variable.py index a008d937eca..f7ddb9e93b1 100644 --- a/api/controllers/console/app/workflow_draft_variable.py +++ b/api/controllers/console/app/workflow_draft_variable.py @@ -7,12 +7,13 @@ from uuid import UUID from flask import Response from flask_restx import Resource from pydantic import BaseModel, ConfigDict, Field, model_validator -from sqlalchemy.orm import sessionmaker +from sqlalchemy.orm import Session, sessionmaker from controllers.common.errors import InvalidArgumentError, NotFoundError from controllers.common.fields import SimpleResultResponse from controllers.common.rbac import PlainApp, RBACCheck from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models +from controllers.common.session import with_session from controllers.console import console_ns from controllers.console.app.error import ( DraftWorkflowNotExist, @@ -41,6 +42,8 @@ from fields.workflow_draft_variable_fields import ( WorkflowDraftVariableListWithoutValueResponse, WorkflowDraftVariableResponse, WorkflowDraftVariableWithoutValueResponse, + draft_variable_list_response_source, + draft_variable_response_source, ) from graphon.variables.types import SegmentType from libs.external_api import error_responses @@ -248,12 +251,12 @@ class WorkflowVariableCollectionApi(Resource): @console_ns.doc(description="Delete all draft workflow variables") @console_ns.response(204, "Workflow variables deleted successfully") @_api_prerequisite - def delete(self, current_user: Account, app_model: App): + @with_session + def delete(self, session: Session, current_user: Account, app_model: App): draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) draft_var_srv.delete_user_workflow_variables(app_model.id, user_id=current_user.id) - db.session.commit() return Response("", 204) @@ -286,25 +289,27 @@ class NodeVariableCollectionApi(Resource): ) @_api_prerequisite @rbac_permission_required(RBACCheck(RBACPermission.APP_VIEW_LAYOUT, PlainApp())) - def get(self, current_user: Account, app_model: App, node_id: str): + @with_session(write=False) + def get(self, session: Session, current_user: Account, app_model: App, node_id: str): validate_node_id(node_id) - with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session: - draft_var_srv = WorkflowDraftVariableService( - session=session, - ) - node_vars = draft_var_srv.list_node_variables(app_model.id, node_id, user_id=current_user.id) - - return dump_response(WorkflowDraftVariableListResponse, node_vars) + draft_var_srv = WorkflowDraftVariableService( + session=session, + ) + node_vars = draft_var_srv.list_node_variables(app_model.id, node_id, user_id=current_user.id) + return dump_response( + WorkflowDraftVariableListResponse, + draft_variable_list_response_source(node_vars, session=session), + ) @console_ns.doc("delete_node_variables") @console_ns.doc(description="Delete all variables for a specific node") @console_ns.response(204, "Node variables deleted successfully") @_api_prerequisite - def delete(self, current_user: Account, app_model: App, node_id: str): + @with_session + def delete(self, session: Session, current_user: Account, app_model: App, node_id: str): validate_node_id(node_id) - srv = WorkflowDraftVariableService(db.session()) + srv = WorkflowDraftVariableService(session) srv.delete_node_variables(app_model.id, node_id, user_id=current_user.id) - db.session.commit() return Response("", 204) @@ -324,9 +329,10 @@ class VariableApi(Resource): @console_ns.response(404, "Variable not found") @_api_prerequisite @rbac_permission_required(RBACCheck(RBACPermission.APP_VIEW_LAYOUT, PlainApp())) - def get(self, current_user: Account, app_model: App, variable_id: UUID): + @with_session(write=False) + def get(self, session: Session, current_user: Account, app_model: App, variable_id: UUID): draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) variable_id_str = str(variable_id) variable = ensure_variable_access( @@ -335,7 +341,7 @@ class VariableApi(Resource): variable_id=variable_id_str, current_user_id=current_user.id, ) - return dump_response(WorkflowDraftVariableResponse, variable) + return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session)) @console_ns.doc("update_variable") @console_ns.doc(description="Update a workflow variable") @@ -348,8 +354,10 @@ class VariableApi(Resource): @console_ns.response(404, "Variable not found") @_api_prerequisite @model_validate(WorkflowDraftVariableUpdatePayload) + @with_session def patch( self, + session: Session, req_data: WorkflowDraftVariableUpdatePayload, current_user: Account, app_model: App, @@ -377,7 +385,7 @@ class VariableApi(Resource): # } draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) variable_id_str = str(variable_id) @@ -391,7 +399,9 @@ class VariableApi(Resource): new_name = req_data.name raw_value = req_data.value if new_name is None and raw_value is None: - return dump_response(WorkflowDraftVariableResponse, variable) + return dump_response( + WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session) + ) new_value = None if raw_value is not None: @@ -418,17 +428,17 @@ class VariableApi(Resource): pass new_value = build_segment_with_type(variable.value_type, raw_value) draft_var_srv.update_variable(variable, name=new_name, value=new_value) - db.session.commit() - return dump_response(WorkflowDraftVariableResponse, variable) + return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session)) @console_ns.doc("delete_variable") @console_ns.doc(description="Delete a workflow variable") @console_ns.response(204, "Variable deleted successfully") @console_ns.response(404, "Variable not found") @_api_prerequisite - def delete(self, current_user: Account, app_model: App, variable_id: UUID): + @with_session + def delete(self, session: Session, current_user: Account, app_model: App, variable_id: UUID): draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) variable_id_str = str(variable_id) variable = ensure_variable_access( @@ -438,7 +448,6 @@ class VariableApi(Resource): current_user_id=current_user.id, ) draft_var_srv.delete_variable(variable) - db.session.commit() return Response("", 204) @@ -455,13 +464,14 @@ class VariableResetApi(Resource): @console_ns.response(204, "Variable reset (no content)") @console_ns.response(404, "Variable not found") @_api_prerequisite - def put(self, current_user: Account, app_model: App, variable_id: UUID): + @with_session + def put(self, session: Session, current_user: Account, app_model: App, variable_id: UUID): draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) workflow_srv = WorkflowService() - draft_workflow = workflow_srv.get_draft_workflow(app_model, session=db.session()) + draft_workflow = workflow_srv.get_draft_workflow(app_model, session=session) if draft_workflow is None: raise NotFoundError( f"Draft workflow not found, app_id={app_model.id}", @@ -475,27 +485,27 @@ class VariableResetApi(Resource): ) resetted = draft_var_srv.reset_variable(draft_workflow, variable) - db.session.commit() if resetted is None: return Response("", 204) - return dump_response(WorkflowDraftVariableResponse, resetted) + return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(resetted, session=session)) -def _get_variable_list(app_model: App, node_id: str, current_user_id: str) -> WorkflowDraftVariableList: - with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session: - draft_var_srv = WorkflowDraftVariableService( - session=session, +def _get_variable_list( + app_model: App, node_id: str, current_user_id: str, *, session: Session +) -> WorkflowDraftVariableList: + draft_var_srv = WorkflowDraftVariableService( + session=session, + ) + if node_id == CONVERSATION_VARIABLE_NODE_ID: + draft_vars = draft_var_srv.list_conversation_variables(app_model.id, user_id=current_user_id) + elif node_id == SYSTEM_VARIABLE_NODE_ID: + draft_vars = draft_var_srv.list_system_variables(app_model.id, user_id=current_user_id) + else: + draft_vars = draft_var_srv.list_node_variables( + app_id=app_model.id, + node_id=node_id, + user_id=current_user_id, ) - if node_id == CONVERSATION_VARIABLE_NODE_ID: - draft_vars = draft_var_srv.list_conversation_variables(app_model.id, user_id=current_user_id) - elif node_id == SYSTEM_VARIABLE_NODE_ID: - draft_vars = draft_var_srv.list_system_variables(app_model.id, user_id=current_user_id) - else: - draft_vars = draft_var_srv.list_node_variables( - app_id=app_model.id, - node_id=node_id, - user_id=current_user_id, - ) return draft_vars @@ -512,19 +522,22 @@ class ConversationVariableCollectionApi(Resource): @console_ns.response(404, "Draft workflow not found") @_api_prerequisite @rbac_permission_required(RBACCheck(RBACPermission.APP_VIEW_LAYOUT, PlainApp())) - def get(self, current_user: Account, app_model: App): + @with_session + def get(self, session: Session, current_user: Account, app_model: App): # NOTE(QuantumGhost): Prefill conversation variables into the draft variables table # so their IDs can be returned to the caller. workflow_srv = WorkflowService() - draft_workflow = workflow_srv.get_draft_workflow(app_model, session=db.session()) + draft_workflow = workflow_srv.get_draft_workflow(app_model, session=session) if draft_workflow is None: raise NotFoundError(description=f"draft workflow not found, id={app_model.id}") - draft_var_srv = WorkflowDraftVariableService(db.session()) + draft_var_srv = WorkflowDraftVariableService(session) draft_var_srv.prefill_conversation_variable_default_values(draft_workflow, user_id=current_user.id) - db.session.commit() return dump_response( WorkflowDraftVariableListResponse, - _get_variable_list(app_model, CONVERSATION_VARIABLE_NODE_ID, current_user.id), + draft_variable_list_response_source( + _get_variable_list(app_model, CONVERSATION_VARIABLE_NODE_ID, current_user.id, session=session), + session=session, + ), ) @console_ns.expect(console_ns.models[ConversationVariableUpdatePayload.__name__]) @@ -576,10 +589,14 @@ class SystemVariableCollectionApi(Resource): ) @_api_prerequisite @rbac_permission_required(RBACCheck(RBACPermission.APP_VIEW_LAYOUT, PlainApp())) - def get(self, current_user: Account, app_model: App): + @with_session(write=False) + def get(self, session: Session, current_user: Account, app_model: App): return dump_response( WorkflowDraftVariableListResponse, - _get_variable_list(app_model, SYSTEM_VARIABLE_NODE_ID, current_user.id), + draft_variable_list_response_source( + _get_variable_list(app_model, SYSTEM_VARIABLE_NODE_ID, current_user.id, session=session), + session=session, + ), ) diff --git a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_draft_variable.py b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_draft_variable.py index aa69080197c..d5dafd1f6c4 100644 --- a/api/controllers/console/datasets/rag_pipeline/rag_pipeline_draft_variable.py +++ b/api/controllers/console/datasets/rag_pipeline/rag_pipeline_draft_variable.py @@ -7,11 +7,12 @@ from uuid import UUID from flask import Response from flask_restx import Resource from pydantic import BaseModel, Field -from sqlalchemy.orm import sessionmaker +from sqlalchemy.orm import Session, sessionmaker from controllers.common.errors import InvalidArgumentError, NotFoundError from controllers.common.rbac import DatasetByPipeline, RBACCheck from controllers.common.schema import query_params_from_model, register_schema_models +from controllers.common.session import with_session from controllers.console import console_ns from controllers.console.app.error import ( DraftWorkflowNotExist, @@ -40,6 +41,8 @@ from fields.workflow_draft_variable_fields import ( WorkflowDraftVariableListResponse, WorkflowDraftVariableListWithoutValueResponse, WorkflowDraftVariableResponse, + draft_variable_list_response_source, + draft_variable_response_source, ) from graphon.variables.types import SegmentType from libs.external_api import error_responses @@ -132,12 +135,12 @@ class RagPipelineVariableCollectionApi(Resource): @console_ns.response(204, "Workflow variables deleted successfully") @_api_prerequisite - def delete(self, current_user: Account, pipeline: Pipeline): + @with_session + def delete(self, session: Session, current_user: Account, pipeline: Pipeline): draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) draft_var_srv.delete_user_workflow_variables(pipeline.id, user_id=current_user.id) - db.session.commit() return Response("", 204) @@ -167,23 +170,25 @@ class RagPipelineNodeVariableCollectionApi(Resource): console_ns.models[WorkflowDraftVariableListResponse.__name__], ) @_api_prerequisite - def get(self, current_user: Account, pipeline: Pipeline, node_id: str): + @with_session(write=False) + def get(self, session: Session, current_user: Account, pipeline: Pipeline, node_id: str): validate_node_id(node_id) - with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session: - draft_var_srv = WorkflowDraftVariableService( - session=session, - ) - node_vars = draft_var_srv.list_node_variables(pipeline.id, node_id, user_id=current_user.id) - - return dump_response(WorkflowDraftVariableListResponse, node_vars) + draft_var_srv = WorkflowDraftVariableService( + session=session, + ) + node_vars = draft_var_srv.list_node_variables(pipeline.id, node_id, user_id=current_user.id) + return dump_response( + WorkflowDraftVariableListResponse, + draft_variable_list_response_source(node_vars, session=session), + ) @console_ns.response(204, "Node variables deleted successfully") @_api_prerequisite - def delete(self, current_user: Account, pipeline: Pipeline, node_id: str): + @with_session + def delete(self, session: Session, current_user: Account, pipeline: Pipeline, node_id: str): validate_node_id(node_id) - srv = WorkflowDraftVariableService(db.session()) + srv = WorkflowDraftVariableService(session) srv.delete_node_variables(pipeline.id, node_id, user_id=current_user.id) - db.session.commit() return Response("", 204) @@ -198,9 +203,10 @@ class RagPipelineVariableApi(Resource): console_ns.models[WorkflowDraftVariableResponse.__name__], ) @_api_prerequisite - def get(self, current_user: Account, pipeline: Pipeline, variable_id: UUID): + @with_session(write=False) + def get(self, session: Session, current_user: Account, pipeline: Pipeline, variable_id: UUID): draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) variable_id_str = str(variable_id) variable = ensure_variable_access( @@ -209,7 +215,7 @@ class RagPipelineVariableApi(Resource): variable_id=variable_id_str, current_user_id=current_user.id, ) - return dump_response(WorkflowDraftVariableResponse, variable) + return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session)) @console_ns.response( 200, @@ -219,8 +225,10 @@ class RagPipelineVariableApi(Resource): @_api_prerequisite @console_ns.expect(console_ns.models[WorkflowDraftVariablePatchPayload.__name__]) @model_validate(WorkflowDraftVariablePatchPayload) + @with_session def patch( self, + session: Session, req_data: WorkflowDraftVariablePatchPayload, current_user: Account, pipeline: Pipeline, @@ -248,7 +256,7 @@ class RagPipelineVariableApi(Resource): # } draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) args = req_data.model_dump(exclude_none=True) @@ -263,7 +271,9 @@ class RagPipelineVariableApi(Resource): new_name = args.get(self._PATCH_NAME_FIELD, None) raw_value = args.get(self._PATCH_VALUE_FIELD, None) if new_name is None and raw_value is None: - return dump_response(WorkflowDraftVariableResponse, variable) + return dump_response( + WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session) + ) new_value = None if raw_value is not None: @@ -290,14 +300,14 @@ class RagPipelineVariableApi(Resource): pass new_value = build_segment_with_type(variable.value_type, raw_value) draft_var_srv.update_variable(variable, name=new_name, value=new_value) - db.session.commit() - return dump_response(WorkflowDraftVariableResponse, variable) + return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session)) @console_ns.response(204, "Variable deleted successfully") @_api_prerequisite - def delete(self, current_user: Account, pipeline: Pipeline, variable_id: UUID): + @with_session + def delete(self, session: Session, current_user: Account, pipeline: Pipeline, variable_id: UUID): draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) variable_id_str = str(variable_id) variable = ensure_variable_access( @@ -307,7 +317,6 @@ class RagPipelineVariableApi(Resource): current_user_id=current_user.id, ) draft_var_srv.delete_variable(variable) - db.session.commit() return Response("", 204) @@ -320,12 +329,13 @@ class RagPipelineVariableResetApi(Resource): ) @console_ns.response(204, "Variable reset (no content)") @_api_prerequisite - def put(self, current_user: Account, pipeline: Pipeline, variable_id: UUID): + @with_session + def put(self, session: Session, current_user: Account, pipeline: Pipeline, variable_id: UUID): draft_var_srv = WorkflowDraftVariableService( - session=db.session(), + session=session, ) - rag_pipeline_service = RagPipelineService(db.session()) + rag_pipeline_service = RagPipelineService(session) draft_workflow = rag_pipeline_service.get_draft_workflow(pipeline=pipeline) if draft_workflow is None: raise NotFoundError( @@ -340,23 +350,23 @@ class RagPipelineVariableResetApi(Resource): ) resetted = draft_var_srv.reset_variable(draft_workflow, variable) - db.session.commit() if resetted is None: return Response("", 204) - return dump_response(WorkflowDraftVariableResponse, resetted) + return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(resetted, session=session)) -def _get_variable_list(pipeline: Pipeline, node_id: str, current_user_id: str) -> WorkflowDraftVariableList: - with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session: - draft_var_srv = WorkflowDraftVariableService( - session=session, - ) - if node_id == CONVERSATION_VARIABLE_NODE_ID: - draft_vars = draft_var_srv.list_conversation_variables(pipeline.id, user_id=current_user_id) - elif node_id == SYSTEM_VARIABLE_NODE_ID: - draft_vars = draft_var_srv.list_system_variables(pipeline.id, user_id=current_user_id) - else: - draft_vars = draft_var_srv.list_node_variables(app_id=pipeline.id, node_id=node_id, user_id=current_user_id) +def _get_variable_list( + pipeline: Pipeline, node_id: str, current_user_id: str, *, session: Session +) -> WorkflowDraftVariableList: + draft_var_srv = WorkflowDraftVariableService( + session=session, + ) + if node_id == CONVERSATION_VARIABLE_NODE_ID: + draft_vars = draft_var_srv.list_conversation_variables(pipeline.id, user_id=current_user_id) + elif node_id == SYSTEM_VARIABLE_NODE_ID: + draft_vars = draft_var_srv.list_system_variables(pipeline.id, user_id=current_user_id) + else: + draft_vars = draft_var_srv.list_node_variables(app_id=pipeline.id, node_id=node_id, user_id=current_user_id) return draft_vars @@ -368,10 +378,14 @@ class RagPipelineSystemVariableCollectionApi(Resource): console_ns.models[WorkflowDraftVariableListResponse.__name__], ) @_api_prerequisite - def get(self, current_user: Account, pipeline: Pipeline): + @with_session(write=False) + def get(self, session: Session, current_user: Account, pipeline: Pipeline): return dump_response( WorkflowDraftVariableListResponse, - _get_variable_list(pipeline, SYSTEM_VARIABLE_NODE_ID, current_user.id), + draft_variable_list_response_source( + _get_variable_list(pipeline, SYSTEM_VARIABLE_NODE_ID, current_user.id, session=session), + session=session, + ), ) diff --git a/api/controllers/console/snippets/snippet_workflow_draft_variable.py b/api/controllers/console/snippets/snippet_workflow_draft_variable.py index a9dc94fe26b..f83215b2e8a 100644 --- a/api/controllers/console/snippets/snippet_workflow_draft_variable.py +++ b/api/controllers/console/snippets/snippet_workflow_draft_variable.py @@ -48,6 +48,8 @@ from fields.workflow_draft_variable_fields import ( WorkflowDraftVariableListResponse, WorkflowDraftVariableListWithoutValueResponse, WorkflowDraftVariableResponse, + draft_variable_list_response_source, + draft_variable_response_source, ) from graphon.variables.types import SegmentType from libs.external_api import error_responses @@ -160,7 +162,9 @@ class SnippetNodeVariableCollectionApi(Resource): draft_var_srv = WorkflowDraftVariableService(session=session) node_vars = draft_var_srv.list_node_variables(snippet.id, node_id, user_id=current_user.id) - return dump_response(WorkflowDraftVariableListResponse, node_vars) + return dump_response( + WorkflowDraftVariableListResponse, draft_variable_list_response_source(node_vars, session=session) + ) @console_ns.doc("delete_snippet_node_variables") @console_ns.doc(description="Delete all variables for a specific node (snippet draft workflow)") @@ -197,7 +201,7 @@ class SnippetVariableApi(Resource): current_user_id=current_user.id, ) _ensure_snippet_draft_variable_row_allowed(variable=variable, variable_id=variable_id) - return dump_response(WorkflowDraftVariableResponse, variable) + return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session)) @console_ns.doc("update_snippet_workflow_variable") @console_ns.doc(description="Update a draft workflow variable (snippet scope)") @@ -232,7 +236,9 @@ class SnippetVariableApi(Resource): new_name = req_data.name raw_value = req_data.value if new_name is None and raw_value is None: - return dump_response(WorkflowDraftVariableResponse, variable) + return dump_response( + WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session) + ) new_value = None if raw_value is not None: @@ -256,7 +262,7 @@ class SnippetVariableApi(Resource): ) new_value = build_segment_with_type(variable.value_type, raw_value) draft_var_srv.update_variable(variable, name=new_name, value=new_value) - return dump_response(WorkflowDraftVariableResponse, variable) + return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(variable, session=session)) @console_ns.doc("delete_snippet_workflow_variable") @console_ns.doc(description="Delete a draft workflow variable (snippet scope)") @@ -311,7 +317,7 @@ class SnippetVariableResetApi(Resource): resetted = draft_var_srv.reset_variable(draft_workflow, variable) if resetted is None: return Response("", 204) - return dump_response(WorkflowDraftVariableResponse, resetted) + return dump_response(WorkflowDraftVariableResponse, draft_variable_response_source(resetted, session=session)) @console_ns.route("/snippets//workflows/draft/conversation-variables") diff --git a/api/fields/workflow_draft_variable_fields.py b/api/fields/workflow_draft_variable_fields.py index ba3b754a078..c7e55a8472c 100644 --- a/api/fields/workflow_draft_variable_fields.py +++ b/api/fields/workflow_draft_variable_fields.py @@ -3,8 +3,9 @@ from __future__ import annotations from typing import Any, override from pydantic import model_validator +from sqlalchemy.orm import Session -from fields.base import ResponseModel +from fields.base import ResponseModel, SessionResponseSource from graphon.file import helpers as file_helpers from graphon.variables.segment_group import SegmentGroup from graphon.variables.segments import ArrayFileSegment, FileSegment, Segment @@ -25,7 +26,27 @@ def _convert_values_to_json_serializable_object(value: Segment) -> RecursiveJSON return value.value -def _serialize_var_value(variable: WorkflowDraftVariable) -> RecursiveJSONValue: +class WorkflowDraftVariableResponseSource(SessionResponseSource[WorkflowDraftVariable]): + """Expose the session-backed value decoding during response validation.""" + + def get_value(self) -> Segment: + return self._source.get_value(session=self._session) + + +def draft_variable_response_source( + variable: WorkflowDraftVariable, *, session: Session +) -> WorkflowDraftVariableResponseSource: + return WorkflowDraftVariableResponseSource(variable, session=session) + + +def draft_variable_list_response_source(variable_list: Any, *, session: Session) -> dict[str, Any]: + """Wrap each variable in a list payload so value decoding resolves via the session.""" + return { + "items": [draft_variable_response_source(variable, session=session) for variable in variable_list.variables], + } + + +def _serialize_var_value(variable: WorkflowDraftVariableResponseSource) -> RecursiveJSONValue: value = variable.get_value() # Create a copy to avoid mutating the model's cached deserialized value. value = value.model_copy(deep=True) @@ -105,12 +126,18 @@ class WorkflowDraftVariableResponse(WorkflowDraftVariableWithoutValueResponse): @classmethod @override def _from_workflow_draft_variable(cls, value: Any) -> Any: - if isinstance(value, WorkflowDraftVariable): + if isinstance(value, WorkflowDraftVariableResponseSource): return { - **_serialize_without_value(value), + **_serialize_without_value(value._source), "value": _serialize_var_value(value), - "full_content": _serialize_full_content(value), + "full_content": _serialize_full_content(value._source), } + if isinstance(value, WorkflowDraftVariable): + # Decoding the value needs a database session; a bare model would fall through to + # from_attributes and expose the raw serialized JSON string instead of the decoded value. + raise TypeError( + "WorkflowDraftVariableResponse requires draft_variable_response_source(variable, session=...)" + ) return value diff --git a/api/models/workflow.py b/api/models/workflow.py index 8ae5e4a2ab6..f36808b4ecf 100644 --- a/api/models/workflow.py +++ b/api/models/workflow.py @@ -22,7 +22,7 @@ from sqlalchemy import ( orm, select, ) -from sqlalchemy.orm import Mapped, mapped_column +from sqlalchemy.orm import Mapped, Session, mapped_column from constants.trigger import TRIGGER_PLUGIN_NODE_TYPE, is_trigger_node_type from core.workflow.environment_variables import load_environment_variables @@ -74,7 +74,6 @@ from libs import helper from .account import Account from .base import DefaultFieldsDCMixin, TypeBase -from .engine import db from .enums import CreatorUserRole, DraftVariableType, ExecutionOffLoadType, WorkflowRunTriggeredFrom # UploadFile and workflow execution models use TypeBase; importing the class object keeps @@ -92,10 +91,10 @@ SerializedWorkflowValue = dict[str, Any] SerializedWorkflowVariables = dict[str, SerializedWorkflowValue] -def _resolve_workflow_app_tenant_id(app_id: str) -> str: +def _resolve_workflow_app_tenant_id(session: Session, app_id: str) -> str: from .model import App - tenant_id = db.session.scalar(select(App.tenant_id).where(App.id == app_id)) + tenant_id = session.scalar(select(App.tenant_id).where(App.id == app_id)) if not tenant_id: raise ValueError(f"Unable to resolve tenant_id for app {app_id}") return tenant_id @@ -1775,11 +1774,11 @@ class WorkflowDraftVariable(TypeBase): def _set_selector(self, value: list[str]): self.selector = json.dumps(value) - def _loads_value(self) -> Segment: + def _loads_value(self, *, session: Session) -> Segment: value = json.loads(self.value) - return self.build_segment_from_serialized_value(self.value_type, value) + return self.build_segment_from_serialized_value(self.value_type, value, session=session) - def _rebuild_file_types(self, value: Any): + def _rebuild_file_types(self, value: Any, *, session: Session): # NOTE(QuantumGhost): Temporary workaround for structured data handling. # By this point, `output` has been converted to dict by # `WorkflowEntry.handle_special_values`, so we need to @@ -1794,7 +1793,7 @@ class WorkflowDraftVariable(TypeBase): case dict(): if not maybe_file_object(value): return cast(Any, value) - tenant_id = _resolve_workflow_app_tenant_id(self.app_id) + tenant_id = _resolve_workflow_app_tenant_id(session, self.app_id) return build_file_from_stored_mapping( file_mapping=cast(dict[str, Any], value), tenant_id=tenant_id, @@ -1804,7 +1803,7 @@ class WorkflowDraftVariable(TypeBase): first: Any = value_list[0] if not maybe_file_object(first): return cast(Any, value) - tenant_id = _resolve_workflow_app_tenant_id(self.app_id) + tenant_id = _resolve_workflow_app_tenant_id(session, self.app_id) file_list: list[File] = [] for item in value_list: file_list.append( @@ -1817,7 +1816,9 @@ class WorkflowDraftVariable(TypeBase): case _: return cast(Any, value) - def build_segment_from_serialized_value(self, segment_type: SegmentType, value: Any) -> Segment: + def build_segment_from_serialized_value( + self, segment_type: SegmentType, value: Any, *, session: Session + ) -> Segment: # Persisted draft variable rows may contain historical file payloads. # Rebuild them through the file factory so tenant ownership, signed URLs, # and storage-backed metadata come from canonical records instead of the @@ -1828,14 +1829,14 @@ class WorkflowDraftVariable(TypeBase): case File(): return build_segment_with_type(segment_type, value) case dict(): - file = self._rebuild_file_types(value) + file = self._rebuild_file_types(value, session=session) return build_segment_with_type(segment_type, file) case _: raise TypeMismatchError(f"expected dict or File for FileSegment, got {type(value)}") case SegmentType.ARRAY_FILE: if not isinstance(value, list): raise TypeMismatchError(f"expected list for ArrayFileSegment, got {type(value)}") - file_list = self._rebuild_file_types(value) + file_list = self._rebuild_file_types(value, session=session) return build_segment_with_type(segment_type=segment_type, value=file_list) case _: return build_segment_with_type(segment_type=segment_type, value=value) @@ -1890,7 +1891,7 @@ class WorkflowDraftVariable(TypeBase): case _: return build_segment_with_type(segment_type=segment_type, value=value) - def get_value(self) -> Segment: + def get_value(self, *, session: Session) -> Segment: """Decode the serialized value into its corresponding `Segment` object. This method caches the result, so repeated calls will return the same @@ -1908,7 +1909,7 @@ class WorkflowDraftVariable(TypeBase): if self.__value is not None: return self.__value - value = self._loads_value() + value = self._loads_value(session=session) self.__value = value return value diff --git a/api/services/workflow_draft_variable_service.py b/api/services/workflow_draft_variable_service.py index b83076d552d..081cedb1476 100644 --- a/api/services/workflow_draft_variable_service.py +++ b/api/services/workflow_draft_variable_service.py @@ -114,21 +114,37 @@ class DraftVarLoader(VariableLoader): # Map each selector (as a tuple via `_selector_to_tuple`) to its corresponding variable instance. variable_by_selector: dict[tuple[str, str], VariableBase] = {} + files: list[File] = [] + offloaded_draft_vars: list[WorkflowDraftVariable] = [] with Session(bind=self._engine, expire_on_commit=False) as session: srv = WorkflowDraftVariableService(session) draft_vars = srv.get_draft_variables_by_selectors(self._app_id, selectors, user_id=self._user_id) - # Important: - files: list[File] = [] - # FileSegment and ArrayFileSegment are not subject to offloading, so their values - # can be safely accessed before any offloading logic is applied. - for draft_var in draft_vars: - value = draft_var.get_value() - match value: - case FileSegment(): - files.append(value.value) - case ArrayFileSegment(): - files.extend(value.value) + for draft_var in draft_vars: + if draft_var.is_truncated(): + offloaded_draft_vars.append(draft_var) + continue + + # FileSegment and ArrayFileSegment are not subject to offloading, so their values + # can be safely accessed before any offloading logic is applied. + segment = draft_var.get_value(session=session) + match segment: + case FileSegment(): + files.append(segment.value) + case ArrayFileSegment(): + files.extend(segment.value) + variable = segment_to_variable( + segment=segment, + selector=draft_var.get_selector(), + variable_id=draft_var.id, + name=draft_var.name, + description=draft_var.description, + ) + selector_tuple = self._selector_to_tuple(variable.selector) + variable_by_selector[selector_tuple] = variable + + # The files above are the same instances held by the variables, so loading their + # storage keys here is visible through `variable_by_selector` as well. with Session(bind=self._engine) as session: storage_key_loader = StorageKeyLoader( session, @@ -137,23 +153,6 @@ class DraftVarLoader(VariableLoader): ) storage_key_loader.load_storage_keys(files) - offloaded_draft_vars = [] - for draft_var in draft_vars: - if draft_var.is_truncated(): - offloaded_draft_vars.append(draft_var) - continue - - segment = draft_var.get_value() - variable = segment_to_variable( - segment=segment, - selector=draft_var.get_selector(), - variable_id=draft_var.id, - name=draft_var.name, - description=draft_var.description, - ) - selector_tuple = self._selector_to_tuple(variable.selector) - variable_by_selector[selector_tuple] = variable - # Load offloaded variables using multithreading. # This approach reduces loading time by querying external systems concurrently. with ThreadPoolExecutor(max_workers=10) as executor: @@ -188,7 +187,11 @@ class DraftVarLoader(VariableLoader): return (draft_var.node_id, draft_var.name), variable deserialized = json.loads(content) - segment = draft_var.build_segment_from_serialized_value(variable_file.value_type, deserialized) + # Runs on a worker thread, so it opens its own session instead of sharing one. + with Session(bind=self._engine, expire_on_commit=False) as session: + segment = draft_var.build_segment_from_serialized_value( + variable_file.value_type, deserialized, session=session + ) variable = segment_to_variable( segment=segment, selector=draft_var.get_selector(), @@ -573,7 +576,7 @@ class WorkflowDraftVariableService: ) if draft_var is None: return None - segment = draft_var.get_value() + segment = draft_var.get_value(session=self._session) if not isinstance(segment, StringSegment): logger.warning( "sys.conversation_id variable is not a string: app_id=%s, user_id=%s, id=%s", diff --git a/api/tests/e2e_tests/containers/controllers/console/app/test_workflow_draft_variable.py b/api/tests/e2e_tests/containers/controllers/console/app/test_workflow_draft_variable.py index fbaf21c6dc0..02db201a8cd 100644 --- a/api/tests/e2e_tests/containers/controllers/console/app/test_workflow_draft_variable.py +++ b/api/tests/e2e_tests/containers/controllers/console/app/test_workflow_draft_variable.py @@ -288,6 +288,9 @@ def test_variable_api_patch_success( assert payload["id"] == variable.id assert payload["name"] == "renamed_var" + # The handler commits through its own decorator-injected session, so expire this + # session's cached instance to read the persisted row instead of the identity map. + db_session_with_containers.expire_all() refreshed = db_session_with_containers.scalar( select(WorkflowDraftVariable).where(WorkflowDraftVariable.id == variable.id) ) diff --git a/api/tests/e2e_tests/containers/services/test_workflow_draft_variable_service.py b/api/tests/e2e_tests/containers/services/test_workflow_draft_variable_service.py index 21da6f6d6d8..2355c615a4c 100644 --- a/api/tests/e2e_tests/containers/services/test_workflow_draft_variable_service.py +++ b/api/tests/e2e_tests/containers/services/test_workflow_draft_variable_service.py @@ -213,7 +213,7 @@ class TestWorkflowDraftVariableService: assert retrieved_variable.id == variable.id assert retrieved_variable.name == "test_var" assert retrieved_variable.app_id == app.id - assert retrieved_variable.get_value().value == test_value.value + assert retrieved_variable.get_value(session=db_session_with_containers).value == test_value.value def test_get_variable_not_found(self, db_session_with_containers: Session, mock_external_service_dependencies): """ @@ -274,11 +274,11 @@ class TestWorkflowDraftVariableService: assert "var3" in var_names for var in retrieved_variables: if var.name == "var1": - assert var.get_value().value == var1_value.value + assert var.get_value(session=db_session_with_containers).value == var1_value.value elif var.name == "var2": - assert var.get_value().value == var2_value.value + assert var.get_value(session=db_session_with_containers).value == var2_value.value elif var.name == "var3": - assert var.get_value().value == var3_value.value + assert var.get_value(session=db_session_with_containers).value == var3_value.value def test_list_variables_without_values_success( self, db_session_with_containers: Session, mock_external_service_dependencies @@ -431,12 +431,12 @@ class TestWorkflowDraftVariableService: service = WorkflowDraftVariableService(db_session_with_containers) updated_variable = service.update_variable(variable, name="new_name", value=new_value) assert updated_variable.name == "new_name" - assert updated_variable.get_value().value == new_value.value + assert updated_variable.get_value(session=db_session_with_containers).value == new_value.value assert updated_variable.last_edited_at is not None db_session_with_containers.refresh(variable) assert variable.name == "new_name" - assert variable.get_value().value == new_value.value + assert variable.get_value(session=db_session_with_containers).value == new_value.value assert variable.last_edited_at is not None def test_update_variable_not_editable( @@ -509,10 +509,10 @@ class TestWorkflowDraftVariableService: service = WorkflowDraftVariableService(db_session_with_containers) reset_variable = service.reset_variable(workflow, variable) assert reset_variable is not None - assert reset_variable.get_value().value == "default_value" + assert reset_variable.get_value(session=db_session_with_containers).value == "default_value" assert reset_variable.last_edited_at is None db_session_with_containers.refresh(variable) - assert variable.get_value().value == "default_value" + assert variable.get_value(session=db_session_with_containers).value == "default_value" assert variable.last_edited_at is None def test_delete_variable_success(self, db_session_with_containers: Session, mock_external_service_dependencies): diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow_draft_variable_handlers.py b/api/tests/unit_tests/controllers/console/app/test_workflow_draft_variable_handlers.py new file mode 100644 index 00000000000..618b94ef561 --- /dev/null +++ b/api/tests/unit_tests/controllers/console/app/test_workflow_draft_variable_handlers.py @@ -0,0 +1,219 @@ +"""Direct-call tests for the app draft-variable handlers converted to `@with_session`.""" + +import json +from dataclasses import dataclass +from inspect import unwrap + +import pytest +from flask import Flask +from sqlalchemy import Engine +from sqlalchemy.orm import Session + +from controllers.console.app import workflow_draft_variable as module +from graphon.variables.segments import StringSegment +from models.account import Account, AccountStatus +from models.model import App, AppMode +from models.workflow import Workflow, WorkflowDraftVariable + +_APP_ID = "app-1" +_USER_ID = "user-1" + + +def _make_account() -> Account: + account = Account( + name="tester", + email="tester@example.com", + status=AccountStatus.ACTIVE, + ) + account.id = _USER_ID # type: ignore[assignment] + return account + + +def _make_app_model() -> App: + return App( + id=_APP_ID, + tenant_id="tenant-1", + name="App", + mode=AppMode.ADVANCED_CHAT, + enable_site=True, + enable_api=True, + ) + + +def _make_draft_workflow() -> Workflow: + workflow = Workflow.new( + tenant_id="tenant-1", + app_id=_APP_ID, + type="workflow", + version=Workflow.VERSION_DRAFT, + graph=json.dumps({"nodes": [], "edges": []}), + features="{}", + created_by=_USER_ID, + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + workflow.id = "workflow-1" + return workflow + + +def _make_node_variable(variable_id: str, *, node_id: str = "llm-1") -> WorkflowDraftVariable: + variable = WorkflowDraftVariable.new_node_variable( + app_id=_APP_ID, + user_id=_USER_ID, + node_id=node_id, + name=variable_id, + value=StringSegment(value=f"value-{variable_id}"), + node_execution_id="execution-1", + ) + variable.id = variable_id + return variable + + +def _make_system_variable(variable_id: str, *, name: str) -> WorkflowDraftVariable: + variable = WorkflowDraftVariable.new_sys_variable( + app_id=_APP_ID, + user_id=_USER_ID, + name=name, + value=StringSegment(value=f"value-{name}"), + node_execution_id="execution-1", + ) + variable.id = variable_id + return variable + + +@dataclass +class _Database: + engine: Engine + + +@pytest.fixture +def app(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> Flask: + app = Flask("test_workflow_draft_variable_handlers") + monkeypatch.setattr("services.workflow_service.db", _Database(sqlite_engine)) + app.config["TESTING"] = True + return app + + +def test_node_variable_collection_get_returns_decoded_values(app: Flask, sqlite_session: Session) -> None: + variable = _make_node_variable("var-1") + sqlite_session.add(variable) + sqlite_session.commit() + api = module.NodeVariableCollectionApi() + handler = unwrap(api.get) + + with app.test_request_context("/"): + result = handler(api, sqlite_session, _make_account(), app_model=_make_app_model(), node_id="llm-1") + + assert [item["id"] for item in result["items"]] == ["var-1"] + assert result["items"][0]["value"] == "value-var-1" + + +def test_variable_get_returns_decoded_value(app: Flask, sqlite_session: Session) -> None: + variable = _make_node_variable("var-1") + sqlite_session.add(variable) + sqlite_session.commit() + api = module.VariableApi() + handler = unwrap(api.get) + + with app.test_request_context("/"): + result = handler(api, sqlite_session, _make_account(), app_model=_make_app_model(), variable_id="var-1") + + assert result["id"] == "var-1" + assert result["value"] == "value-var-1" + + +def test_variable_patch_without_changes_returns_variable(app: Flask, sqlite_session: Session) -> None: + variable = _make_node_variable("var-1") + sqlite_session.add(variable) + sqlite_session.commit() + api = module.VariableApi() + handler = unwrap(api.patch) + + with app.test_request_context("/", method="PATCH", json={}): + result = handler( + api, + sqlite_session, + module.WorkflowDraftVariableUpdatePayload(), + _make_account(), + app_model=_make_app_model(), + variable_id="var-1", + ) + + assert result["id"] == "var-1" + assert result["value"] == "value-var-1" + + +def test_variable_patch_renames_variable(app: Flask, sqlite_session: Session) -> None: + variable = _make_node_variable("var-1") + sqlite_session.add(variable) + sqlite_session.commit() + api = module.VariableApi() + handler = unwrap(api.patch) + + with app.test_request_context("/", method="PATCH", json={"name": "renamed"}): + result = handler( + api, + sqlite_session, + module.WorkflowDraftVariableUpdatePayload(name="renamed"), + _make_account(), + app_model=_make_app_model(), + variable_id="var-1", + ) + + assert result["name"] == "renamed" + stored = sqlite_session.get(WorkflowDraftVariable, "var-1") + assert stored is not None + assert stored.name == "renamed" + + +def test_variable_reset_returns_default_value(app: Flask, sqlite_session: Session) -> None: + workflow = _make_draft_workflow() + variable = _make_node_variable("var-1") + variable.node_execution_id = None + sqlite_session.add_all([workflow, variable]) + sqlite_session.commit() + api = module.VariableResetApi() + handler = unwrap(api.put) + + with app.test_request_context("/", method="PUT"): + response = handler(api, sqlite_session, _make_account(), app_model=_make_app_model(), variable_id="var-1") + + # A node variable without an associated execution has no default to restore, so the + # handler answers 204 after removing the row. + assert response.status_code == 204 + + +def test_conversation_variable_collection_get_lists_conversation_variables(app: Flask, sqlite_session: Session) -> None: + workflow = _make_draft_workflow() + conv_var = WorkflowDraftVariable.new_conversation_variable( + app_id=_APP_ID, + name="conv_var", + value=StringSegment(value="conv-value"), + ) + conv_var.id = "conv-1" + conv_var.user_id = _USER_ID + sqlite_session.add_all([workflow, conv_var]) + sqlite_session.commit() + api = module.ConversationVariableCollectionApi() + handler = unwrap(api.get) + + with app.test_request_context("/"): + result = handler(api, sqlite_session, _make_account(), app_model=_make_app_model()) + + assert [item["id"] for item in result["items"]] == ["conv-1"] + assert result["items"][0]["value"] == "conv-value" + + +def test_system_variable_collection_get_lists_system_variables(app: Flask, sqlite_session: Session) -> None: + sys_var = _make_system_variable("sys-1", name="query") + sqlite_session.add(sys_var) + sqlite_session.commit() + api = module.SystemVariableCollectionApi() + handler = unwrap(api.get) + + with app.test_request_context("/"): + result = handler(api, sqlite_session, _make_account(), app_model=_make_app_model()) + + assert [item["id"] for item in result["items"]] == ["sys-1"] + assert result["items"][0]["value"] == "value-query" diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_draft_variable.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_draft_variable.py index e9011cde1e4..06ab9508272 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_draft_variable.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_draft_variable.py @@ -178,7 +178,7 @@ class TestRagPipelineVariableCollectionApi: patch("controllers.console.datasets.rag_pipeline.rag_pipeline_draft_variable.db", sqlite_db), patch("controllers.console.datasets.rag_pipeline.rag_pipeline_draft_variable.WorkflowDraftVariableService"), ): - result = method(api, editor_user, pipeline) + result = method(api, sqlite_db.session(), editor_user, pipeline) assert isinstance(result, Response) assert result.status_code == 204 @@ -206,12 +206,12 @@ class TestRagPipelineNodeVariableCollectionApi: return_value=srv, ), ): - result = method(api, editor_user, pipeline, "node1") + result = method(api, sqlite_db.session(), editor_user, pipeline, "node1") assert result == {"items": []} srv.list_node_variables.assert_called_once_with("p1", "node1", user_id="account-1") - def test_get_node_variables_invalid_node(self, app: Flask, editor_user): + def test_get_node_variables_invalid_node(self, app: Flask, sqlite_db, editor_user): api = RagPipelineNodeVariableCollectionApi() method = unwrap(api.get) @@ -219,7 +219,7 @@ class TestRagPipelineNodeVariableCollectionApi: app.test_request_context("/"), ): with pytest.raises(InvalidArgumentError): - method(api, editor_user, _pipeline(), SYSTEM_VARIABLE_NODE_ID) + method(api, sqlite_db.session(), editor_user, _pipeline(), SYSTEM_VARIABLE_NODE_ID) class TestRagPipelineVariableApi: @@ -239,7 +239,7 @@ class TestRagPipelineVariableApi: ), ): with pytest.raises(NotFoundError): - method(api, editor_user, _pipeline(), "v1") + method(api, sqlite_db.session(), editor_user, _pipeline(), "v1") def test_patch_variable_invalid_file_payload(self, app: Flask, sqlite_db, editor_user): api = RagPipelineVariableApi() @@ -263,7 +263,14 @@ class TestRagPipelineVariableApi: ), ): with pytest.raises(InvalidArgumentError): - method(api, WorkflowDraftVariablePatchPayload.model_validate(payload), editor_user, pipeline, "v1") + method( + api, + sqlite_db.session(), + WorkflowDraftVariablePatchPayload.model_validate(payload), + editor_user, + pipeline, + "v1", + ) def test_delete_variable_success(self, app: Flask, sqlite_db, editor_user): api = RagPipelineVariableApi() @@ -283,7 +290,7 @@ class TestRagPipelineVariableApi: return_value=srv, ), ): - result = method(api, editor_user, pipeline, "v1") + result = method(api, sqlite_db.session(), editor_user, pipeline, "v1") assert result.status_code == 204 @@ -318,9 +325,9 @@ def test_direct_variable_access_rejects_different_user( rag_service = MagicMock() rag_service.get_draft_workflow.return_value = _workflow(pipeline) if payload is not None: - call_args = (api, payload, editor_user, pipeline, "v1") + call_args = (api, sqlite_db.session(), payload, editor_user, pipeline, "v1") else: - call_args = (api, editor_user, pipeline, "v1") + call_args = (api, sqlite_db.session(), editor_user, pipeline, "v1") with ( app.test_request_context("/"), @@ -376,7 +383,7 @@ class TestRagPipelineVariableResetApi: return_value={"id": "v1"}, ), ): - result = method(api, editor_user, pipeline, "v1") + result = method(api, sqlite_db.session(), editor_user, pipeline, "v1") assert result == {"id": "v1"} @@ -403,7 +410,7 @@ class TestSystemAndEnvironmentVariablesApi: return_value=srv, ), ): - result = method(api, editor_user, pipeline) + result = method(api, sqlite_db.session(), editor_user, pipeline) assert result == {"items": []} srv.list_system_variables.assert_called_once_with("p1", user_id="account-1") diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_draft_variable_handlers.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_draft_variable_handlers.py new file mode 100644 index 00000000000..1c411fcd9d7 --- /dev/null +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_draft_variable_handlers.py @@ -0,0 +1,166 @@ +"""Direct-call tests for the rag-pipeline draft-variable handlers converted to `@with_session`.""" + +import json +from dataclasses import dataclass +from inspect import unwrap + +import pytest +from flask import Flask +from sqlalchemy import Engine +from sqlalchemy.orm import Session + +from controllers.console.datasets.rag_pipeline import rag_pipeline_draft_variable as module +from graphon.variables.segments import StringSegment +from models.account import Account, AccountStatus +from models.dataset import Pipeline +from models.workflow import Workflow, WorkflowDraftVariable + +_PIPELINE_ID = "pipeline-1" +_USER_ID = "user-1" + + +def _make_account() -> Account: + account = Account( + name="tester", + email="tester@example.com", + status=AccountStatus.ACTIVE, + ) + account.id = _USER_ID # type: ignore[assignment] + return account + + +def _make_pipeline() -> Pipeline: + pipeline = Pipeline( + tenant_id="tenant-1", + name="Pipeline", + ) + pipeline.id = _PIPELINE_ID + return pipeline + + +def _make_draft_workflow() -> Workflow: + workflow = Workflow.new( + tenant_id="tenant-1", + app_id=_PIPELINE_ID, + type="rag-pipeline", + version=Workflow.VERSION_DRAFT, + graph=json.dumps({"nodes": [], "edges": []}), + features="{}", + created_by=_USER_ID, + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + workflow.id = "workflow-1" + return workflow + + +def _make_node_variable(variable_id: str, *, node_id: str = "extract-1") -> WorkflowDraftVariable: + variable = WorkflowDraftVariable.new_node_variable( + app_id=_PIPELINE_ID, + user_id=_USER_ID, + node_id=node_id, + name=variable_id, + value=StringSegment(value=f"value-{variable_id}"), + node_execution_id="execution-1", + ) + variable.id = variable_id + return variable + + +@dataclass +class _Database: + engine: Engine + + +@pytest.fixture +def app(sqlite_engine: Engine, monkeypatch: pytest.MonkeyPatch) -> Flask: + app = Flask("test_rag_pipeline_draft_variable_handlers") + monkeypatch.setattr("services.rag_pipeline.rag_pipeline.db", _Database(sqlite_engine)) + app.config["TESTING"] = True + return app + + +def test_node_variable_collection_get_returns_decoded_values(app: Flask, sqlite_session: Session) -> None: + variable = _make_node_variable("var-1") + sqlite_session.add(variable) + sqlite_session.commit() + api = module.RagPipelineNodeVariableCollectionApi() + handler = unwrap(api.get) + + with app.test_request_context("/"): + result = handler(api, sqlite_session, _make_account(), pipeline=_make_pipeline(), node_id="extract-1") + + assert [item["id"] for item in result["items"]] == ["var-1"] + assert result["items"][0]["value"] == "value-var-1" + + +def test_variable_get_returns_decoded_value(app: Flask, sqlite_session: Session) -> None: + variable = _make_node_variable("var-1") + sqlite_session.add(variable) + sqlite_session.commit() + api = module.RagPipelineVariableApi() + handler = unwrap(api.get) + + with app.test_request_context("/"): + result = handler(api, sqlite_session, _make_account(), pipeline=_make_pipeline(), variable_id="var-1") + + assert result["id"] == "var-1" + assert result["value"] == "value-var-1" + + +def test_variable_patch_without_changes_returns_variable(app: Flask, sqlite_session: Session) -> None: + variable = _make_node_variable("var-1") + sqlite_session.add(variable) + sqlite_session.commit() + api = module.RagPipelineVariableApi() + handler = unwrap(api.patch) + + with app.test_request_context("/", method="PATCH", json={}): + result = handler( + api, + sqlite_session, + module.WorkflowDraftVariablePatchPayload(), + _make_account(), + pipeline=_make_pipeline(), + variable_id="var-1", + ) + + assert result["id"] == "var-1" + assert result["value"] == "value-var-1" + + +def test_variable_reset_removes_variable_without_default(app: Flask, sqlite_session: Session) -> None: + workflow = _make_draft_workflow() + variable = _make_node_variable("var-1") + variable.node_execution_id = None + sqlite_session.add_all([workflow, variable]) + sqlite_session.commit() + api = module.RagPipelineVariableResetApi() + handler = unwrap(api.put) + + with app.test_request_context("/", method="PUT"): + response = handler(api, sqlite_session, _make_account(), pipeline=_make_pipeline(), variable_id="var-1") + + assert response.status_code == 204 + + +def test_system_variable_collection_get_lists_system_variables(app: Flask, sqlite_session: Session) -> None: + sys_var = WorkflowDraftVariable.new_sys_variable( + app_id=_PIPELINE_ID, + user_id=_USER_ID, + name="query", + value=StringSegment(value="sys-value"), + node_execution_id="execution-1", + ) + sys_var.id = "sys-1" + sqlite_session.add(sys_var) + sqlite_session.commit() + api = module.RagPipelineSystemVariableCollectionApi() + handler = unwrap(api.get) + + with app.test_request_context("/"): + result = handler(api, sqlite_session, _make_account(), pipeline=_make_pipeline()) + + assert [item["id"] for item in result["items"]] == ["sys-1"] + assert result["items"][0]["value"] == "sys-value" diff --git a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py index 3f926b9ec95..e4a0058eee3 100644 --- a/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py +++ b/api/tests/unit_tests/controllers/console/datasets/rag_pipeline/test_rag_pipeline_workflow.py @@ -18,7 +18,6 @@ from sqlalchemy import Engine from sqlalchemy.orm import Session, scoped_session, sessionmaker from werkzeug.exceptions import BadRequest, Forbidden, HTTPException, NotFound -import models.workflow as workflow_models import services import services.rag_pipeline.rag_pipeline as rag_pipeline_service from controllers.console.app.error import DraftWorkflowNotExist, DraftWorkflowNotSync @@ -989,7 +988,6 @@ class TestRagPipelineWorkflowApisCases: database_session = scoped_session(sessionmaker(bind=sqlite_engine, expire_on_commit=False)) database = SQLiteDatabase(engine=sqlite_engine, session=database_session) monkeypatch.setattr(workflow_controller, "db", database) - monkeypatch.setattr(workflow_models, "db", database) monkeypatch.setattr(rag_pipeline_service, "db", database) with database_session() as session: diff --git a/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow_draft_variable.py b/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow_draft_variable.py index 578b2b1c03b..8f914c32a0a 100644 --- a/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow_draft_variable.py +++ b/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow_draft_variable.py @@ -241,6 +241,46 @@ def test_variable_patch_returns_persisted_variable_without_committing_when_no_ch assert result["id"] == variable.id +def test_variable_get_returns_decoded_persisted_value(app: Flask, sqlite_session: Session) -> None: + """Given a persisted snippet variable, return its decoded value through the caller's session.""" + variable = _make_node_variable("var-1") + _persist_variables(sqlite_session, variable) + sqlite_session.expunge_all() + api = module.SnippetVariableApi() + handler = unwrap(api.get) + + with app.test_request_context("/"): + result = handler(api, sqlite_session, _make_account(), snippet=_make_snippet(), variable_id="var-1") + + assert result["id"] == "var-1" + assert result["value"] == "value-var-1" + + +def test_variable_patch_persists_and_returns_decoded_updated_value(app: Flask, sqlite_session: Session) -> None: + """Given a persisted snippet variable, update its value and return the decoded result without losing the row.""" + variable = _make_node_variable("var-1") + _persist_variables(sqlite_session, variable) + sqlite_session.expunge_all() + api = module.SnippetVariableApi() + handler = unwrap(api.patch) + + with app.test_request_context("/", method="PATCH", json={"value": "updated-value"}): + result = handler( + api, + module.WorkflowDraftVariableUpdatePayload(value="updated-value"), + sqlite_session, + _make_account(), + snippet=_make_snippet(), + variable_id="var-1", + ) + + assert result["value"] == "updated-value" + sqlite_session.expire_all() + stored = sqlite_session.get(WorkflowDraftVariable, "var-1") + assert stored is not None + assert stored.get_value(session=sqlite_session).value == "updated-value" + + def test_variable_delete_deletes_persisted_variable( app: Flask, sqlite_session: Session, diff --git a/api/tests/unit_tests/controllers/service_api/app/test_workflow.py b/api/tests/unit_tests/controllers/service_api/app/test_workflow.py index fcdc5807b7a..bd74c261c9d 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_workflow.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_workflow.py @@ -160,7 +160,6 @@ def _bind_sqlite_database( """Bind controller- and model-owned database access to the test engine.""" database = SimpleNamespace(engine=sqlite_engine, session=sqlite_session) monkeypatch.setattr(sys.modules["controllers.service_api.app.workflow"], "db", database) - monkeypatch.setattr(sys.modules["models.workflow"], "db", database) def _persist_workflow_log( diff --git a/api/tests/unit_tests/fields/test_workflow_draft_variable_fields.py b/api/tests/unit_tests/fields/test_workflow_draft_variable_fields.py index c7b4d15494f..14c7f4e7286 100644 --- a/api/tests/unit_tests/fields/test_workflow_draft_variable_fields.py +++ b/api/tests/unit_tests/fields/test_workflow_draft_variable_fields.py @@ -4,6 +4,7 @@ from typing import Any, NamedTuple from unittest.mock import patch import pytest +from sqlalchemy.orm import Session from core.workflow.variable_prefixes import CONVERSATION_VARIABLE_NODE_ID, SYSTEM_VARIABLE_NODE_ID from factories.variable_factory import build_segment @@ -13,6 +14,8 @@ from fields.workflow_draft_variable_fields import ( WorkflowDraftVariableResponse, WorkflowDraftVariableWithoutValueResponse, _serialize_full_content, + draft_variable_list_response_source, + draft_variable_response_source, ) from graphon.variables.types import SegmentType from libs.datetime_utils import naive_utc_now @@ -82,7 +85,7 @@ class TestWorkflowDraftVariableFields: with pytest.raises(AssertionError): _serialize_full_content(draft_var) - def test_conversation_variable(self) -> None: + def test_conversation_variable(self, sqlite_session: Session) -> None: conv_var = WorkflowDraftVariable.new_conversation_variable( app_id=_TEST_APP_ID, name="conv_var", value=build_segment(1) ) @@ -108,9 +111,14 @@ class TestWorkflowDraftVariableFields: expected_with_value = expected_without_value.copy() expected_with_value["value"] = 1 expected_with_value["full_content"] = None - assert dump_response(WorkflowDraftVariableResponse, conv_var) == expected_with_value + assert ( + dump_response( + WorkflowDraftVariableResponse, draft_variable_response_source(conv_var, session=sqlite_session) + ) + == expected_with_value + ) - def test_create_sys_variable(self) -> None: + def test_create_sys_variable(self, sqlite_session: Session) -> None: sys_var = WorkflowDraftVariable.new_sys_variable( app_id=_TEST_APP_ID, name="sys_var", @@ -140,9 +148,14 @@ class TestWorkflowDraftVariableFields: expected_with_value = expected_without_value.copy() expected_with_value["value"] = "a" expected_with_value["full_content"] = None - assert dump_response(WorkflowDraftVariableResponse, sys_var) == expected_with_value + assert ( + dump_response( + WorkflowDraftVariableResponse, draft_variable_response_source(sys_var, session=sqlite_session) + ) + == expected_with_value + ) - def test_node_variable(self) -> None: + def test_node_variable(self, sqlite_session: Session) -> None: node_var = WorkflowDraftVariable.new_node_variable( app_id=_TEST_APP_ID, node_id="test_node", @@ -173,9 +186,14 @@ class TestWorkflowDraftVariableFields: expected_with_value = expected_without_value.copy() expected_with_value["value"] = [1, "a"] expected_with_value["full_content"] = None - assert dump_response(WorkflowDraftVariableResponse, node_var) == expected_with_value + assert ( + dump_response( + WorkflowDraftVariableResponse, draft_variable_response_source(node_var, session=sqlite_session) + ) + == expected_with_value + ) - def test_node_variable_with_file(self) -> None: + def test_node_variable_with_file(self, sqlite_session: Session) -> None: node_var = WorkflowDraftVariable.new_node_variable( app_id=_TEST_APP_ID, node_id="test_node", @@ -225,7 +243,12 @@ class TestWorkflowDraftVariableFields: "length": 10, "download_url": "http://example.com/signed-url", } - assert dump_response(WorkflowDraftVariableResponse, node_var) == expected_with_value + assert ( + dump_response( + WorkflowDraftVariableResponse, draft_variable_response_source(node_var, session=sqlite_session) + ) + == expected_with_value + ) class TestWorkflowDraftVariableList: @@ -307,11 +330,42 @@ class TestWorkflowDraftVariableList: ) -def test_workflow_node_variables_fields() -> None: +def test_value_response_rejects_bare_model() -> None: + """A bare model would expose the raw serialized JSON string as the value, so it must fail loudly.""" conv_var = WorkflowDraftVariable.new_conversation_variable( app_id=_TEST_APP_ID, name="conv_var", value=build_segment(1) ) - resp = dump_response(WorkflowDraftVariableListResponse, WorkflowDraftVariableList(variables=[conv_var])) + with pytest.raises(TypeError, match="draft_variable_response_source"): + dump_response(WorkflowDraftVariableResponse, conv_var) + + +def test_value_response_accepts_plain_mapping() -> None: + resp = WorkflowDraftVariableResponse.model_validate( + { + "id": "variable-id", + "type": "conversation", + "name": "conv_var", + "description": "", + "selector": ["conversation", "conv_var"], + "value_type": "number", + "edited": False, + "visible": True, + "is_truncated": False, + "value": 1, + "full_content": None, + } + ) + assert resp.value == 1 + + +def test_workflow_node_variables_fields(sqlite_session: Session) -> None: + conv_var = WorkflowDraftVariable.new_conversation_variable( + app_id=_TEST_APP_ID, name="conv_var", value=build_segment(1) + ) + resp = dump_response( + WorkflowDraftVariableListResponse, + draft_variable_list_response_source(WorkflowDraftVariableList(variables=[conv_var]), session=sqlite_session), + ) assert isinstance(resp, dict) assert len(resp["items"]) == 1 item_dict = resp["items"][0] @@ -319,7 +373,7 @@ def test_workflow_node_variables_fields() -> None: assert item_dict["value"] == 1 -def test_workflow_file_variable_with_signed_url() -> None: +def test_workflow_file_variable_with_signed_url(sqlite_session: Session) -> None: """Test that File type variables include signed URLs in API responses.""" from graphon.file import File, FileTransferMethod, FileType @@ -345,7 +399,10 @@ def test_workflow_file_variable_with_signed_url() -> None: ) # Marshal the variable using the API fields - resp = dump_response(WorkflowDraftVariableListResponse, WorkflowDraftVariableList(variables=[file_var])) + resp = dump_response( + WorkflowDraftVariableListResponse, + draft_variable_list_response_source(WorkflowDraftVariableList(variables=[file_var]), session=sqlite_session), + ) # Verify the response structure assert isinstance(resp, dict) @@ -375,7 +432,7 @@ def test_workflow_file_variable_with_signed_url() -> None: assert "sign=" in remote_url -def test_workflow_file_variable_remote_url() -> None: +def test_workflow_file_variable_remote_url(sqlite_session: Session) -> None: """Test that File type variables with REMOTE_URL transfer method return the remote URL.""" from graphon.file import File, FileTransferMethod, FileType @@ -401,7 +458,10 @@ def test_workflow_file_variable_remote_url() -> None: ) # Marshal the variable using the API fields - resp = dump_response(WorkflowDraftVariableListResponse, WorkflowDraftVariableList(variables=[file_var])) + resp = dump_response( + WorkflowDraftVariableListResponse, + draft_variable_list_response_source(WorkflowDraftVariableList(variables=[file_var]), session=sqlite_session), + ) # Verify the response structure assert isinstance(resp, dict) diff --git a/api/tests/unit_tests/models/test_workflow.py b/api/tests/unit_tests/models/test_workflow.py index 32195d9fc03..678c4c7594e 100644 --- a/api/tests/unit_tests/models/test_workflow.py +++ b/api/tests/unit_tests/models/test_workflow.py @@ -36,7 +36,7 @@ from models import Workflow as _workflow_feature_compatibility_Workflow from models.account import Account from models.base import Base, TypeBase from models.enums import CreatorUserRole, EndUserType, ExecutionOffLoadType, WorkflowRunTriggeredFrom -from models.model import EndUser, UploadFile +from models.model import App, AppMode, EndUser, UploadFile from models.tools import WorkflowToolProvider from models.workflow import ( Workflow as _workflow_Workflow, @@ -54,6 +54,7 @@ from models.workflow import ( WorkflowRun, WorkflowType, WorkflowVersionCounter, + _resolve_workflow_app_tenant_id, is_system_variable_editable, ) from models.workflow import WorkflowType as ModelWorkflowType @@ -379,7 +380,7 @@ class TestWorkflowCases: assert is_system_variable_editable("invalid_or_new_system_variable") is False class TestWorkflowDraftVariableGetValue: - def test_get_value_by_case(self) -> None: + def test_get_value_by_case(self, sqlite_session: Session) -> None: @dataclasses.dataclass class TestCase: name: str @@ -449,9 +450,9 @@ class TestWorkflowCases: fail_msg = f"test case {c.name} failed, index={idx}" draft_var = WorkflowDraftVariable() draft_var.set_value(c.value) - assert c.value == draft_var.get_value(), fail_msg + assert c.value == draft_var.get_value(session=sqlite_session), fail_msg - def test_file_variable_preserves_all_fields(self) -> None: + def test_file_variable_preserves_all_fields(self, sqlite_session: Session) -> None: """Test that File type variables preserve all fields during encoding/decoding.""" # Create a File with specific field values @@ -473,7 +474,7 @@ class TestWorkflowCases: draft_var.set_value(file_segment) # Retrieve the value and verify all fields are preserved - retrieved_segment = draft_var.get_value() + retrieved_segment = draft_var.get_value(session=sqlite_session) retrieved_file = retrieved_segment.value # Verify all important fields are preserved @@ -490,7 +491,7 @@ class TestWorkflowCases: # Verify the segments have the same type and the important fields match assert file_segment.value_type == retrieved_segment.value_type - def test_file_variable_rebuilds_storage_backed_payloads_with_app_tenant(self) -> None: + def test_file_variable_rebuilds_storage_backed_payloads_with_app_tenant(self, sqlite_session: Session) -> None: persisted_file = File( file_id="test_file_id", file_type=FileType.DOCUMENT, @@ -522,17 +523,58 @@ class TestWorkflowCases: mock.patch("models.workflow._resolve_workflow_app_tenant_id", return_value="tenant-1"), mock.patch("models.workflow.build_file_from_stored_mapping", return_value=rebuilt_file) as rebuild_file, ): - retrieved_segment = draft_var.get_value() + retrieved_segment = draft_var.get_value(session=sqlite_session) assert retrieved_segment.value == rebuilt_file rebuild_file.assert_called_once() assert rebuild_file.call_args.kwargs["tenant_id"] == "tenant-1" - def test_get_and_set_value(self) -> None: + def test_array_file_variable_rebuilds_each_stored_payload_with_app_tenant( + self, sqlite_session: Session + ) -> None: + files = [ + File( + file_id=f"test_file_id_{index}", + file_type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.LOCAL_FILE, + reference=build_file_reference(record_id=f"upload-{index}", storage_key="legacy-storage-key"), + filename=f"test_{index}.txt", + extension=".txt", + mime_type="text/plain", + size=12, + ) + for index in range(2) + ] + rebuilt_files = [ + file.model_copy(update={"storage_key": "canonical-storage-key"}, deep=True) for file in files + ] + encoded = WorkflowDraftVariable( + app_id="app-1", + ) + encoded.set_value(build_segment(files)) + # A fresh instance holding only the persisted columns, so get_value decodes from scratch. + draft_var = WorkflowDraftVariable( + app_id="app-1", + ) + draft_var.value = encoded.value + draft_var.value_type = encoded.value_type + + with ( + mock.patch("models.workflow._resolve_workflow_app_tenant_id", return_value="tenant-1") as resolve, + mock.patch("models.workflow.build_file_from_stored_mapping", side_effect=rebuilt_files) as rebuild_file, + ): + retrieved_segment = draft_var.get_value(session=sqlite_session) + + assert retrieved_segment.value == rebuilt_files + resolve.assert_called_once_with(sqlite_session, "app-1") + assert rebuild_file.call_count == 2 + assert all(call.kwargs["tenant_id"] == "tenant-1" for call in rebuild_file.call_args_list) + + def test_get_and_set_value(self, sqlite_session: Session) -> None: draft_var = WorkflowDraftVariable() int_var = IntegerSegment(value=1) draft_var.set_value(int_var) - value = draft_var.get_value() + value = draft_var.get_value(session=sqlite_session) assert value == int_var @@ -2177,3 +2219,22 @@ class TestWorkflowTypeBaseCases: relationship = inspect(WorkflowNodeExecutionModel).relationships["offload_data"] assert relationship.mapper.class_ is WorkflowNodeExecutionOffload assert relationship.back_populates == "execution" + + +@pytest.mark.parametrize("sqlite_session", [(App,)], indirect=True) +def test_resolve_workflow_app_tenant_id_reads_through_caller_session(sqlite_session: Session) -> None: + app = App( + id="app-1", + tenant_id="tenant-1", + name="Workflow", + mode=AppMode.WORKFLOW, + enable_site=True, + enable_api=True, + ) + sqlite_session.add(app) + sqlite_session.flush() + + assert _resolve_workflow_app_tenant_id(sqlite_session, "app-1") == "tenant-1" + + with pytest.raises(ValueError, match="Unable to resolve tenant_id"): + _resolve_workflow_app_tenant_id(sqlite_session, "missing-app") diff --git a/api/tests/unit_tests/services/test_workflow_draft_variable_service.py b/api/tests/unit_tests/services/test_workflow_draft_variable_service.py index 9a3495c8d31..42cf9b82906 100644 --- a/api/tests/unit_tests/services/test_workflow_draft_variable_service.py +++ b/api/tests/unit_tests/services/test_workflow_draft_variable_service.py @@ -27,7 +27,7 @@ from graphon.enums import BuiltinNodeTypes as _workflow_draft_variable_service_B from graphon.enums import WorkflowNodeExecutionStatus from graphon.file import File, FileTransferMethod, FileType from graphon.nodes import BuiltinNodeTypes as _workflow_draft_variable_service_sqlite_BuiltinNodeTypes -from graphon.variables.segments import StringSegment +from graphon.variables.segments import ArrayFileSegment, FileSegment, StringSegment from graphon.variables.types import SegmentType from graphon.variables.variables import StringVariable from libs import datetime_utils @@ -216,6 +216,70 @@ class TestDraftVarLoaderSimpleCases: with pytest.raises(AssertionError): draft_var_loader._load_offloaded_variable(draft_var) + @pytest.mark.parametrize( + "sqlite_session", + [(WorkflowDraftVariable, WorkflowDraftVariableFile, UploadFile)], + indirect=True, + ) + def test_load_variables_collects_file_values_for_storage_key_loading( + self, + draft_var_loader: DraftVarLoader, + sqlite_session: Session, + ): + """File and array-file variables surface their File objects to the storage-key loader.""" + + def _make_file(index: int) -> File: + return File( + file_id=f"file-{index}", + file_type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.LOCAL_FILE, + reference=build_file_reference(record_id=f"upload-{index}", storage_key="legacy-storage-key"), + filename=f"test_{index}.txt", + extension=".txt", + mime_type="text/plain", + size=12, + ) + + single_file = _make_file(0) + array_files = [_make_file(1), _make_file(2)] + file_var = WorkflowDraftVariable.new_node_variable( + app_id="test-app-id", + user_id="test-user-id", + node_id="node1", + name="file_var", + value=FileSegment(value=single_file), + node_execution_id="execution-node1", + ) + array_var = WorkflowDraftVariable.new_node_variable( + app_id="test-app-id", + user_id="test-user-id", + node_id="node1", + name="array_var", + value=ArrayFileSegment(value=array_files), + node_execution_id="execution-node1", + ) + sqlite_session.add_all([file_var, array_var]) + sqlite_session.commit() + + rebuilt = { + file.id: file.model_copy(update={"storage_key": "canonical"}, deep=True) + for file in [single_file, *array_files] + } + + with ( + patch("models.workflow._resolve_workflow_app_tenant_id", return_value="test-tenant-id"), + patch( + "models.workflow.build_file_from_stored_mapping", + side_effect=lambda **kwargs: rebuilt[kwargs["file_mapping"]["id"]], + ), + patch("services.workflow_draft_variable_service.StorageKeyLoader") as storage_key_loader_cls, + ): + result = draft_var_loader.load_variables([["node1", "file_var"], ["node1", "array_var"]]) + + loaded_files = storage_key_loader_cls.return_value.load_storage_keys.call_args.args[0] + assert {file.id for file in loaded_files} == {"file-0", "file-1", "file-2"} + assert {variable.name for variable in result} == {"file_var", "array_var"} + def test_load_variables_empty_selectors_unit(self, draft_var_loader): """Test load_variables returns empty list for empty selectors.""" result = draft_var_loader.load_variables([]) @@ -616,7 +680,7 @@ class TestWorkflowDraftVariableServiceCases: draft_vars = saver._build_variables_from_start_mapping({"sys.files": [raw_file]}) sys_var = draft_vars[0] - assert sys_var.get_value().value[0] == rebuilt_file + assert sys_var.get_value(session=sqlite_session).value[0] == rebuilt_file rebuild_file.assert_called_once_with(file_mapping=raw_file, tenant_id="tenant-1") @pytest.fixture @@ -1307,7 +1371,7 @@ class TestWorkflowDraftVariableServiceSqliteCases: assert node_var is not None assert node_var.id == self._node1_str_var_id assert node_var.name == "str_var" - assert node_var.get_value() == build_segment("str_value") + assert node_var.get_value(session=self._session) == build_segment("str_value") def test_get_system_variable(self) -> None: srv = self._get_test_srv() @@ -1315,7 +1379,7 @@ class TestWorkflowDraftVariableServiceSqliteCases: assert sys_var is not None assert sys_var.id == self._sys_var_id assert sys_var.name == "sys_var" - assert sys_var.get_value() == build_segment("sys_value") + assert sys_var.get_value(session=self._session) == build_segment("sys_value") def test_get_conversation_variable(self) -> None: srv = self._get_test_srv() @@ -1323,7 +1387,7 @@ class TestWorkflowDraftVariableServiceSqliteCases: assert conv_var is not None assert conv_var.id == self._conv_var_id assert conv_var.name == "conv_var" - assert conv_var.get_value() == build_segment("conv_value") + assert conv_var.get_value(session=self._session) == build_segment("conv_value") def test_delete_node_variables(self) -> None: srv = self._get_test_srv() diff --git a/api/tests/unit_tests/services/test_workflow_node_execution_trace_service.py b/api/tests/unit_tests/services/test_workflow_node_execution_trace_service.py index 42778758e02..13b6b64b480 100644 --- a/api/tests/unit_tests/services/test_workflow_node_execution_trace_service.py +++ b/api/tests/unit_tests/services/test_workflow_node_execution_trace_service.py @@ -2,11 +2,9 @@ from __future__ import annotations import json from datetime import UTC, datetime -from types import SimpleNamespace from typing import cast from unittest.mock import MagicMock -import pytest from sqlalchemy.orm import Session, sessionmaker from domain.workflow.retry_history import RETRY_HISTORY_PROCESS_DATA_KEY, WorkflowNodeRetryAttempt @@ -18,13 +16,6 @@ from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTri from repositories.api_workflow_node_execution_repository import DifyAPIWorkflowNodeExecutionRepository from services.workflow_node_execution_trace_service import assemble_workflow_node_execution_traces -pytestmark = pytest.mark.parametrize("sqlite_session", [(Account,)], indirect=True) - - -@pytest.fixture(autouse=True) -def _bind_model_database(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None: - monkeypatch.setattr("models.workflow.db", SimpleNamespace(session=sqlite_session)) - def _retry_attempt(retry_index: int, **overrides: object) -> dict[str, object]: values = {