Don't Check Partitions for Non-Partitioned Queues (#531)

This fixes an issue where a workflow is stuck when moved from a
partitioned queue to the non-partitioned internal queue (for example
when resuming a workflow).

Author: kraftp

dbos/_sys_db.py

@@ -1862,10 +1862,6 @@ def start_queued_workflows(
                     sa.select(sa.func.count())
                     .select_from(SystemSchema.workflow_status)
                     .where(SystemSchema.workflow_status.c.queue_name == queue.name)
-                    .where(
-                        SystemSchema.workflow_status.c.queue_partition_key
-                        == queue_partition_key
-                    )
                     .where(
                         SystemSchema.workflow_status.c.status
                         != WorkflowStatusString.ENQUEUED.value
@@ -1875,6 +1871,11 @@ def start_queued_workflows(
                         > start_time_ms - limiter_period_ms
                     )
                 )
+                if queue_partition_key is not None:
+                    query = query.where(
+                        SystemSchema.workflow_status.c.queue_partition_key
+                        == queue_partition_key
+                    )
                 num_recent_queries = c.execute(query).fetchone()[0]  # type: ignore
                 if num_recent_queries >= queue.limiter["limit"]:
                     return []
@@ -1890,16 +1891,17 @@ def start_queued_workflows(
                     )
                     .select_from(SystemSchema.workflow_status)
                     .where(SystemSchema.workflow_status.c.queue_name == queue.name)
-                    .where(
-                        SystemSchema.workflow_status.c.queue_partition_key
-                        == queue_partition_key
-                    )
                     .where(
                         SystemSchema.workflow_status.c.status
                         == WorkflowStatusString.PENDING.value
                     )
                     .group_by(SystemSchema.workflow_status.c.executor_id)
                 )
+                if queue_partition_key is not None:
+                    pending_tasks_query = pending_tasks_query.where(
+                        SystemSchema.workflow_status.c.queue_partition_key
+                        == queue_partition_key
+                    )
                 pending_workflows = c.execute(pending_tasks_query).fetchall()
                 pending_workflows_dict = {row[0]: row[1] for row in pending_workflows}
                 local_pending_workflows = pending_workflows_dict.get(executor_id, 0)
@@ -1935,10 +1937,6 @@ def start_queued_workflows(
                 )
                 .select_from(SystemSchema.workflow_status)
                 .where(SystemSchema.workflow_status.c.queue_name == queue.name)
-                .where(
-                    SystemSchema.workflow_status.c.queue_partition_key
-                    == queue_partition_key
-                )
                 .where(
                     SystemSchema.workflow_status.c.status
                     == WorkflowStatusString.ENQUEUED.value
@@ -1955,6 +1953,11 @@ def start_queued_workflows(
                 # to ensure all processes have a consistent view of the table.
                 .with_for_update(skip_locked=skip_locks, nowait=(not skip_locks))
             )
+            if queue_partition_key is not None:
+                query = query.where(
+                    SystemSchema.workflow_status.c.queue_partition_key
+                    == queue_partition_key
+                )
             if queue.priority_enabled:
                 query = query.order_by(
                     SystemSchema.workflow_status.c.priority.asc(),

tests/test_queue.py

@@ -1000,6 +1000,47 @@ def regular_workflow() -> None:
     assert queue_entries_are_cleaned_up(dbos)
 
 
+def test_resuming_queued_partitioned_workflows(dbos: DBOS) -> None:
+    start_event = threading.Event()
+    blocking_event = threading.Event()
+
+    @DBOS.workflow()
+    def stuck_workflow() -> None:
+        start_event.set()
+        blocking_event.wait()
+
+    @DBOS.workflow()
+    def regular_workflow() -> None:
+        return
+
+    # Enqueue a blocked workflow and two regular workflows on a queue with concurrency 1
+    queue = Queue("test_queue", concurrency=1, partition_queue=True)
+    wfid = str(uuid.uuid4())
+    with SetEnqueueOptions(queue_partition_key="key"):
+        blocked_handle = queue.enqueue(stuck_workflow)
+        with SetWorkflowID(wfid):
+            regular_handle_1 = queue.enqueue(regular_workflow)
+        regular_handle_2 = queue.enqueue(regular_workflow)
+
+    # Verify that the blocked workflow starts and is PENDING while the regular workflows remain ENQUEUED.
+    start_event.wait()
+    assert blocked_handle.get_status().status == WorkflowStatusString.PENDING.value
+    assert regular_handle_1.get_status().status == WorkflowStatusString.ENQUEUED.value
+    assert regular_handle_2.get_status().status == WorkflowStatusString.ENQUEUED.value
+
+    # Resume a regular workflow. Verify it completes.
+    dbos.resume_workflow(wfid)
+    assert regular_handle_1.get_result() == None
+
+    # Complete the blocked workflow. Verify the second regular workflow also completes.
+    blocking_event.set()
+    assert blocked_handle.get_result() == None
+    assert regular_handle_2.get_result() == None
+
+    # Verify all queue entries eventually get cleaned up.
+    assert queue_entries_are_cleaned_up(dbos)
+
+
 def test_dlq_enqueued_workflows(dbos: DBOS) -> None:
     start_event = threading.Event()
     blocking_event = threading.Event()