Partitioned Queues (#458)

This PR allows you to create "partitioned queues." 

Partitioned queues let you distribute work across dynamically created
queue partitions. When you enqueue a workflow on a partitioned queue,
you must supply a "queue partition key." Partitioned queues dequeue
workflows and apply flow control limits for individual partitions, not
for the entire queue. Essentially, you can think of each partition as a
"virtual queue" you can create dynamically by just enqueueing a workflow
with a new partition key.

For example, let's say you want to run only once task at once per user.
You can do this with a partitioned queue where the partition key is the
user ID:

```python
queue = Queue("queue", partition_queue=True, concurrency=1)

@DBOS.workflow()
def process_task(task: Task):
  ...

def on_user_task_submission(user_id: str, task: Task):
    with SetEnqueueOptions(queue_partition_key=user_id):
        queue.enqueue(process_task, task)
```

Author: kraftp

dbos/_client.py

@@ -62,6 +62,7 @@ class EnqueueOptions(_EnqueueOptionsRequired, total=False):
     deduplication_id: str
     priority: int
     max_recovery_attempts: int
+    queue_partition_key: str
 
 
 def validate_enqueue_options(options: EnqueueOptions) -> None:
@@ -185,6 +186,7 @@ def _enqueue(self, options: EnqueueOptions, *args: Any, **kwargs: Any) -> str:
             "deduplication_id": options.get("deduplication_id"),
             "priority": options.get("priority"),
             "app_version": options.get("app_version"),
+            "queue_partition_key": options.get("queue_partition_key"),
         }
 
         inputs: WorkflowInputs = {
@@ -221,6 +223,7 @@ def _enqueue(self, options: EnqueueOptions, *args: Any, **kwargs: Any) -> str:
                 else 0
             ),
             "inputs": self._serializer.serialize(inputs),
+            "queue_partition_key": enqueue_options_internal["queue_partition_key"],
         }
 
         self._sys_db.init_workflow(
@@ -286,6 +289,7 @@ def send(
             "deduplication_id": None,
             "priority": 0,
             "inputs": self._serializer.serialize({"args": (), "kwargs": {}}),
+            "queue_partition_key": None,
         }
         with self._sys_db.engine.begin() as conn:
             self._sys_db._insert_workflow_status(

dbos/_context.py

@@ -120,6 +120,8 @@ def __init__(self) -> None:
         self.deduplication_id: Optional[str] = None
         # A user-specified priority for the enqueuing workflow.
         self.priority: Optional[int] = None
+        # If the workflow is enqueued on a partitioned queue, its partition key
+        self.queue_partition_key: Optional[str] = None
 
     def create_child(self) -> DBOSContext:
         rv = DBOSContext()
@@ -479,6 +481,7 @@ def __init__(
         deduplication_id: Optional[str] = None,
         priority: Optional[int] = None,
         app_version: Optional[str] = None,
+        queue_partition_key: Optional[str] = None,
     ) -> None:
         self.created_ctx = False
         self.deduplication_id: Optional[str] = deduplication_id
@@ -491,6 +494,8 @@ def __init__(
         self.saved_priority: Optional[int] = None
         self.app_version: Optional[str] = app_version
         self.saved_app_version: Optional[str] = None
+        self.queue_partition_key = queue_partition_key
+        self.saved_queue_partition_key: Optional[str] = None
 
     def __enter__(self) -> SetEnqueueOptions:
         # Code to create a basic context
@@ -505,6 +510,8 @@ def __enter__(self) -> SetEnqueueOptions:
         ctx.priority = self.priority
         self.saved_app_version = ctx.app_version
         ctx.app_version = self.app_version
+        self.saved_queue_partition_key = ctx.queue_partition_key
+        ctx.queue_partition_key = self.queue_partition_key
         return self
 
     def __exit__(
@@ -517,6 +524,7 @@ def __exit__(
         curr_ctx.deduplication_id = self.saved_deduplication_id
         curr_ctx.priority = self.saved_priority
         curr_ctx.app_version = self.saved_app_version
+        curr_ctx.queue_partition_key = self.saved_queue_partition_key
         # Code to clean up the basic context if we created it
         if self.created_ctx:
             _clear_local_dbos_context()

dbos/_core.py

@@ -303,6 +303,11 @@ def _init_workflow(
             else 0
         ),
         "inputs": dbos._serializer.serialize(inputs),
+        "queue_partition_key": (
+            enqueue_options["queue_partition_key"]
+            if enqueue_options is not None
+            else None
+        ),
     }
 
     # Synchronously record the status and inputs for workflows
@@ -571,6 +576,9 @@ def start_workflow(
         deduplication_id=local_ctx.deduplication_id if local_ctx is not None else None,
         priority=local_ctx.priority if local_ctx is not None else None,
         app_version=local_ctx.app_version if local_ctx is not None else None,
+        queue_partition_key=(
+            local_ctx.queue_partition_key if local_ctx is not None else None
+        ),
     )
     new_wf_id, new_wf_ctx = _get_new_wf()
 
@@ -664,6 +672,9 @@ async def start_workflow_async(
         deduplication_id=local_ctx.deduplication_id if local_ctx is not None else None,
         priority=local_ctx.priority if local_ctx is not None else None,
         app_version=local_ctx.app_version if local_ctx is not None else None,
+        queue_partition_key=(
+            local_ctx.queue_partition_key if local_ctx is not None else None
+        ),
     )
     new_wf_id, new_wf_ctx = _get_new_wf()
 

dbos/_migration.py

@@ -203,8 +203,14 @@ def get_dbos_migration_one(schema: str) -> str:
 """
 
 
+def get_dbos_migration_two(schema: str) -> str:
+    return f"""
+ALTER TABLE \"{schema}\".workflow_status ADD COLUMN queue_partition_key TEXT;
+"""
+
+
 def get_dbos_migrations(schema: str) -> list[str]:
-    return [get_dbos_migration_one(schema)]
+    return [get_dbos_migration_one(schema), get_dbos_migration_two(schema)]
 
 
 def get_sqlite_timestamp_expr() -> str:
@@ -293,4 +299,8 @@ def get_sqlite_timestamp_expr() -> str:
 );
 """
 
-sqlite_migrations = [sqlite_migration_one]
+sqlite_migration_two = """
+ALTER TABLE workflow_status ADD COLUMN queue_partition_key TEXT;
+"""
+
+sqlite_migrations = [sqlite_migration_one, sqlite_migration_two]

dbos/_queue.py

@@ -43,6 +43,7 @@ def __init__(
         *,  # Disable positional arguments from here on
         worker_concurrency: Optional[int] = None,
         priority_enabled: bool = False,
+        partition_queue: bool = False,
     ) -> None:
         if (
             worker_concurrency is not None
@@ -57,6 +58,7 @@ def __init__(
         self.worker_concurrency = worker_concurrency
         self.limiter = limiter
         self.priority_enabled = priority_enabled
+        self.partition_queue = partition_queue
         from ._dbos import _get_or_create_dbos_registry
 
         registry = _get_or_create_dbos_registry()
@@ -78,6 +80,18 @@ def enqueue(
             raise Exception(
                 f"Priority is not enabled for queue {self.name}. Setting priority will not have any effect."
             )
+        if self.partition_queue and (
+            context is None or context.queue_partition_key is None
+        ):
+            raise Exception(
+                f"A workflow cannot be enqueued on partitioned queue {self.name} without a partition key"
+            )
+        if context and context.queue_partition_key and not self.partition_queue:
+            raise Exception(
+                f"You can only use a partition key on a partition-enabled queue. Key {context.queue_partition_key} was used with non-partitioned queue {self.name}"
+            )
+        if context and context.queue_partition_key and context.deduplication_id:
+            raise Exception("Deduplication is not supported for partitioned queues")
 
         dbos = _get_dbos_instance()
         return start_workflow(dbos, func, self.name, False, *args, **kwargs)
@@ -105,10 +119,21 @@ def queue_thread(stop_event: threading.Event, dbos: "DBOS") -> None:
         queues = dict(dbos._registry.queue_info_map)
         for _, queue in queues.items():
             try:
-                wf_ids = dbos._sys_db.start_queued_workflows(
-                    queue, GlobalParams.executor_id, GlobalParams.app_version
-                )
-                for id in wf_ids:
+                if queue.partition_queue:
+                    dequeued_workflows = []
+                    queue_partition_keys = dbos._sys_db.get_queue_partitions(queue.name)
+                    for key in queue_partition_keys:
+                        dequeued_workflows += dbos._sys_db.start_queued_workflows(
+                            queue,
+                            GlobalParams.executor_id,
+                            GlobalParams.app_version,
+                            key,
+                        )
+                else:
+                    dequeued_workflows = dbos._sys_db.start_queued_workflows(
+                        queue, GlobalParams.executor_id, GlobalParams.app_version, None
+                    )
+                for id in dequeued_workflows:
                     execute_workflow_by_id(dbos, id)
             except OperationalError as e:
                 if isinstance(

dbos/_schemas/system_database.py

@@ -77,6 +77,7 @@ def set_schema(cls, schema_name: Optional[str]) -> None:
         Column("deduplication_id", Text(), nullable=True),
         Column("inputs", Text()),
         Column("priority", Integer(), nullable=False, server_default=text("'0'::int")),
+        Column("queue_partition_key", Text()),
         Index("workflow_status_created_at_index", "created_at"),
         Index("workflow_status_executor_id_index", "executor_id"),
         Index("workflow_status_status_index", "status"),

dbos/_sys_db.py

@@ -152,6 +152,8 @@ class WorkflowStatusInternal(TypedDict):
     priority: int
     # Serialized workflow inputs
     inputs: str
+    # If this workflow is enqueued on a partitioned queue, its partition key
+    queue_partition_key: Optional[str]
 
 
 class EnqueueOptionsInternal(TypedDict):
@@ -161,6 +163,8 @@ class EnqueueOptionsInternal(TypedDict):
     priority: Optional[int]
     # On what version the workflow is enqueued. Current version if not specified.
     app_version: Optional[str]
+    # If the workflow is enqueued on a partitioned queue, its partition key
+    queue_partition_key: Optional[str]
 
 
 class RecordedResult(TypedDict):
@@ -490,6 +494,7 @@ def _insert_workflow_status(
                 deduplication_id=status["deduplication_id"],
                 priority=status["priority"],
                 inputs=status["inputs"],
+                queue_partition_key=status["queue_partition_key"],
             )
             .on_conflict_do_update(
                 index_elements=["workflow_uuid"],
@@ -761,6 +766,7 @@ def get_workflow_status(
                     SystemSchema.workflow_status.c.deduplication_id,
                     SystemSchema.workflow_status.c.priority,
                     SystemSchema.workflow_status.c.inputs,
+                    SystemSchema.workflow_status.c.queue_partition_key,
                 ).where(SystemSchema.workflow_status.c.workflow_uuid == workflow_uuid)
             ).fetchone()
             if row is None:
@@ -788,6 +794,7 @@ def get_workflow_status(
                 "deduplication_id": row[16],
                 "priority": row[17],
                 "inputs": row[18],
+                "queue_partition_key": row[19],
             }
             return status
 
@@ -1714,8 +1721,41 @@ def get_event(
             )
         return value
 
+    @db_retry()
+    def get_queue_partitions(self, queue_name: str) -> List[str]:
+        """
+        Get all unique partition names associated with a queue for ENQUEUED workflows.
+
+        Args:
+            queue_name: The name of the queue to get partitions for
+
+        Returns:
+            A list of unique partition names for the queue
+        """
+        with self.engine.begin() as c:
+            query = (
+                sa.select(SystemSchema.workflow_status.c.queue_partition_key)
+                .distinct()
+                .where(SystemSchema.workflow_status.c.queue_name == queue_name)
+                .where(
+                    SystemSchema.workflow_status.c.status.in_(
+                        [
+                            WorkflowStatusString.ENQUEUED.value,
+                        ]
+                    )
+                )
+                .where(SystemSchema.workflow_status.c.queue_partition_key.isnot(None))
+            )
+
+            rows = c.execute(query).fetchall()
+            return [row[0] for row in rows]
+
     def start_queued_workflows(
-        self, queue: "Queue", executor_id: str, app_version: str
+        self,
+        queue: "Queue",
+        executor_id: str,
+        app_version: str,
+        queue_partition_key: Optional[str],
     ) -> List[str]:
         if self._debug_mode:
             return []
@@ -1734,6 +1774,10 @@ 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
@@ -1758,6 +1802,10 @@ 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
@@ -1799,6 +1847,10 @@ 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