Fix: Resolve workflow_node_execution primary key conflicts with UUID v7 (#24643)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2025-09-02 14:18:29 +08:00
parent 067b0d07c4
commit a32dde5428
4 changed files with 325 additions and 20 deletions
--- a/api/tests/unit_tests/core/repositories/test_workflow_node_execution_conflict_handling.py
+++ b/api/tests/unit_tests/core/repositories/test_workflow_node_execution_conflict_handling.py
@@ -0,0 +1,210 @@
+"""Unit tests for workflow node execution conflict handling."""
+
+from datetime import datetime
+from unittest.mock import MagicMock, Mock
+
+import psycopg2.errors
+import pytest
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.orm import sessionmaker
+
+from core.repositories.sqlalchemy_workflow_node_execution_repository import (
+    SQLAlchemyWorkflowNodeExecutionRepository,
+)
+from core.workflow.entities.workflow_node_execution import (
+    WorkflowNodeExecution,
+    WorkflowNodeExecutionStatus,
+)
+from core.workflow.nodes.enums import NodeType
+from models import Account, WorkflowNodeExecutionTriggeredFrom
+
+
+class TestWorkflowNodeExecutionConflictHandling:
+    """Test cases for handling duplicate key conflicts in workflow node execution."""
+
+    def setup_method(self):
+        """Set up test fixtures."""
+        # Create a mock user with tenant_id
+        self.mock_user = Mock(spec=Account)
+        self.mock_user.id = "test-user-id"
+        self.mock_user.current_tenant_id = "test-tenant-id"
+
+        # Create mock session factory
+        self.mock_session_factory = Mock(spec=sessionmaker)
+
+        # Create repository instance
+        self.repository = SQLAlchemyWorkflowNodeExecutionRepository(
+            session_factory=self.mock_session_factory,
+            user=self.mock_user,
+            app_id="test-app-id",
+            triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
+        )
+
+    def test_save_with_duplicate_key_retries_with_new_uuid(self):
+        """Test that save retries with a new UUID v7 when encountering duplicate key error."""
+        # Create a mock session
+        mock_session = MagicMock()
+        mock_session.__enter__ = Mock(return_value=mock_session)
+        mock_session.__exit__ = Mock(return_value=None)
+        self.mock_session_factory.return_value = mock_session
+
+        # Mock session.get to return None (no existing record)
+        mock_session.get.return_value = None
+
+        # Create IntegrityError for duplicate key with proper psycopg2.errors.UniqueViolation
+        mock_unique_violation = Mock(spec=psycopg2.errors.UniqueViolation)
+        duplicate_error = IntegrityError(
+            "duplicate key value violates unique constraint",
+            params=None,
+            orig=mock_unique_violation,
+        )
+
+        # First call to session.add raises IntegrityError, second succeeds
+        mock_session.add.side_effect = [duplicate_error, None]
+        mock_session.commit.side_effect = [None, None]
+
+        # Create test execution
+        execution = WorkflowNodeExecution(
+            id="original-id",
+            workflow_id="test-workflow-id",
+            workflow_execution_id="test-workflow-execution-id",
+            node_execution_id="test-node-execution-id",
+            node_id="test-node-id",
+            node_type=NodeType.START,
+            title="Test Node",
+            index=1,
+            status=WorkflowNodeExecutionStatus.RUNNING,
+            created_at=datetime.utcnow(),
+        )
+
+        original_id = execution.id
+
+        # Save should succeed after retry
+        self.repository.save(execution)
+
+        # Verify that session.add was called twice (initial attempt + retry)
+        assert mock_session.add.call_count == 2
+
+        # Verify that the ID was changed (new UUID v7 generated)
+        assert execution.id != original_id
+
+    def test_save_with_existing_record_updates_instead_of_insert(self):
+        """Test that save updates existing record instead of inserting duplicate."""
+        # Create a mock session
+        mock_session = MagicMock()
+        mock_session.__enter__ = Mock(return_value=mock_session)
+        mock_session.__exit__ = Mock(return_value=None)
+        self.mock_session_factory.return_value = mock_session
+
+        # Mock existing record
+        mock_existing = MagicMock()
+        mock_session.get.return_value = mock_existing
+        mock_session.commit.return_value = None
+
+        # Create test execution
+        execution = WorkflowNodeExecution(
+            id="existing-id",
+            workflow_id="test-workflow-id",
+            workflow_execution_id="test-workflow-execution-id",
+            node_execution_id="test-node-execution-id",
+            node_id="test-node-id",
+            node_type=NodeType.START,
+            title="Test Node",
+            index=1,
+            status=WorkflowNodeExecutionStatus.SUCCEEDED,
+            created_at=datetime.utcnow(),
+        )
+
+        # Save should update existing record
+        self.repository.save(execution)
+
+        # Verify that session.add was not called (update path)
+        mock_session.add.assert_not_called()
+
+        # Verify that session.commit was called
+        mock_session.commit.assert_called_once()
+
+    def test_save_exceeds_max_retries_raises_error(self):
+        """Test that save raises error after exceeding max retries."""
+        # Create a mock session
+        mock_session = MagicMock()
+        mock_session.__enter__ = Mock(return_value=mock_session)
+        mock_session.__exit__ = Mock(return_value=None)
+        self.mock_session_factory.return_value = mock_session
+
+        # Mock session.get to return None (no existing record)
+        mock_session.get.return_value = None
+
+        # Create IntegrityError for duplicate key with proper psycopg2.errors.UniqueViolation
+        mock_unique_violation = Mock(spec=psycopg2.errors.UniqueViolation)
+        duplicate_error = IntegrityError(
+            "duplicate key value violates unique constraint",
+            params=None,
+            orig=mock_unique_violation,
+        )
+
+        # All attempts fail with duplicate error
+        mock_session.add.side_effect = duplicate_error
+
+        # Create test execution
+        execution = WorkflowNodeExecution(
+            id="test-id",
+            workflow_id="test-workflow-id",
+            workflow_execution_id="test-workflow-execution-id",
+            node_execution_id="test-node-execution-id",
+            node_id="test-node-id",
+            node_type=NodeType.START,
+            title="Test Node",
+            index=1,
+            status=WorkflowNodeExecutionStatus.RUNNING,
+            created_at=datetime.utcnow(),
+        )
+
+        # Save should raise IntegrityError after max retries
+        with pytest.raises(IntegrityError):
+            self.repository.save(execution)
+
+        # Verify that session.add was called 3 times (max_retries)
+        assert mock_session.add.call_count == 3
+
+    def test_save_non_duplicate_integrity_error_raises_immediately(self):
+        """Test that non-duplicate IntegrityErrors are raised immediately without retry."""
+        # Create a mock session
+        mock_session = MagicMock()
+        mock_session.__enter__ = Mock(return_value=mock_session)
+        mock_session.__exit__ = Mock(return_value=None)
+        self.mock_session_factory.return_value = mock_session
+
+        # Mock session.get to return None (no existing record)
+        mock_session.get.return_value = None
+
+        # Create IntegrityError for non-duplicate constraint
+        other_error = IntegrityError(
+            "null value in column violates not-null constraint",
+            params=None,
+            orig=None,
+        )
+
+        # First call raises non-duplicate error
+        mock_session.add.side_effect = other_error
+
+        # Create test execution
+        execution = WorkflowNodeExecution(
+            id="test-id",
+            workflow_id="test-workflow-id",
+            workflow_execution_id="test-workflow-execution-id",
+            node_execution_id="test-node-execution-id",
+            node_id="test-node-id",
+            node_type=NodeType.START,
+            title="Test Node",
+            index=1,
+            status=WorkflowNodeExecutionStatus.RUNNING,
+            created_at=datetime.utcnow(),
+        )
+
+        # Save should raise error immediately
+        with pytest.raises(IntegrityError):
+            self.repository.save(execution)
+
+        # Verify that session.add was called only once (no retry)
+        assert mock_session.add.call_count == 1
--- a/api/tests/unit_tests/repositories/workflow_node_execution/test_sqlalchemy_repository.py
+++ b/api/tests/unit_tests/repositories/workflow_node_execution/test_sqlalchemy_repository.py
@@ -86,6 +86,8 @@ def test_save(repository, session):
    session_obj, _ = session
    # Create a mock execution
    execution = MagicMock(spec=WorkflowNodeExecutionModel)
+    execution.id = "test-id"
+    execution.node_execution_id = "test-node-execution-id"
    execution.tenant_id = None
    execution.app_id = None
    execution.inputs = None
@@ -95,7 +97,13 @@ def test_save(repository, session):

    # Mock the to_db_model method to return the execution itself
    # This simulates the behavior of setting tenant_id and app_id
-    repository.to_db_model = MagicMock(return_value=execution)
+    db_model = MagicMock(spec=WorkflowNodeExecutionModel)
+    db_model.id = "test-id"
+    db_model.node_execution_id = "test-node-execution-id"
+    repository.to_db_model = MagicMock(return_value=db_model)
+
+    # Mock session.get to return None (no existing record)
+    session_obj.get.return_value = None

    # Call save method
    repository.save(execution)
@@ -103,8 +111,14 @@ def test_save(repository, session):
    # Assert to_db_model was called with the execution
    repository.to_db_model.assert_called_once_with(execution)

-    # Assert session.merge was called (now using merge for both save and update)
-    session_obj.merge.assert_called_once_with(execution)
+    # Assert session.get was called to check for existing record
+    session_obj.get.assert_called_once_with(WorkflowNodeExecutionModel, db_model.id)
+
+    # Assert session.add was called for new record
+    session_obj.add.assert_called_once_with(db_model)
+
+    # Assert session.commit was called
+    session_obj.commit.assert_called_once()


 def test_save_with_existing_tenant_id(repository, session):
@@ -112,6 +126,8 @@ def test_save_with_existing_tenant_id(repository, session):
    session_obj, _ = session
    # Create a mock execution with existing tenant_id
    execution = MagicMock(spec=WorkflowNodeExecutionModel)
+    execution.id = "existing-id"
+    execution.node_execution_id = "existing-node-execution-id"
    execution.tenant_id = "existing-tenant"
    execution.app_id = None
    execution.inputs = None
@@ -121,20 +137,39 @@ def test_save_with_existing_tenant_id(repository, session):

    # Create a modified execution that will be returned by _to_db_model
    modified_execution = MagicMock(spec=WorkflowNodeExecutionModel)
+    modified_execution.id = "existing-id"
+    modified_execution.node_execution_id = "existing-node-execution-id"
    modified_execution.tenant_id = "existing-tenant"  # Tenant ID should not change
    modified_execution.app_id = repository._app_id  # App ID should be set
+    # Create a dictionary to simulate __dict__ for updating attributes
+    modified_execution.__dict__ = {
+        "id": "existing-id",
+        "node_execution_id": "existing-node-execution-id",
+        "tenant_id": "existing-tenant",
+        "app_id": repository._app_id,
+    }

    # Mock the to_db_model method to return the modified execution
    repository.to_db_model = MagicMock(return_value=modified_execution)

+    # Mock session.get to return an existing record
+    existing_model = MagicMock(spec=WorkflowNodeExecutionModel)
+    session_obj.get.return_value = existing_model
+
    # Call save method
    repository.save(execution)

    # Assert to_db_model was called with the execution
    repository.to_db_model.assert_called_once_with(execution)

-    # Assert session.merge was called with the modified execution (now using merge for both save and update)
-    session_obj.merge.assert_called_once_with(modified_execution)
+    # Assert session.get was called to check for existing record
+    session_obj.get.assert_called_once_with(WorkflowNodeExecutionModel, modified_execution.id)
+
+    # Assert session.add was NOT called since we're updating existing
+    session_obj.add.assert_not_called()
+
+    # Assert session.commit was called
+    session_obj.commit.assert_called_once()


 def test_get_by_workflow_run(repository, session, mocker: MockerFixture):