refactor: port WorkflowDraftVariableFile (#30923)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

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

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:

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": "",

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
 

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)