dbos 2.2.0a2__tar.gz → 2.2.0a3__tar.gz

{dbos-2.2.0a2 → dbos-2.2.0a3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dbos
-Version: 2.2.0a2
+Version: 2.2.0a3
 Summary: Ultra-lightweight durable execution in Python
 Author-Email: "DBOS, Inc." <contact@dbos.dev>
 License: MIT

{dbos-2.2.0a2 → dbos-2.2.0a3}/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 @@ class DBOSClient:
             "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 @@ class DBOSClient:
                 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 @@ class DBOSClient:
             "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-2.2.0a2 → dbos-2.2.0a3}/dbos/_context.py

@@ -120,6 +120,8 @@ class DBOSContext:
         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 @@ class SetEnqueueOptions:
         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 @@ class SetEnqueueOptions:
         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 @@ class 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 @@ class SetEnqueueOptions:
         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-2.2.0a2 → dbos-2.2.0a3}/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-2.2.0a2 → dbos-2.2.0a3}/dbos/_migration.py

@@ -203,8 +203,14 @@ CREATE TABLE \"{schema}\".event_dispatch_kv (
 """
 
 
+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 @@ CREATE TABLE streams (
 );
 """
 
-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-2.2.0a2 → dbos-2.2.0a3}/dbos/_queue.py

@@ -43,6 +43,7 @@ class Queue:
         *,  # 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 @@ class Queue:
         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 @@ class Queue:
             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-2.2.0a2 → dbos-2.2.0a3}/dbos/_schemas/system_database.py

@@ -77,6 +77,7 @@ class SystemSchema:
         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-2.2.0a2 → dbos-2.2.0a3}/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 @@ class SystemDatabase(ABC):
                 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 @@ class SystemDatabase(ABC):
                     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 @@ class SystemDatabase(ABC):
                 "deduplication_id": row[16],
                 "priority": row[17],
                 "inputs": row[18],
+                "queue_partition_key": row[19],
             }
             return status
 
@@ -1714,8 +1721,41 @@ class SystemDatabase(ABC):
             )
         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 @@ class SystemDatabase(ABC):
                     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 @@ class SystemDatabase(ABC):
                     )
                     .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 @@ class SystemDatabase(ABC):
                 )
                 .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

{dbos-2.2.0a2 → dbos-2.2.0a3}/pyproject.toml

@@ -34,7 +34,7 @@ classifiers = [
     "Topic :: Software Development :: Libraries :: Python Modules",
     "Framework :: AsyncIO",
 ]
-version = "2.2.0a2"
+version = "2.2.0a3"
 
 [project.license]
 text = "MIT"

{dbos-2.2.0a2 → dbos-2.2.0a3}/tests/test_queue.py

@@ -1612,3 +1612,80 @@ async def test_enqueue_version_async(dbos: DBOS) -> None:
     # Change the global version, verify it works
     GlobalParams.app_version = future_version
     assert await handle.get_result() == input
+
+
+def test_queue_partitions(dbos: DBOS, client: DBOSClient) -> None:
+
+    blocking_event = threading.Event()
+    waiting_event = threading.Event()
+
+    @DBOS.workflow()
+    def blocked_workflow() -> str:
+        waiting_event.set()
+        blocking_event.wait()
+        assert DBOS.workflow_id
+        return DBOS.workflow_id
+
+    @DBOS.workflow()
+    def normal_workflow() -> str:
+        assert DBOS.workflow_id
+        return DBOS.workflow_id
+
+    queue = Queue("queue", partition_queue=True, worker_concurrency=1)
+
+    blocked_partition_key = "blocked"
+    normal_partition_key = "normal"
+
+    # Enqueue a blocked workflow and a normal workflow on
+    # the blocked partition. Verify the blocked workflow starts
+    # but the normal workflow is stuck behind it.
+    with SetEnqueueOptions(queue_partition_key=blocked_partition_key):
+        blocked_blocked_handle = queue.enqueue(blocked_workflow)
+        blocked_normal_handle = queue.enqueue(normal_workflow)
+
+    waiting_event.wait()
+    assert (
+        blocked_blocked_handle.get_status().status == WorkflowStatusString.PENDING.value
+    )
+    assert (
+        blocked_normal_handle.get_status().status == WorkflowStatusString.ENQUEUED.value
+    )
+
+    # Enqueue a normal workflow on the other partition and verify it runs normally
+    with SetEnqueueOptions(queue_partition_key=normal_partition_key):
+        normal_handle = queue.enqueue(normal_workflow)
+
+    assert normal_handle.get_result()
+
+    # Unblock the blocked partition and verify its workflows complete
+    blocking_event.set()
+    assert blocked_blocked_handle.get_result()
+    assert blocked_normal_handle.get_result()
+
+    # Confirm client enqueue works with partitions
+    client_handle: WorkflowHandle[None] = client.enqueue(
+        {
+            "queue_name": queue.name,
+            "workflow_name": normal_workflow.__qualname__,
+            "queue_partition_key": blocked_partition_key,
+        }
+    )
+    assert client_handle.get_result()
+
+    # You can only enqueue on a partitioned queue with a partition key
+    with pytest.raises(Exception):
+        queue.enqueue(normal_workflow)
+
+    # Deduplication is not supported for partitioned queues
+    with pytest.raises(Exception):
+        with SetEnqueueOptions(
+            queue_partition_key=normal_partition_key, deduplication_id="key"
+        ):
+            queue.enqueue(normal_workflow)
+
+    # You can only enqueue with a partition key on a partitioned queue
+    partitionless_queue = Queue("partitionless-queue")
+
+    with pytest.raises(Exception):
+        with SetEnqueueOptions(queue_partition_key="test"):
+            partitionless_queue.enqueue(normal_workflow)