dbos 2.1.0a3__py3-none-any.whl → 2.2.0__py3-none-any.whl

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/_serialization.py

@@ -1,6 +1,6 @@
 import base64
 import pickle
-import types
+from abc import ABC, abstractmethod
 from typing import Any, Dict, Optional, Tuple, TypedDict
 
 from ._logger import dbos_logger
@@ -11,47 +11,31 @@ class WorkflowInputs(TypedDict):
     kwargs: Dict[str, Any]
 
 
-def serialize(data: Any) -> str:
-    pickled_data: bytes = pickle.dumps(data)
-    encoded_data: str = base64.b64encode(pickled_data).decode("utf-8")
-    return encoded_data
+class Serializer(ABC):
 
+    @abstractmethod
+    def serialize(self, data: Any) -> str:
+        pass
 
-def serialize_args(data: WorkflowInputs) -> str:
-    """Serialize args to a base64-encoded string using pickle."""
-    pickled_data: bytes = pickle.dumps(data)
-    encoded_data: str = base64.b64encode(pickled_data).decode("utf-8")
-    return encoded_data
+    @abstractmethod
+    def deserialize(cls, serialized_data: str) -> Any:
+        pass
 
 
-def serialize_exception(data: Exception) -> str:
-    """Serialize an Exception object to a base64-encoded string using pickle."""
-    pickled_data: bytes = pickle.dumps(data)
-    encoded_data: str = base64.b64encode(pickled_data).decode("utf-8")
-    return encoded_data
+class DefaultSerializer(Serializer):
 
+    def serialize(self, data: Any) -> str:
+        pickled_data: bytes = pickle.dumps(data)
+        encoded_data: str = base64.b64encode(pickled_data).decode("utf-8")
+        return encoded_data
 
-def deserialize(serialized_data: str) -> Any:
-    """Deserialize a base64-encoded string back to a Python object using pickle."""
-    pickled_data: bytes = base64.b64decode(serialized_data)
-    return pickle.loads(pickled_data)
-
-
-def deserialize_args(serialized_data: str) -> WorkflowInputs:
-    """Deserialize a base64-encoded string back to a Python object list using pickle."""
-    pickled_data: bytes = base64.b64decode(serialized_data)
-    args: WorkflowInputs = pickle.loads(pickled_data)
-    return args
-
-
-def deserialize_exception(serialized_data: str) -> Exception:
-    """Deserialize a base64-encoded string back to a Python Exception using pickle."""
-    pickled_data: bytes = base64.b64decode(serialized_data)
-    exc: Exception = pickle.loads(pickled_data)
-    return exc
+    def deserialize(cls, serialized_data: str) -> Any:
+        pickled_data: bytes = base64.b64decode(serialized_data)
+        return pickle.loads(pickled_data)
 
 
 def safe_deserialize(
+    serializer: Serializer,
     workflow_id: str,
     *,
     serialized_input: Optional[str],
@@ -68,7 +52,9 @@ def safe_deserialize(
     input: Optional[WorkflowInputs]
     try:
         input = (
-            deserialize_args(serialized_input) if serialized_input is not None else None
+            serializer.deserialize(serialized_input)
+            if serialized_input is not None
+            else None
         )
     except Exception as e:
         dbos_logger.warning(
@@ -78,7 +64,9 @@ def safe_deserialize(
     output: Optional[Any]
     try:
         output = (
-            deserialize(serialized_output) if serialized_output is not None else None
+            serializer.deserialize(serialized_output)
+            if serialized_output is not None
+            else None
         )
     except Exception as e:
         dbos_logger.warning(
@@ -88,7 +76,7 @@ def safe_deserialize(
     exception: Optional[Exception]
     try:
         exception = (
-            deserialize_exception(serialized_exception)
+            serializer.deserialize(serialized_exception)
             if serialized_exception is not None
             else None
         )

dbos/_sys_db.py

@@ -30,7 +30,6 @@ from dbos._utils import (
     retriable_sqlite_exception,
 )
 
-from . import _serialization
 from ._context import get_local_dbos_context
 from ._error import (
     DBOSAwaitedWorkflowCancelledError,
@@ -44,6 +43,7 @@ from ._error import (
 )
 from ._logger import dbos_logger
 from ._schemas.system_database import SystemSchema
+from ._serialization import Serializer, WorkflowInputs, safe_deserialize
 
 if TYPE_CHECKING:
     from ._queue import Queue
@@ -95,7 +95,7 @@ class WorkflowStatus:
     # All roles which the authenticated user could assume
     authenticated_roles: Optional[list[str]]
     # The deserialized workflow input object
-    input: Optional[_serialization.WorkflowInputs]
+    input: Optional[WorkflowInputs]
     # The workflow's output, if any
     output: Optional[Any] = None
     # The error the workflow threw, if any
@@ -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):
@@ -341,6 +345,39 @@ def db_retry(
 
 class SystemDatabase(ABC):
 
+    @staticmethod
+    def create(
+        system_database_url: str,
+        engine_kwargs: Dict[str, Any],
+        engine: Optional[sa.Engine],
+        schema: Optional[str],
+        serializer: Serializer,
+        debug_mode: bool = False,
+    ) -> "SystemDatabase":
+        """Factory method to create the appropriate SystemDatabase implementation based on URL."""
+        if system_database_url.startswith("sqlite"):
+            from ._sys_db_sqlite import SQLiteSystemDatabase
+
+            return SQLiteSystemDatabase(
+                system_database_url=system_database_url,
+                engine_kwargs=engine_kwargs,
+                engine=engine,
+                schema=schema,
+                serializer=serializer,
+                debug_mode=debug_mode,
+            )
+        else:
+            from ._sys_db_postgres import PostgresSystemDatabase
+
+            return PostgresSystemDatabase(
+                system_database_url=system_database_url,
+                engine_kwargs=engine_kwargs,
+                engine=engine,
+                schema=schema,
+                serializer=serializer,
+                debug_mode=debug_mode,
+            )
+
     def __init__(
         self,
         *,
@@ -348,6 +385,7 @@ class SystemDatabase(ABC):
         engine_kwargs: Dict[str, Any],
         engine: Optional[sa.Engine],
         schema: Optional[str],
+        serializer: Serializer,
         debug_mode: bool = False,
     ):
         import sqlalchemy.dialects.postgresql as pg
@@ -355,6 +393,8 @@ class SystemDatabase(ABC):
 
         self.dialect = sq if system_database_url.startswith("sqlite") else pg
 
+        self.serializer = serializer
+
         if system_database_url.startswith("sqlite"):
             self.schema = None
         else:
@@ -454,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"],
@@ -725,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:
@@ -752,6 +794,7 @@ class SystemDatabase(ABC):
                 "deduplication_id": row[16],
                 "priority": row[17],
                 "inputs": row[18],
+                "queue_partition_key": row[19],
             }
             return status
 
@@ -797,10 +840,11 @@ class SystemDatabase(ABC):
                     status = row[0]
                     if status == WorkflowStatusString.SUCCESS.value:
                         output = row[1]
-                        return _serialization.deserialize(output)
+                        return self.serializer.deserialize(output)
                     elif status == WorkflowStatusString.ERROR.value:
                         error = row[2]
-                        raise _serialization.deserialize_exception(error)
+                        e: Exception = self.serializer.deserialize(error)
+                        raise e
                     elif status == WorkflowStatusString.CANCELLED.value:
                         # Raise AwaitedWorkflowCancelledError here, not the cancellation exception
                         # because the awaiting workflow is not being cancelled.
@@ -917,7 +961,8 @@ class SystemDatabase(ABC):
             raw_input = row[17] if load_input else None
             raw_output = row[18] if load_output else None
             raw_error = row[19] if load_output else None
-            inputs, output, exception = _serialization.safe_deserialize(
+            inputs, output, exception = safe_deserialize(
+                self.serializer,
                 info.workflow_id,
                 serialized_input=raw_input,
                 serialized_output=raw_output,
@@ -1028,7 +1073,8 @@ class SystemDatabase(ABC):
             raw_input = row[17] if load_input else None
 
             # Error and Output are not loaded because they should always be None for queued workflows.
-            inputs, output, exception = _serialization.safe_deserialize(
+            inputs, output, exception = safe_deserialize(
+                self.serializer,
                 info.workflow_id,
                 serialized_input=raw_input,
                 serialized_output=None,
@@ -1079,7 +1125,8 @@ class SystemDatabase(ABC):
             ).fetchall()
             steps = []
             for row in rows:
-                _, output, exception = _serialization.safe_deserialize(
+                _, output, exception = safe_deserialize(
+                    self.serializer,
                     workflow_id,
                     serialized_input=None,
                     serialized_output=row[2],
@@ -1278,7 +1325,8 @@ class SystemDatabase(ABC):
         if row is None:
             return None
         elif row[1]:
-            raise _serialization.deserialize_exception(row[1])
+            e: Exception = self.serializer.deserialize(row[1])
+            raise e
         else:
             return str(row[0])
 
@@ -1317,7 +1365,7 @@ class SystemDatabase(ABC):
                     sa.insert(SystemSchema.notifications).values(
                         destination_uuid=destination_uuid,
                         topic=topic,
-                        message=_serialization.serialize(message),
+                        message=self.serializer.serialize(message),
                     )
                 )
             except DBAPIError as dbapi_error:
@@ -1354,7 +1402,7 @@ class SystemDatabase(ABC):
         if recorded_output is not None:
             dbos_logger.debug(f"Replaying recv, id: {function_id}, topic: {topic}")
             if recorded_output["output"] is not None:
-                return _serialization.deserialize(recorded_output["output"])
+                return self.serializer.deserialize(recorded_output["output"])
             else:
                 raise Exception("No output recorded in the last recv")
         else:
@@ -1421,13 +1469,13 @@ class SystemDatabase(ABC):
             rows = c.execute(delete_stmt).fetchall()
             message: Any = None
             if len(rows) > 0:
-                message = _serialization.deserialize(rows[0][0])
+                message = self.serializer.deserialize(rows[0][0])
             self._record_operation_result_txn(
                 {
                     "workflow_uuid": workflow_uuid,
                     "function_id": function_id,
                     "function_name": function_name,
-                    "output": _serialization.serialize(
+                    "output": self.serializer.serialize(
                         message
                     ),  # None will be serialized to 'null'
                     "error": None,
@@ -1453,36 +1501,6 @@ class SystemDatabase(ABC):
 
             PostgresSystemDatabase._reset_system_database(database_url)
 
-    @staticmethod
-    def create(
-        system_database_url: str,
-        engine_kwargs: Dict[str, Any],
-        engine: Optional[sa.Engine],
-        schema: Optional[str],
-        debug_mode: bool = False,
-    ) -> "SystemDatabase":
-        """Factory method to create the appropriate SystemDatabase implementation based on URL."""
-        if system_database_url.startswith("sqlite"):
-            from ._sys_db_sqlite import SQLiteSystemDatabase
-
-            return SQLiteSystemDatabase(
-                system_database_url=system_database_url,
-                engine_kwargs=engine_kwargs,
-                engine=engine,
-                schema=schema,
-                debug_mode=debug_mode,
-            )
-        else:
-            from ._sys_db_postgres import PostgresSystemDatabase
-
-            return PostgresSystemDatabase(
-                system_database_url=system_database_url,
-                engine_kwargs=engine_kwargs,
-                engine=engine,
-                schema=schema,
-                debug_mode=debug_mode,
-            )
-
     @db_retry()
     def sleep(
         self,
@@ -1502,7 +1520,7 @@ class SystemDatabase(ABC):
         if recorded_output is not None:
             dbos_logger.debug(f"Replaying sleep, id: {function_id}, seconds: {seconds}")
             assert recorded_output["output"] is not None, "no recorded end time"
-            end_time = _serialization.deserialize(recorded_output["output"])
+            end_time = self.serializer.deserialize(recorded_output["output"])
         else:
             dbos_logger.debug(f"Running sleep, id: {function_id}, seconds: {seconds}")
             end_time = time.time() + seconds
@@ -1512,7 +1530,7 @@ class SystemDatabase(ABC):
                         "workflow_uuid": workflow_uuid,
                         "function_id": function_id,
                         "function_name": function_name,
-                        "output": _serialization.serialize(end_time),
+                        "output": self.serializer.serialize(end_time),
                         "error": None,
                     }
                 )
@@ -1550,11 +1568,11 @@ class SystemDatabase(ABC):
                 .values(
                     workflow_uuid=workflow_uuid,
                     key=key,
-                    value=_serialization.serialize(message),
+                    value=self.serializer.serialize(message),
                 )
                 .on_conflict_do_update(
                     index_elements=["workflow_uuid", "key"],
-                    set_={"value": _serialization.serialize(message)},
+                    set_={"value": self.serializer.serialize(message)},
                 )
             )
             output: OperationResultInternal = {
@@ -1578,11 +1596,11 @@ class SystemDatabase(ABC):
                 .values(
                     workflow_uuid=workflow_uuid,
                     key=key,
-                    value=_serialization.serialize(message),
+                    value=self.serializer.serialize(message),
                 )
                 .on_conflict_do_update(
                     index_elements=["workflow_uuid", "key"],
-                    set_={"value": _serialization.serialize(message)},
+                    set_={"value": self.serializer.serialize(message)},
                 )
             )
 
@@ -1607,7 +1625,7 @@ class SystemDatabase(ABC):
             events: Dict[str, Any] = {}
             for row in rows:
                 key = row[0]
-                value = _serialization.deserialize(row[1])
+                value = self.serializer.deserialize(row[1])
                 events[key] = value
 
             return events
@@ -1641,7 +1659,7 @@ class SystemDatabase(ABC):
                     f"Replaying get_event, id: {caller_ctx['function_id']}, key: {key}"
                 )
                 if recorded_output["output"] is not None:
-                    return _serialization.deserialize(recorded_output["output"])
+                    return self.serializer.deserialize(recorded_output["output"])
                 else:
                     raise Exception("No output recorded in the last get_event")
             else:
@@ -1666,7 +1684,7 @@ class SystemDatabase(ABC):
 
         value: Any = None
         if len(init_recv) > 0:
-            value = _serialization.deserialize(init_recv[0][0])
+            value = self.serializer.deserialize(init_recv[0][0])
         else:
             # Wait for the notification
             actual_timeout = timeout_seconds
@@ -1684,7 +1702,7 @@ class SystemDatabase(ABC):
             with self.engine.begin() as c:
                 final_recv = c.execute(get_sql).fetchall()
                 if len(final_recv) > 0:
-                    value = _serialization.deserialize(final_recv[0][0])
+                    value = self.serializer.deserialize(final_recv[0][0])
         condition.release()
         self.workflow_events_map.pop(payload)
 
@@ -1695,7 +1713,7 @@ class SystemDatabase(ABC):
                     "workflow_uuid": caller_ctx["workflow_uuid"],
                     "function_id": caller_ctx["function_id"],
                     "function_name": function_name,
-                    "output": _serialization.serialize(
+                    "output": self.serializer.serialize(
                         value
                     ),  # None will be serialized to 'null'
                     "error": None,
@@ -1703,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 []
@@ -1723,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
@@ -1747,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
@@ -1788,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
@@ -1897,12 +1960,13 @@ class SystemDatabase(ABC):
             )
             if res is not None:
                 if res["output"] is not None:
-                    resstat: SystemDatabase.T = _serialization.deserialize(
+                    resstat: SystemDatabase.T = self.serializer.deserialize(
                         res["output"]
                     )
                     return resstat
                 elif res["error"] is not None:
-                    raise _serialization.deserialize_exception(res["error"])
+                    e: Exception = self.serializer.deserialize(res["error"])
+                    raise e
                 else:
                     raise Exception(
                         f"Recorded output and error are both None for {function_name}"
@@ -1914,7 +1978,7 @@ class SystemDatabase(ABC):
                     "workflow_uuid": ctx.workflow_id,
                     "function_id": ctx.function_id,
                     "function_name": function_name,
-                    "output": _serialization.serialize(result),
+                    "output": self.serializer.serialize(result),
                     "error": None,
                 }
             )
@@ -1968,7 +2032,7 @@ class SystemDatabase(ABC):
             )
 
             # Serialize the value before storing
-            serialized_value = _serialization.serialize(value)
+            serialized_value = self.serializer.serialize(value)
 
             # Insert the new stream entry
             c.execute(
@@ -2023,7 +2087,7 @@ class SystemDatabase(ABC):
             )
 
             # Serialize the value before storing
-            serialized_value = _serialization.serialize(value)
+            serialized_value = self.serializer.serialize(value)
 
             # Insert the new stream entry
             c.execute(
@@ -2068,7 +2132,7 @@ class SystemDatabase(ABC):
                 )
 
             # Deserialize the value before returning
-            return _serialization.deserialize(result[0])
+            return self.serializer.deserialize(result[0])
 
     def garbage_collect(
         self, cutoff_epoch_timestamp_ms: Optional[int], rows_threshold: Optional[int]

dbos/cli/migration.py

@@ -4,6 +4,7 @@ import sqlalchemy as sa
 import typer
 
 from dbos._app_db import ApplicationDatabase
+from dbos._serialization import DefaultSerializer
 from dbos._sys_db import SystemDatabase
 
 
@@ -22,6 +23,7 @@ def migrate_dbos_databases(
             },
             engine=None,
             schema=schema,
+            serializer=DefaultSerializer(),
         )
         sys_db.run_migrations()
         if app_database_url:
@@ -33,6 +35,7 @@ def migrate_dbos_databases(
                     "pool_size": 2,
                 },
                 schema=schema,
+                serializer=DefaultSerializer(),
             )
             app_db.run_migrations()
     except Exception as e:

{dbos-2.1.0a3.dist-info → dbos-2.2.0.dist-info}/METADATA

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