dbos 1.15.0a9__py3-none-any.whl → 2.4.0a5__py3-none-any.whl

dbos/_migration.py

@@ -203,8 +203,39 @@ 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_migration_three(schema: str) -> str:
+    return f"""
+create index "idx_workflow_status_queue_status_started" on \"{schema}\"."workflow_status" ("queue_name", "status", "started_at_epoch_ms")
+"""
+
+
+def get_dbos_migration_four(schema: str) -> str:
+    return f"""
+ALTER TABLE \"{schema}\".workflow_status ADD COLUMN forked_from TEXT;
+CREATE INDEX "idx_workflow_status_forked_from" ON \"{schema}\"."workflow_status" ("forked_from")
+"""
+
+
+def get_dbos_migration_five(schema: str) -> str:
+    return f"""
+ALTER TABLE \"{schema}\".operation_outputs ADD COLUMN started_at_epoch_ms BIGINT, ADD COLUMN completed_at_epoch_ms BIGINT;
+"""
+
+
 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),
+        get_dbos_migration_three(schema),
+        get_dbos_migration_four(schema),
+        get_dbos_migration_five(schema),
+    ]
 
 
 def get_sqlite_timestamp_expr() -> str:
@@ -293,4 +324,30 @@ CREATE TABLE streams (
 );
 """
 
-sqlite_migrations = [sqlite_migration_one]
+sqlite_migration_two = """
+ALTER TABLE workflow_status ADD COLUMN queue_partition_key TEXT;
+"""
+
+sqlite_migration_three = """
+CREATE INDEX "idx_workflow_status_queue_status_started"
+ON "workflow_status" ("queue_name", "status", "started_at_epoch_ms")
+"""
+
+sqlite_migration_four = """
+ALTER TABLE workflow_status ADD COLUMN forked_from TEXT;
+CREATE INDEX "idx_workflow_status_forked_from" ON "workflow_status" ("forked_from")
+"""
+
+sqlite_migration_five = """
+ALTER TABLE operation_outputs ADD COLUMN started_at_epoch_ms BIGINT;
+ALTER TABLE operation_outputs ADD COLUMN completed_at_epoch_ms BIGINT;
+"""
+
+
+sqlite_migrations = [
+    sqlite_migration_one,
+    sqlite_migration_two,
+    sqlite_migration_three,
+    sqlite_migration_four,
+    sqlite_migration_five,
+]

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

@@ -1,7 +1,8 @@
+import random
 import threading
 import traceback
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING, Any, Callable, Coroutine
 
 from ._logger import dbos_logger
 from ._queue import Queue
@@ -13,30 +14,45 @@ from ._context import SetWorkflowID
 from ._croniter import croniter  # type: ignore
 from ._registrations import get_dbos_func_name
 
-ScheduledWorkflow = Callable[[datetime, datetime], None]
-
-scheduler_queue: Queue
+ScheduledWorkflow = (
+    Callable[[datetime, datetime], None]
+    | Callable[[datetime, datetime], Coroutine[Any, Any, None]]
+)
 
 
 def scheduler_loop(
     func: ScheduledWorkflow, cron: str, stop_event: threading.Event
 ) -> None:
+    from dbos._dbos import _get_dbos_instance
+
+    dbos = _get_dbos_instance()
+    scheduler_queue = dbos._registry.get_internal_queue()
     try:
         iter = croniter(cron, datetime.now(timezone.utc), second_at_beginning=True)
-    except Exception as e:
+    except Exception:
         dbos_logger.error(
             f'Cannot run scheduled function {get_dbos_func_name(func)}. Invalid crontab "{cron}"'
         )
+        raise
     while not stop_event.is_set():
-        nextExecTime = iter.get_next(datetime)
-        sleepTime = nextExecTime - datetime.now(timezone.utc)
-        if stop_event.wait(timeout=sleepTime.total_seconds()):
+        next_exec_time = iter.get_next(datetime)
+        sleep_time = (next_exec_time - datetime.now(timezone.utc)).total_seconds()
+        sleep_time = max(0, sleep_time)
+        # To prevent a "thundering herd" problem in a distributed setting,
+        # apply jitter of up to 10% the sleep time, capped at 10 seconds
+        max_jitter = min(sleep_time / 10, 10)
+        jitter = random.uniform(0, max_jitter)
+        if stop_event.wait(timeout=sleep_time + jitter):
             return
         try:
-            with SetWorkflowID(
-                f"sched-{get_dbos_func_name(func)}-{nextExecTime.isoformat()}"
-            ):
-                scheduler_queue.enqueue(func, nextExecTime, datetime.now(timezone.utc))
+            workflowID = (
+                f"sched-{get_dbos_func_name(func)}-{next_exec_time.isoformat()}"
+            )
+            if not dbos._sys_db.get_workflow_status(workflowID):
+                with SetWorkflowID(workflowID):
+                    scheduler_queue.enqueue(
+                        func, next_exec_time, datetime.now(timezone.utc)
+                    )
         except Exception:
             dbos_logger.warning(
                 f"Exception encountered in scheduler thread: {traceback.format_exc()})"
@@ -49,13 +65,10 @@ def scheduled(
     def decorator(func: ScheduledWorkflow) -> ScheduledWorkflow:
         try:
             croniter(cron, datetime.now(timezone.utc), second_at_beginning=True)
-        except Exception as e:
+        except Exception:
             raise ValueError(
                 f'Invalid crontab "{cron}" for scheduled function function {get_dbos_func_name(func)}.'
             )
-
-        global scheduler_queue
-        scheduler_queue = dbosreg.get_internal_queue()
         stop_event = threading.Event()
         dbosreg.register_poller(stop_event, scheduler_loop, func, cron, stop_event)
         return func

dbos/_schemas/system_database.py

@@ -77,6 +77,8 @@ 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()),
+        Column("forked_from", Text()),
         Index("workflow_status_created_at_index", "created_at"),
         Index("workflow_status_executor_id_index", "executor_id"),
         Index("workflow_status_status_index", "status"),
@@ -103,6 +105,8 @@ class SystemSchema:
         Column("output", Text, nullable=True),
         Column("error", Text, nullable=True),
         Column("child_workflow_id", Text, nullable=True),
+        Column("started_at_epoch_ms", BigInteger, nullable=True),
+        Column("completed_at_epoch_ms", BigInteger, nullable=True),
         PrimaryKeyConstraint("workflow_uuid", "function_id"),
     )
 

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,35 @@ 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:
+        try:
+            pickled_data: bytes = pickle.dumps(data)
+            encoded_data: str = base64.b64encode(pickled_data).decode("utf-8")
+            return encoded_data
+        except Exception as e:
+            dbos_logger.error(f"Error serializing object: {data}", exc_info=e)
+            raise
 
-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 +56,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 +68,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 +80,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
         )