Better Introspection (#501)

- Add all missing database fields to `WorkflowStatus`
- Add `forked_from` and `forked_to` fields to `WorkflowStatus`, tracing
workflow fork lineage and closing
#498
- Use `pool_pre_ping` in the client, closing
#500
- Record step start and end times, enabling visualization of step
duration

Author: kraftp

dbos/_app_db.py

@@ -201,6 +201,8 @@ def get_transactions(self, workflow_uuid: str) -> List[StepInfo]:
                     else row[3]
                 ),
                 child_workflow_id=None,
+                started_at_epoch_ms=None,
+                completed_at_epoch_ms=None,
             )
             for row in rows
         ]

dbos/_client.py

@@ -149,9 +149,11 @@ def __init__(
         self._sys_db = SystemDatabase.create(
             system_database_url=system_database_url,
             engine_kwargs={
+                "connect_args": {"application_name": "dbos_transact_client"},
                 "pool_timeout": 30,
                 "max_overflow": 0,
                 "pool_size": 2,
+                "pool_pre_ping": True,
             },
             engine=system_database_engine,
             schema=dbos_system_schema,
@@ -162,9 +164,11 @@ def __init__(
             self._app_db = ApplicationDatabase.create(
                 database_url=application_database_url,
                 engine_kwargs={
+                    "connect_args": {"application_name": "dbos_transact_client"},
                     "pool_timeout": 30,
                     "max_overflow": 0,
                     "pool_size": 2,
+                    "pool_pre_ping": True,
                 },
                 schema=dbos_system_schema,
                 serializer=serializer,
@@ -234,6 +238,7 @@ def _enqueue(self, options: EnqueueOptions, *args: Any, **kwargs: Any) -> str:
             ),
             "inputs": self._serializer.serialize(inputs),
             "queue_partition_key": enqueue_options_internal["queue_partition_key"],
+            "forked_from": None,
         }
 
         self._sys_db.init_workflow(
@@ -300,6 +305,7 @@ def send(
             "priority": 0,
             "inputs": self._serializer.serialize({"args": (), "kwargs": {}}),
             "queue_partition_key": None,
+            "forked_from": None,
         }
         with self._sys_db.engine.begin() as conn:
             self._sys_db._insert_workflow_status(

dbos/_conductor/protocol.py

@@ -143,6 +143,13 @@ class WorkflowsOutput:
     QueueName: Optional[str]
     ApplicationVersion: Optional[str]
     ExecutorID: Optional[str]
+    WorkflowTimeoutMS: Optional[str]
+    WorkflowDeadlineEpochMS: Optional[str]
+    DeduplicationID: Optional[str]
+    Priority: Optional[str]
+    QueuePartitionKey: Optional[str]
+    ForkedFrom: Optional[str]
+    ForkedTo: Optional[list[str]]
 
     @classmethod
     def from_workflow_information(cls, info: WorkflowStatus) -> "WorkflowsOutput":
@@ -152,12 +159,22 @@ def from_workflow_information(cls, info: WorkflowStatus) -> "WorkflowsOutput":
         inputs_str = str(info.input) if info.input is not None else None
         outputs_str = str(info.output) if info.output is not None else None
         error_str = str(info.error) if info.error is not None else None
-        request_str = None
         roles_str = (
             str(info.authenticated_roles)
             if info.authenticated_roles is not None
             else None
         )
+        workflow_timeout_ms_str = (
+            str(info.workflow_timeout_ms)
+            if info.workflow_timeout_ms is not None
+            else None
+        )
+        workflow_deadline_epoch_ms_str = (
+            str(info.workflow_deadline_epoch_ms)
+            if info.workflow_deadline_epoch_ms is not None
+            else None
+        )
+        priority_str = str(info.priority) if info.priority is not None else None
 
         return cls(
             WorkflowUUID=info.workflow_id,
@@ -176,6 +193,13 @@ def from_workflow_information(cls, info: WorkflowStatus) -> "WorkflowsOutput":
             QueueName=info.queue_name,
             ApplicationVersion=info.app_version,
             ExecutorID=info.executor_id,
+            WorkflowTimeoutMS=workflow_timeout_ms_str,
+            WorkflowDeadlineEpochMS=workflow_deadline_epoch_ms_str,
+            DeduplicationID=info.deduplication_id,
+            Priority=priority_str,
+            QueuePartitionKey=info.queue_partition_key,
+            ForkedFrom=info.forked_from,
+            ForkedTo=info.forked_to,
         )
 
 
@@ -186,14 +210,28 @@ class WorkflowSteps:
     output: Optional[str]
     error: Optional[str]
     child_workflow_id: Optional[str]
+    started_at_epoch_ms: Optional[str]
+    completed_at_epoch_ms: Optional[str]
 
     @classmethod
     def from_step_info(cls, info: StepInfo) -> "WorkflowSteps":
         output_str = str(info["output"]) if info["output"] is not None else None
         error_str = str(info["error"]) if info["error"] is not None else None
+        started_at_str = (
+            str(info["started_at_epoch_ms"])
+            if info["started_at_epoch_ms"] is not None
+            else None
+        )
+        completed_at_str = (
+            str(info["completed_at_epoch_ms"])
+            if info["completed_at_epoch_ms"] is not None
+            else None
+        )
         return cls(
             function_id=info["function_id"],
             function_name=info["function_name"],
+            started_at_epoch_ms=started_at_str,
+            completed_at_epoch_ms=completed_at_str,
             output=output_str,
             error=error_str,
             child_workflow_id=info["child_workflow_id"],

dbos/_core.py

@@ -300,6 +300,7 @@ def _init_workflow(
             if enqueue_options is not None
             else None
         ),
+        "forked_from": None,
     }
 
     # Synchronously record the status and inputs for workflows
@@ -316,6 +317,7 @@ def _init_workflow(
                 "function_name": wf_name,
                 "output": None,
                 "error": dbos._serializer.serialize(e),
+                "started_at_epoch_ms": int(time.time() * 1000),
             }
             dbos._sys_db.record_operation_result(result)
         raise
@@ -1118,6 +1120,7 @@ def record_step_result(func: Callable[[], R]) -> R:
                     "function_name": step_name,
                     "output": None,
                     "error": None,
+                    "started_at_epoch_ms": int(time.time() * 1000),
                 }
 
                 try:

dbos/_dbos.py

@@ -1128,6 +1128,7 @@ def list_workflows(
         name: Optional[str] = None,
         app_version: Optional[str] = None,
         user: Optional[str] = None,
+        queue_name: Optional[str] = None,
         limit: Optional[int] = None,
         offset: Optional[int] = None,
         sort_desc: bool = False,
@@ -1151,6 +1152,7 @@ def fn() -> List[WorkflowStatus]:
                 workflow_id_prefix=workflow_id_prefix,
                 load_input=load_input,
                 load_output=load_output,
+                queue_name=queue_name,
             )
 
         return _get_dbos_instance()._sys_db.call_function_as_step(

dbos/_migration.py

@@ -215,11 +215,25 @@ def get_dbos_migration_three(schema: str) -> str:
 """
 
 
+def get_dbos_migration_four(schema: str) -> str:
+    return f"""
+ALTER TABLE \"{schema}\".workflow_status ADD COLUMN forked_from TEXT;
+"""
+
+
+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),
         get_dbos_migration_two(schema),
         get_dbos_migration_three(schema),
+        get_dbos_migration_four(schema),
+        get_dbos_migration_five(schema),
     ]
 
 
@@ -318,4 +332,20 @@ def get_sqlite_timestamp_expr() -> str:
 ON "workflow_status" ("queue_name", "status", "started_at_epoch_ms")
 """
 
-sqlite_migrations = [sqlite_migration_one, sqlite_migration_two, sqlite_migration_three]
+sqlite_migration_four = """
+ALTER TABLE workflow_status ADD COLUMN forked_from TEXT;
+"""
+
+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/_schemas/system_database.py

@@ -78,6 +78,7 @@ def set_schema(cls, schema_name: Optional[str]) -> None:
         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"),
@@ -104,6 +105,8 @@ def set_schema(cls, schema_name: Optional[str]) -> None:
         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"),
     )