dbos 2.2.0a2__py3-none-any.whl → 2.2.0a3__py3-none-any.whl

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/_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/_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 @@ 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/_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/_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/_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.dist-info → dbos-2.2.0a3.dist-info}/METADATA

@@ -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.dist-info → dbos-2.2.0a3.dist-info}/RECORD

@@ -1,17 +1,17 @@
-dbos-2.2.0a2.dist-info/METADATA,sha256=N8ToPPsy7wCWFBPg5BbVEM_XYJ82Q5LFEFOu7WdlwIU,14532
-dbos-2.2.0a2.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
-dbos-2.2.0a2.dist-info/entry_points.txt,sha256=_QOQ3tVfEjtjBlr1jS4sHqHya9lI2aIEIWkz8dqYp14,58
-dbos-2.2.0a2.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
+dbos-2.2.0a3.dist-info/METADATA,sha256=BIvdTxFdjJo9CSVgU7cplDHr5pR0p74ZuFoUbm3EP9c,14532
+dbos-2.2.0a3.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
+dbos-2.2.0a3.dist-info/entry_points.txt,sha256=_QOQ3tVfEjtjBlr1jS4sHqHya9lI2aIEIWkz8dqYp14,58
+dbos-2.2.0a3.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
 dbos/__init__.py,sha256=M7FdFSBGhcvaLIXrNw_0eR68ijwMWV7_UEyimHMP_F4,1039
 dbos/__main__.py,sha256=G7Exn-MhGrVJVDbgNlpzhfh8WMX_72t3_oJaFT9Lmt8,653
 dbos/_admin_server.py,sha256=hubQJw5T8zGKCPNS6FQTXy8jQ8GTJxoYQaDTMlICl9k,16267
 dbos/_app_db.py,sha256=mvWQ66ebdbiD9fpGKHZBWNVEza6Ulo1D-3UoTB_LwRc,16378
 dbos/_classproperty.py,sha256=f0X-_BySzn3yFDRKB2JpCbLYQ9tLwt1XftfshvY7CBs,626
-dbos/_client.py,sha256=l15aNXTnNQrjmchBNvJypdwdc53dRlMQlUF40iOqDWo,19092
+dbos/_client.py,sha256=oUpIYcUsWHSH_w1fXQ4ZvbtmgnWZNgDmMj7UqUC0iRA,19317
 dbos/_conductor/conductor.py,sha256=3E_hL3c9g9yWqKZkvI6KA0-ZzPMPRo06TOzT1esMiek,24114
 dbos/_conductor/protocol.py,sha256=q3rgLxINFtWFigdOONc-4gX4vn66UmMlJQD6Kj8LnL4,7420
-dbos/_context.py,sha256=cJDxVbswTLXKE5MV4Hmg6gpIX3Dd5mBTG-4lmofWP9E,27668
-dbos/_core.py,sha256=ThzKR7LCqLlR4eNkoGKrssddM9iVJC0UfshmDbx1TEA,50728
+dbos/_context.py,sha256=XKllmsDR_oMcWOuZnoe1X4yv2JeOi_vsAuyWC-mWs_o,28164
+dbos/_core.py,sha256=6OU3SMW5x8CvO7c0LBlHhF1eLiHPLs6nfkkasP73IEo,51124
 dbos/_croniter.py,sha256=XHAyUyibs_59sJQfSNWkP7rqQY6_XrlfuuCxk4jYqek,47559
 dbos/_dbos.py,sha256=dr32Z_NT36JkUxWGyYVX7xkl3bYJmgsxVMOX8H9_mpM,59394
 dbos/_dbos_config.py,sha256=NIMQfxkznoyscyeMFLrfrPAS1W_PHXXWrxqpvvrbp3E,24923
@@ -25,18 +25,18 @@ dbos/_flask.py,sha256=Npnakt-a3W5OykONFRkDRnumaDhTQmA0NPdUCGRYKXE,1652
 dbos/_kafka.py,sha256=Gm4fHWl7gYb-i5BMvwNwm5Km3z8zQpseqdMgqgFjlGI,4252
 dbos/_kafka_message.py,sha256=NYvOXNG3Qn7bghn1pv3fg4Pbs86ILZGcK4IB-MLUNu0,409
 dbos/_logger.py,sha256=djnCp147QoQ1iG9Bt3Uz8RyGaXGmi6gebccXsrA6Cps,4660
-dbos/_migration.py,sha256=VAQxZXWQISifW0JpIG78lowV1MTBJ5ZC4P0YIwqxQhM,10013
+dbos/_migration.py,sha256=Fvc3m4dC4oDpjPMHX-tUZVnXklVB9OMMojSLuVyV9ak,10312
 dbos/_outcome.py,sha256=7HvosMfEHTh1U5P6xok7kFTGLwa2lPaul0YApb3UnN4,8191
-dbos/_queue.py,sha256=cgFFwVPUeQtrTgk7ivoTZb0v9ya8rZK4m7-G-h5gIb4,4846
+dbos/_queue.py,sha256=GmqZHl9smES1KSmpauhSdsnZFJHDyfvRArmC-jBibhw,6228
 dbos/_recovery.py,sha256=K-wlFhdf4yGRm6cUzyhcTjQUS0xp2T5rdNMLiiBErYg,2882
 dbos/_registrations.py,sha256=bEOntObnWaBylnebr5ZpcX2hk7OVLDd1z4BvW4_y3zA,7380
 dbos/_roles.py,sha256=kCuhhg8XLtrHCgKgm44I0abIRTGHltf88OwjEKAUggk,2317
 dbos/_scheduler.py,sha256=n96dNzKMr6-2RQvMxRI6BaoExHbLjw0Kr46j1P-DjP4,2620
 dbos/_schemas/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/_schemas/application_database.py,sha256=SypAS9l9EsaBHFn9FR8jmnqt01M74d9AF1AMa4m2hhI,1040
-dbos/_schemas/system_database.py,sha256=aEkjRQDh9xjdke0d9uFx_20-c9UjQtvuLtHZ24aOypA,5497
+dbos/_schemas/system_database.py,sha256=mNsBV0ttlqJArvOqGPY60WvtuiWrHCpYnVxtvMfe2LI,5544
 dbos/_serialization.py,sha256=8TVXB1c2k3keodNcXszqmcOGTQz2r5UBSYtxn2OrYjI,2804
-dbos/_sys_db.py,sha256=ON8P8DZOy9FZ-TFFcBKMhzTdBXbUhE2oEGKKofDJgwE,85571
+dbos/_sys_db.py,sha256=FDboSk58CyQCAFjOF_KMLnRtIw05OL3IpJHT1qwKEKo,87596
 dbos/_sys_db_postgres.py,sha256=GuyGVyZZD_Wl7LjRSkHnOuZ-hOROlO4Xs2UeDhKq10E,6963
 dbos/_sys_db_sqlite.py,sha256=ifjKdy-Z9vlVIBf5L6XnSaNjiBdvqPE73asVHim4A5Q,6998
 dbos/_templates/dbos-db-starter/README.md,sha256=GhxhBj42wjTt1fWEtwNriHbJuKb66Vzu89G4pxNHw2g,930
@@ -56,4 +56,4 @@ dbos/cli/migration.py,sha256=I0_0ngWTuCPQf6Symbpd0lizaxWUKe3uTYEmuCmsrdU,3775
 dbos/dbos-config.schema.json,sha256=47wofTZ5jlFynec7bG0L369tAXbRQQ2euBxBXvg4m9c,1730
 dbos/py.typed,sha256=QfzXT1Ktfk3Rj84akygc7_42z0lRpCq0Ilh8OXI6Zas,44
 version/__init__.py,sha256=L4sNxecRuqdtSFdpUGX3TtBi9KL3k7YsZVIvv-fv9-A,1678
-dbos-2.2.0a2.dist-info/RECORD,,
+dbos-2.2.0a3.dist-info/RECORD,,