diff --git a/api/models/workflow.py b/api/models/workflow.py index cb1723440b..7936c06a5a 100644 --- a/api/models/workflow.py +++ b/api/models/workflow.py @@ -1568,12 +1568,14 @@ class WorkflowDraftVariable(Base): ), ) - # Relationship to WorkflowDraftVariableFile + # WorkflowDraftVariableFile uses TypeBase while WorkflowDraftVariable uses Base, so the relationship + # must resolve the class object lazily instead of relying on string lookup across registries. variable_file: Mapped[Optional["WorkflowDraftVariableFile"]] = orm.relationship( + lambda: WorkflowDraftVariableFile, foreign_keys=[file_id], lazy="raise", uselist=False, - primaryjoin="WorkflowDraftVariableFile.id == WorkflowDraftVariable.file_id", + primaryjoin=lambda: orm.foreign(WorkflowDraftVariable.file_id) == WorkflowDraftVariableFile.id, ) # Cache for deserialized value @@ -1892,7 +1894,7 @@ class WorkflowDraftVariable(Base): return self.last_edited_at is not None -class WorkflowDraftVariableFile(Base): +class WorkflowDraftVariableFile(TypeBase): """Stores metadata about files associated with large workflow draft variables. This model acts as an intermediary between WorkflowDraftVariable and UploadFile, @@ -1906,18 +1908,7 @@ class WorkflowDraftVariableFile(Base): __tablename__ = "workflow_draft_variable_files" # Primary key - id: Mapped[str] = mapped_column( - StringUUID, - primary_key=True, - default=lambda: str(uuidv7()), - ) - - created_at: Mapped[datetime] = mapped_column( - DateTime, - nullable=False, - default=naive_utc_now, - server_default=func.current_timestamp(), - ) + id: Mapped[str] = mapped_column(StringUUID, primary_key=True, default_factory=lambda: str(uuidv7()), init=False) tenant_id: Mapped[str] = mapped_column( StringUUID, @@ -1969,15 +1960,23 @@ class WorkflowDraftVariableFile(Base): nullable=False, ) - # Relationship to UploadFile + # Rows are created with `upload_file_id`; callers should load this relationship explicitly when needed. upload_file: Mapped["UploadFile"] = orm.relationship( UploadFile, foreign_keys=[upload_file_id], lazy="raise", + init=False, uselist=False, primaryjoin=lambda: orm.foreign(WorkflowDraftVariableFile.upload_file_id) == UploadFile.id, ) + created_at: Mapped[datetime] = mapped_column( + DateTime, + nullable=False, + default_factory=naive_utc_now, + server_default=func.current_timestamp(), + ) + def is_system_variable_editable(name: str) -> bool: return name in _EDITABLE_SYSTEM_VARIABLE diff --git a/api/services/workflow_draft_variable_service.py b/api/services/workflow_draft_variable_service.py index 96f936ff9b..a55448e352 100644 --- a/api/services/workflow_draft_variable_service.py +++ b/api/services/workflow_draft_variable_service.py @@ -1083,10 +1083,9 @@ class DraftVariableSaver: mimetype=content_type, user=self._user, ) - + assert self._user.current_tenant_id # Create WorkflowDraftVariableFile record variable_file = WorkflowDraftVariableFile( - id=uuidv7(), upload_file_id=upload_file.id, size=original_size, length=original_length, @@ -1095,6 +1094,7 @@ class DraftVariableSaver: tenant_id=self._user.current_tenant_id, user_id=self._user.id, ) + variable_file.id = str(uuidv7()) engine = bind = self._session.get_bind() assert isinstance(engine, Engine) with sessionmaker(bind=engine, expire_on_commit=False).begin() as session: diff --git a/api/tests/unit_tests/controllers/console/app/workflow_draft_variables_test.py b/api/tests/unit_tests/controllers/console/app/workflow_draft_variables_test.py index 22b80b748e..62fa82e339 100644 --- a/api/tests/unit_tests/controllers/console/app/workflow_draft_variables_test.py +++ b/api/tests/unit_tests/controllers/console/app/workflow_draft_variables_test.py @@ -1,7 +1,7 @@ import uuid from collections import OrderedDict from typing import Any, NamedTuple -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest from flask_restx import marshal @@ -29,15 +29,18 @@ class TestWorkflowDraftVariableFields: def test_serialize_full_content(self): """Test that _serialize_full_content uses pre-loaded relationships.""" # Create mock objects with relationships pre-loaded - mock_variable_file = MagicMock(spec=WorkflowDraftVariableFile) - mock_variable_file.size = 100000 - mock_variable_file.length = 50 - mock_variable_file.value_type = SegmentType.OBJECT - mock_variable_file.upload_file_id = "test-upload-file-id" - - mock_variable = MagicMock(spec=WorkflowDraftVariable) - mock_variable.file_id = "test-file-id" - mock_variable.variable_file = mock_variable_file + mock_variable = WorkflowDraftVariable( + file_id="test-file-id", + variable_file=WorkflowDraftVariableFile( + size=100000, + length=50, + value_type=SegmentType.OBJECT, + upload_file_id="test-upload-file-id", + tenant_id=str(uuid.uuid4()), + app_id=str(uuid.uuid4()), + user_id=str(uuid.uuid4()), + ), + ) # Mock the file helpers with patch("controllers.console.app.workflow_draft_variable.file_helpers", autospec=True) as mock_file_helpers: @@ -84,7 +87,7 @@ class TestWorkflowDraftVariableFields: expected_without_value: OrderedDict[str, Any] = OrderedDict( { - "id": str(conv_var.id), + "id": conv_var.id, "type": conv_var.get_variable_type().value, "name": "conv_var", "description": "", @@ -117,7 +120,7 @@ class TestWorkflowDraftVariableFields: expected_without_value = OrderedDict( { - "id": str(sys_var.id), + "id": sys_var.id, "type": sys_var.get_variable_type().value, "name": "sys_var", "description": "", @@ -149,7 +152,7 @@ class TestWorkflowDraftVariableFields: expected_without_value: OrderedDict[str, Any] = OrderedDict( { - "id": str(node_var.id), + "id": node_var.id, "type": node_var.get_variable_type().value, "name": "node_var", "description": "", @@ -180,19 +183,22 @@ class TestWorkflowDraftVariableFields: node_var.id = str(uuid.uuid4()) node_var.last_edited_at = naive_utc_now() variable_file = WorkflowDraftVariableFile( - id=str(uuidv7()), upload_file_id=str(uuid.uuid4()), size=1024, length=10, value_type=SegmentType.ARRAY_STRING, + tenant_id=str(uuidv7()), + app_id=str(uuidv7()), + user_id=str(uuidv7()), ) + variable_file.id = str(uuidv7()) node_var.variable_file = variable_file node_var.file_id = variable_file.id expected_without_value: OrderedDict[str, Any] = OrderedDict( { - "id": str(node_var.id), - "type": node_var.get_variable_type().value, + "id": node_var.id, + "type": node_var.get_variable_type(), "name": "node_var", "description": "", "selector": ["test_node", "node_var"], @@ -235,7 +241,7 @@ class TestWorkflowDraftVariableList: node_var.id = str(uuid.uuid4()) node_var_dict = OrderedDict( { - "id": str(node_var.id), + "id": node_var.id, "type": node_var.get_variable_type().value, "name": "test_var", "description": "", diff --git a/api/tests/unit_tests/services/workflow/test_draft_var_loader_simple.py b/api/tests/unit_tests/services/workflow/test_draft_var_loader_simple.py index 497c26a9b3..fb5cf7bc6e 100644 --- a/api/tests/unit_tests/services/workflow/test_draft_var_loader_simple.py +++ b/api/tests/unit_tests/services/workflow/test_draft_var_loader_simple.py @@ -33,42 +33,6 @@ class TestDraftVarLoaderSimple: fallback_variables=[], ) - def test_load_offloaded_variable_string_type_unit(self, draft_var_loader): - """Test _load_offloaded_variable with string type - isolated unit test.""" - # Create mock objects - upload_file = Mock(spec=UploadFile) - upload_file.key = "storage/key/test.txt" - - variable_file = Mock(spec=WorkflowDraftVariableFile) - variable_file.value_type = SegmentType.STRING - variable_file.upload_file = upload_file - - draft_var = Mock(spec=WorkflowDraftVariable) - draft_var.id = "draft-var-id" - draft_var.node_id = "test-node-id" - draft_var.name = "test_variable" - draft_var.description = "test description" - draft_var.get_selector.return_value = ["test-node-id", "test_variable"] - draft_var.variable_file = variable_file - - test_content = "This is the full string content" - - with patch("services.workflow_draft_variable_service.storage") as mock_storage: - mock_storage.load.return_value = test_content.encode() - - # Execute the method - selector_tuple, variable = draft_var_loader._load_offloaded_variable(draft_var) - - # Verify results - assert selector_tuple == ("test-node-id", "test_variable") - assert variable.id == "draft-var-id" - assert variable.name == "test_variable" - assert variable.description == "test description" - assert variable.value == test_content - - # Verify storage was called correctly - mock_storage.load.assert_called_once_with("storage/key/test.txt") - def test_load_offloaded_variable_object_type_unit(self, draft_var_loader): """Test _load_offloaded_variable with object type - isolated unit test.""" # Create mock objects @@ -139,47 +103,6 @@ class TestDraftVarLoaderSimple: result = draft_var_loader._selector_to_tuple(selector) assert result == ("node_id", "var_name") - def test_load_offloaded_variable_number_type_unit(self, draft_var_loader): - """Test _load_offloaded_variable with number type - isolated unit test.""" - # Create mock objects - upload_file = Mock(spec=UploadFile) - upload_file.key = "storage/key/test_number.json" - - variable_file = Mock(spec=WorkflowDraftVariableFile) - variable_file.value_type = SegmentType.NUMBER - variable_file.upload_file = upload_file - - draft_var = Mock(spec=WorkflowDraftVariable) - draft_var.id = "draft-var-id" - draft_var.node_id = "test-node-id" - draft_var.name = "test_number" - draft_var.description = "test number description" - draft_var.get_selector.return_value = ["test-node-id", "test_number"] - draft_var.variable_file = variable_file - - test_number = 123.45 - test_json_content = json.dumps(test_number) - - with patch("services.workflow_draft_variable_service.storage") as mock_storage: - mock_storage.load.return_value = test_json_content.encode() - from graphon.variables.segments import FloatSegment - - mock_segment = FloatSegment(value=test_number) - draft_var.build_segment_from_serialized_value.return_value = mock_segment - - # Execute the method - selector_tuple, variable = draft_var_loader._load_offloaded_variable(draft_var) - - # Verify results - assert selector_tuple == ("test-node-id", "test_number") - assert variable.id == "draft-var-id" - assert variable.name == "test_number" - assert variable.description == "test number description" - - # Verify method calls - mock_storage.load.assert_called_once_with("storage/key/test_number.json") - draft_var.build_segment_from_serialized_value.assert_called_once_with(SegmentType.NUMBER, test_number) - def test_load_offloaded_variable_array_type_unit(self, draft_var_loader): """Test _load_offloaded_variable with array type - isolated unit test.""" # Create mock objects @@ -229,12 +152,13 @@ class TestDraftVarLoaderSimple: variable_file.value_type = SegmentType.FILE variable_file.upload_file = upload_file - draft_var = WorkflowDraftVariable() - draft_var.id = "draft-var-id" - draft_var.app_id = "app-1" - draft_var.node_id = "test-node-id" - draft_var.name = "test_file" - draft_var.description = "test file description" + draft_var = WorkflowDraftVariable( + id="draft-var-id", + app_id="app-1", + node_id="test-node-id", + name="test_file", + description="test file description", + ) draft_var._set_selector(["test-node-id", "test_file"]) draft_var.variable_file = variable_file diff --git a/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py b/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py index b14d767568..663eec6a06 100644 --- a/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py +++ b/api/tests/unit_tests/services/workflow/test_workflow_draft_variable_service.py @@ -200,7 +200,7 @@ class TestDraftVariableSaver: user=mock_user, ) - def test_draft_saver_with_small_variables(self, draft_saver, mock_session): + def test_draft_saver_with_small_variables(self, draft_saver: DraftVariableSaver, mock_session): with patch( "services.workflow_draft_variable_service.DraftVariableSaver._try_offload_large_variable", autospec=True ) as _mock_try_offload: @@ -212,18 +212,21 @@ class TestDraftVariableSaver: assert draft_var.file_id is None _mock_try_offload.return_value = None - def test_draft_saver_with_large_variables(self, draft_saver, mock_session): + def test_draft_saver_with_large_variables(self, draft_saver: DraftVariableSaver, mock_session): with patch( "services.workflow_draft_variable_service.DraftVariableSaver._try_offload_large_variable", autospec=True ) as _mock_try_offload: mock_segment = StringSegment(value="small value") mock_draft_var_file = WorkflowDraftVariableFile( - id=str(uuidv7()), + tenant_id=str(uuidv7()), + app_id=str(uuidv7()), + user_id=str(uuidv7()), size=1024, length=10, value_type=SegmentType.ARRAY_STRING, - upload_file_id=str(uuid.uuid4()), + upload_file_id=str(uuidv7()), ) + mock_draft_var_file.id = str(uuidv7()) _mock_try_offload.return_value = mock_segment, mock_draft_var_file draft_var = draft_saver._create_draft_variable(name="small_var", value=mock_segment, visible=True)