dbos 2.4.0a3__py3-none-any.whl → 2.6.0a8__py3-none-any.whl

dbos/_sys_db.py

@@ -158,6 +158,16 @@ class WorkflowStatusInternal(TypedDict):
     forked_from: Optional[str]
 
 
+class MetricData(TypedDict):
+    """
+    Metrics data for workflows and steps within a time range.
+    """
+
+    metric_type: str  # Type of metric: "workflow" or "step"
+    metric_name: str  # Name of the workflow or step
+    value: int  # Number of times the operation ran in the time interval
+
+
 class EnqueueOptionsInternal(TypedDict):
     # Unique ID for deduplication on a queue
     deduplication_id: Optional[str]
@@ -170,16 +180,17 @@ class EnqueueOptionsInternal(TypedDict):
 
 
 class RecordedResult(TypedDict):
-    output: Optional[str]  # JSON (jsonpickle)
-    error: Optional[str]  # JSON (jsonpickle)
+    output: Optional[str]  # Serialized
+    error: Optional[str]  # Serialized
+    child_workflow_id: Optional[str]
 
 
 class OperationResultInternal(TypedDict):
     workflow_uuid: str
     function_id: int
     function_name: str
-    output: Optional[str]  # JSON (jsonpickle)
-    error: Optional[str]  # JSON (jsonpickle)
+    output: Optional[str]  # Serialized
+    error: Optional[str]  # Serialized
     started_at_epoch_ms: int
 
 
@@ -351,6 +362,7 @@ class SystemDatabase(ABC):
         engine: Optional[sa.Engine],
         schema: Optional[str],
         serializer: Serializer,
+        executor_id: Optional[str],
         debug_mode: bool = False,
     ) -> "SystemDatabase":
         """Factory method to create the appropriate SystemDatabase implementation based on URL."""
@@ -363,6 +375,7 @@ class SystemDatabase(ABC):
                 engine=engine,
                 schema=schema,
                 serializer=serializer,
+                executor_id=executor_id,
                 debug_mode=debug_mode,
             )
         else:
@@ -374,6 +387,7 @@ class SystemDatabase(ABC):
                 engine=engine,
                 schema=schema,
                 serializer=serializer,
+                executor_id=executor_id,
                 debug_mode=debug_mode,
             )
 
@@ -385,11 +399,32 @@ class SystemDatabase(ABC):
         engine: Optional[sa.Engine],
         schema: Optional[str],
         serializer: Serializer,
+        executor_id: Optional[str],
         debug_mode: bool = False,
     ):
         import sqlalchemy.dialects.postgresql as pg
         import sqlalchemy.dialects.sqlite as sq
 
+        # Log system database connection information
+        if engine:
+            dbos_logger.info("Initializing DBOS system database with custom engine")
+        else:
+            printable_sys_db_url = sa.make_url(system_database_url).render_as_string(
+                hide_password=True
+            )
+            dbos_logger.info(
+                f"Initializing DBOS system database with URL: {printable_sys_db_url}"
+            )
+            if system_database_url.startswith("sqlite"):
+                dbos_logger.info(
+                    f"Using SQLite as a system database. The SQLite system database is for development and testing. PostgreSQL is recommended for production use."
+                )
+            else:
+                dbos_logger.info(
+                    f"DBOS system database engine parameters: {engine_kwargs}"
+                )
+
+        # Configure and initialize the system database
         self.dialect = sq if system_database_url.startswith("sqlite") else pg
 
         self.serializer = serializer
@@ -410,6 +445,8 @@ class SystemDatabase(ABC):
 
         self.notifications_map = ThreadSafeConditionDict()
         self.workflow_events_map = ThreadSafeConditionDict()
+        self.executor_id = executor_id
+
         self._listener_thread_lock = threading.Lock()
 
         # Now we can run background processes
@@ -694,11 +731,7 @@ class SystemDatabase(ABC):
                     name=status["name"],
                     class_name=status["class_name"],
                     config_name=status["config_name"],
-                    application_version=(
-                        application_version
-                        if application_version is not None
-                        else status["app_version"]
-                    ),
+                    application_version=application_version,
                     application_id=status["app_id"],
                     authenticated_user=status["authenticated_user"],
                     authenticated_roles=status["authenticated_roles"],
@@ -710,34 +743,124 @@ class SystemDatabase(ABC):
             )
 
             if start_step > 1:
+                # Copy the original workflow's step checkpoints
+                c.execute(
+                    sa.insert(SystemSchema.operation_outputs).from_select(
+                        [
+                            "workflow_uuid",
+                            "function_id",
+                            "output",
+                            "error",
+                            "function_name",
+                            "child_workflow_id",
+                            "started_at_epoch_ms",
+                            "completed_at_epoch_ms",
+                        ],
+                        sa.select(
+                            sa.literal(forked_workflow_id).label("workflow_uuid"),
+                            SystemSchema.operation_outputs.c.function_id,
+                            SystemSchema.operation_outputs.c.output,
+                            SystemSchema.operation_outputs.c.error,
+                            SystemSchema.operation_outputs.c.function_name,
+                            SystemSchema.operation_outputs.c.child_workflow_id,
+                            SystemSchema.operation_outputs.c.started_at_epoch_ms,
+                            SystemSchema.operation_outputs.c.completed_at_epoch_ms,
+                        ).where(
+                            (
+                                SystemSchema.operation_outputs.c.workflow_uuid
+                                == original_workflow_id
+                            )
+                            & (
+                                SystemSchema.operation_outputs.c.function_id
+                                < start_step
+                            )
+                        ),
+                    )
+                )
+                # Copy the original workflow's events
+                c.execute(
+                    sa.insert(SystemSchema.workflow_events_history).from_select(
+                        [
+                            "workflow_uuid",
+                            "function_id",
+                            "key",
+                            "value",
+                        ],
+                        sa.select(
+                            sa.literal(forked_workflow_id).label("workflow_uuid"),
+                            SystemSchema.workflow_events_history.c.function_id,
+                            SystemSchema.workflow_events_history.c.key,
+                            SystemSchema.workflow_events_history.c.value,
+                        ).where(
+                            (
+                                SystemSchema.workflow_events_history.c.workflow_uuid
+                                == original_workflow_id
+                            )
+                            & (
+                                SystemSchema.workflow_events_history.c.function_id
+                                < start_step
+                            )
+                        ),
+                    )
+                )
+                # Copy only the latest version of each workflow event from the history table
+                # (the one with the maximum function_id for each key where function_id < start_step)
+                weh1 = SystemSchema.workflow_events_history.alias("weh1")
+                weh2 = SystemSchema.workflow_events_history.alias("weh2")
 
-                # Copy the original workflow's outputs into the forked workflow
-                insert_stmt = sa.insert(SystemSchema.operation_outputs).from_select(
-                    [
-                        "workflow_uuid",
-                        "function_id",
-                        "output",
-                        "error",
-                        "function_name",
-                        "child_workflow_id",
-                    ],
-                    sa.select(
-                        sa.literal(forked_workflow_id).label("workflow_uuid"),
-                        SystemSchema.operation_outputs.c.function_id,
-                        SystemSchema.operation_outputs.c.output,
-                        SystemSchema.operation_outputs.c.error,
-                        SystemSchema.operation_outputs.c.function_name,
-                        SystemSchema.operation_outputs.c.child_workflow_id,
-                    ).where(
-                        (
-                            SystemSchema.operation_outputs.c.workflow_uuid
-                            == original_workflow_id
-                        )
-                        & (SystemSchema.operation_outputs.c.function_id < start_step)
-                    ),
+                max_function_id_subquery = (
+                    sa.select(sa.func.max(weh2.c.function_id))
+                    .where(
+                        (weh2.c.workflow_uuid == original_workflow_id)
+                        & (weh2.c.key == weh1.c.key)
+                        & (weh2.c.function_id < start_step)
+                    )
+                    .scalar_subquery()
+                )
+
+                c.execute(
+                    sa.insert(SystemSchema.workflow_events).from_select(
+                        [
+                            "workflow_uuid",
+                            "key",
+                            "value",
+                        ],
+                        sa.select(
+                            sa.literal(forked_workflow_id).label("workflow_uuid"),
+                            weh1.c.key,
+                            weh1.c.value,
+                        ).where(
+                            (weh1.c.workflow_uuid == original_workflow_id)
+                            & (weh1.c.function_id == max_function_id_subquery)
+                        ),
+                    )
+                )
+                # Copy the original workflow's streams
+                c.execute(
+                    sa.insert(SystemSchema.streams).from_select(
+                        [
+                            "workflow_uuid",
+                            "function_id",
+                            "key",
+                            "value",
+                            "offset",
+                        ],
+                        sa.select(
+                            sa.literal(forked_workflow_id).label("workflow_uuid"),
+                            SystemSchema.streams.c.function_id,
+                            SystemSchema.streams.c.key,
+                            SystemSchema.streams.c.value,
+                            SystemSchema.streams.c.offset,
+                        ).where(
+                            (
+                                SystemSchema.streams.c.workflow_uuid
+                                == original_workflow_id
+                            )
+                            & (SystemSchema.streams.c.function_id < start_step)
+                        ),
+                    )
                 )
 
-                c.execute(insert_stmt)
         return forked_workflow_id
 
     @db_retry()
@@ -828,7 +951,7 @@ class SystemDatabase(ABC):
             return workflow_id
 
     @db_retry()
-    def await_workflow_result(self, workflow_id: str) -> Any:
+    def await_workflow_result(self, workflow_id: str, polling_interval: float) -> Any:
         while True:
             with self.engine.begin() as c:
                 row = c.execute(
@@ -853,7 +976,7 @@ class SystemDatabase(ABC):
                         raise DBOSAwaitedWorkflowCancelledError(workflow_id)
                 else:
                     pass  # CB: I guess we're assuming the WF will show up eventually.
-            time.sleep(1)
+            time.sleep(polling_interval)
 
     def get_workflows(
         self,
@@ -896,11 +1019,12 @@ class SystemDatabase(ABC):
 
         if input.queues_only:
             query = sa.select(*load_columns).where(
-                sa.and_(
-                    SystemSchema.workflow_status.c.queue_name.isnot(None),
-                    SystemSchema.workflow_status.c.status.in_(["ENQUEUED", "PENDING"]),
-                )
+                SystemSchema.workflow_status.c.queue_name.isnot(None),
             )
+            if not input.status:
+                query = query.where(
+                    SystemSchema.workflow_status.c.status.in_(["ENQUEUED", "PENDING"])
+                )
         else:
             query = sa.select(*load_columns)
         if input.sort_desc:
@@ -1027,7 +1151,7 @@ class SystemDatabase(ABC):
                 for row in rows
             ]
 
-    def get_workflow_steps(self, workflow_id: str) -> List[StepInfo]:
+    def list_workflow_steps(self, workflow_id: str) -> List[StepInfo]:
         with self.engine.begin() as c:
             rows = c.execute(
                 sa.select(
@@ -1038,7 +1162,9 @@ class SystemDatabase(ABC):
                     SystemSchema.operation_outputs.c.child_workflow_id,
                     SystemSchema.operation_outputs.c.started_at_epoch_ms,
                     SystemSchema.operation_outputs.c.completed_at_epoch_ms,
-                ).where(SystemSchema.operation_outputs.c.workflow_uuid == workflow_id)
+                )
+                .where(SystemSchema.operation_outputs.c.workflow_uuid == workflow_id)
+                .order_by(SystemSchema.operation_outputs.c.function_id)
             ).fetchall()
             steps = []
             for row in rows:
@@ -1069,17 +1195,44 @@ class SystemDatabase(ABC):
         error = result["error"]
         output = result["output"]
         assert error is None or output is None, "Only one of error or output can be set"
-        sql = sa.insert(SystemSchema.operation_outputs).values(
-            workflow_uuid=result["workflow_uuid"],
-            function_id=result["function_id"],
-            function_name=result["function_name"],
-            started_at_epoch_ms=result["started_at_epoch_ms"],
-            completed_at_epoch_ms=int(time.time() * 1000),
-            output=output,
-            error=error,
-        )
+
+        # Check if the executor ID belong to another process.
+        # Reset it to this process's executor ID if so.
+        wf_executor_id_row = conn.execute(
+            sa.select(
+                SystemSchema.workflow_status.c.executor_id,
+            ).where(
+                SystemSchema.workflow_status.c.workflow_uuid == result["workflow_uuid"]
+            )
+        ).fetchone()
+        assert wf_executor_id_row is not None
+        wf_executor_id = wf_executor_id_row[0]
+        if self.executor_id is not None and wf_executor_id != self.executor_id:
+            dbos_logger.debug(
+                f'Resetting executor_id from {wf_executor_id} to {self.executor_id} for workflow {result["workflow_uuid"]}'
+            )
+            conn.execute(
+                sa.update(SystemSchema.workflow_status)
+                .values(executor_id=self.executor_id)
+                .where(
+                    SystemSchema.workflow_status.c.workflow_uuid
+                    == result["workflow_uuid"]
+                )
+            )
+
+        # Record the outcome, throwing DBOSWorkflowConflictIDError if it is already present
         try:
-            conn.execute(sql)
+            conn.execute(
+                sa.insert(SystemSchema.operation_outputs).values(
+                    workflow_uuid=result["workflow_uuid"],
+                    function_id=result["function_id"],
+                    function_name=result["function_name"],
+                    started_at_epoch_ms=result["started_at_epoch_ms"],
+                    completed_at_epoch_ms=int(time.time() * 1000),
+                    output=output,
+                    error=error,
+                )
+            )
         except DBAPIError as dbapi_error:
             if self._is_unique_constraint_violation(dbapi_error):
                 raise DBOSWorkflowConflictIDError(result["workflow_uuid"])
@@ -1168,6 +1321,7 @@ class SystemDatabase(ABC):
             SystemSchema.operation_outputs.c.output,
             SystemSchema.operation_outputs.c.error,
             SystemSchema.operation_outputs.c.function_name,
+            SystemSchema.operation_outputs.c.child_workflow_id,
         ).where(
             (SystemSchema.operation_outputs.c.workflow_uuid == workflow_id)
             & (SystemSchema.operation_outputs.c.function_id == function_id)
@@ -1196,10 +1350,11 @@ class SystemDatabase(ABC):
             return None
 
         # Extract operation output data
-        output, error, recorded_function_name = (
+        output, error, recorded_function_name, child_workflow_id = (
             operation_output_rows[0][0],
             operation_output_rows[0][1],
             operation_output_rows[0][2],
+            operation_output_rows[0][3],
         )
 
         # If the provided and recorded function name are different, throw an exception
@@ -1214,6 +1369,7 @@ class SystemDatabase(ABC):
         result: RecordedResult = {
             "output": output,
             "error": error,
+            "child_workflow_id": child_workflow_id,
         }
         return result
 
@@ -1226,31 +1382,6 @@ class SystemDatabase(ABC):
                 workflow_id, function_id, function_name, c
             )
 
-    @db_retry()
-    def check_child_workflow(
-        self, workflow_uuid: str, function_id: int
-    ) -> Optional[str]:
-        sql = sa.select(
-            SystemSchema.operation_outputs.c.child_workflow_id,
-            SystemSchema.operation_outputs.c.error,
-        ).where(
-            SystemSchema.operation_outputs.c.workflow_uuid == workflow_uuid,
-            SystemSchema.operation_outputs.c.function_id == function_id,
-        )
-
-        # If in a transaction, use the provided connection
-        row: Any
-        with self.engine.begin() as c:
-            row = c.execute(sql).fetchone()
-
-        if row is None:
-            return None
-        elif row[1]:
-            e: Exception = self.serializer.deserialize(row[1])
-            raise e
-        else:
-            return str(row[0])
-
     @db_retry()
     def send(
         self,
@@ -1503,6 +1634,19 @@ class SystemDatabase(ABC):
                     set_={"value": self.serializer.serialize(message)},
                 )
             )
+            c.execute(
+                self.dialect.insert(SystemSchema.workflow_events_history)
+                .values(
+                    workflow_uuid=workflow_uuid,
+                    function_id=function_id,
+                    key=key,
+                    value=self.serializer.serialize(message),
+                )
+                .on_conflict_do_update(
+                    index_elements=["workflow_uuid", "key", "function_id"],
+                    set_={"value": self.serializer.serialize(message)},
+                )
+            )
             output: OperationResultInternal = {
                 "workflow_uuid": workflow_uuid,
                 "function_id": function_id,
@@ -1516,6 +1660,7 @@ class SystemDatabase(ABC):
     def set_event_from_step(
         self,
         workflow_uuid: str,
+        function_id: int,
         key: str,
         message: Any,
     ) -> None:
@@ -1532,6 +1677,19 @@ class SystemDatabase(ABC):
                     set_={"value": self.serializer.serialize(message)},
                 )
             )
+            c.execute(
+                self.dialect.insert(SystemSchema.workflow_events_history)
+                .values(
+                    workflow_uuid=workflow_uuid,
+                    function_id=function_id,
+                    key=key,
+                    value=self.serializer.serialize(message),
+                )
+                .on_conflict_do_update(
+                    index_elements=["workflow_uuid", "key", "function_id"],
+                    set_={"value": self.serializer.serialize(message)},
+                )
+            )
 
     def get_all_events(self, workflow_id: str) -> Dict[str, Any]:
         """
@@ -1550,7 +1708,6 @@ class SystemDatabase(ABC):
                     SystemSchema.workflow_events.c.value,
                 ).where(SystemSchema.workflow_events.c.workflow_uuid == workflow_id)
             ).fetchall()
-
             events: Dict[str, Any] = {}
             for row in rows:
                 key = row[0]
@@ -1705,10 +1862,6 @@ 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
@@ -1718,6 +1871,11 @@ class SystemDatabase(ABC):
                         > start_time_ms - limiter_period_ms
                     )
                 )
+                if queue_partition_key is not None:
+                    query = query.where(
+                        SystemSchema.workflow_status.c.queue_partition_key
+                        == queue_partition_key
+                    )
                 num_recent_queries = c.execute(query).fetchone()[0]  # type: ignore
                 if num_recent_queries >= queue.limiter["limit"]:
                     return []
@@ -1733,16 +1891,17 @@ 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
                     )
                     .group_by(SystemSchema.workflow_status.c.executor_id)
                 )
+                if queue_partition_key is not None:
+                    pending_tasks_query = pending_tasks_query.where(
+                        SystemSchema.workflow_status.c.queue_partition_key
+                        == queue_partition_key
+                    )
                 pending_workflows = c.execute(pending_tasks_query).fetchall()
                 pending_workflows_dict = {row[0]: row[1] for row in pending_workflows}
                 local_pending_workflows = pending_workflows_dict.get(executor_id, 0)
@@ -1778,10 +1937,6 @@ 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
@@ -1798,6 +1953,11 @@ class SystemDatabase(ABC):
                 # to ensure all processes have a consistent view of the table.
                 .with_for_update(skip_locked=skip_locks, nowait=(not skip_locks))
             )
+            if queue_partition_key is not None:
+                query = query.where(
+                    SystemSchema.workflow_status.c.queue_partition_key
+                    == queue_partition_key
+                )
             if queue.priority_enabled:
                 query = query.order_by(
                     SystemSchema.workflow_status.c.priority.asc(),
@@ -1941,7 +2101,9 @@ class SystemDatabase(ABC):
             dbos_logger.error(f"Error connecting to the DBOS system database: {e}")
             raise
 
-    def write_stream_from_step(self, workflow_uuid: str, key: str, value: Any) -> None:
+    def write_stream_from_step(
+        self, workflow_uuid: str, function_id: int, key: str, value: Any
+    ) -> None:
         """
         Write a key-value pair to the stream at the first unused offset.
         """
@@ -1971,6 +2133,7 @@ class SystemDatabase(ABC):
             c.execute(
                 sa.insert(SystemSchema.streams).values(
                     workflow_uuid=workflow_uuid,
+                    function_id=function_id,
                     key=key,
                     value=serialized_value,
                     offset=next_offset,
@@ -2027,6 +2190,7 @@ class SystemDatabase(ABC):
             c.execute(
                 sa.insert(SystemSchema.streams).values(
                     workflow_uuid=workflow_uuid,
+                    function_id=function_id,
                     key=key,
                     value=serialized_value,
                     offset=next_offset,
@@ -2124,3 +2288,76 @@ class SystemDatabase(ABC):
             return cutoff_epoch_timestamp_ms, [
                 row[0] for row in pending_enqueued_result
             ]
+
+    def get_metrics(self, start_time: str, end_time: str) -> List[MetricData]:
+        """
+        Retrieve the number of workflows and steps that ran in a time range.
+
+        Args:
+            start_time: ISO 8601 formatted start time
+            end_time: ISO 8601 formatted end time
+        """
+        # Convert ISO 8601 times to epoch milliseconds
+        start_epoch_ms = int(
+            datetime.datetime.fromisoformat(start_time).timestamp() * 1000
+        )
+        end_epoch_ms = int(datetime.datetime.fromisoformat(end_time).timestamp() * 1000)
+
+        metrics: List[MetricData] = []
+
+        with self.engine.begin() as c:
+            # Query workflow metrics
+            workflow_query = (
+                sa.select(
+                    SystemSchema.workflow_status.c.name,
+                    func.count(SystemSchema.workflow_status.c.workflow_uuid).label(
+                        "count"
+                    ),
+                )
+                .where(
+                    sa.and_(
+                        SystemSchema.workflow_status.c.created_at >= start_epoch_ms,
+                        SystemSchema.workflow_status.c.created_at < end_epoch_ms,
+                    )
+                )
+                .group_by(SystemSchema.workflow_status.c.name)
+            )
+
+            workflow_results = c.execute(workflow_query).fetchall()
+            for row in workflow_results:
+                metrics.append(
+                    MetricData(
+                        metric_type="workflow_count",
+                        metric_name=row[0],
+                        value=row[1],
+                    )
+                )
+
+            # Query step metrics
+            step_query = (
+                sa.select(
+                    SystemSchema.operation_outputs.c.function_name,
+                    func.count().label("count"),
+                )
+                .where(
+                    sa.and_(
+                        SystemSchema.operation_outputs.c.started_at_epoch_ms
+                        >= start_epoch_ms,
+                        SystemSchema.operation_outputs.c.started_at_epoch_ms
+                        < end_epoch_ms,
+                    )
+                )
+                .group_by(SystemSchema.operation_outputs.c.function_name)
+            )
+
+            step_results = c.execute(step_query).fetchall()
+            for row in step_results:
+                metrics.append(
+                    MetricData(
+                        metric_type="step_count",
+                        metric_name=row[0],
+                        value=row[1],
+                    )
+                )
+
+        return metrics