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

dbos/__init__.py

@@ -14,6 +14,7 @@ from ._kafka_message import KafkaMessage
 from ._queue import Queue
 from ._serialization import Serializer
 from ._sys_db import GetWorkflowsInput, WorkflowStatus, WorkflowStatusString
+from .cli.migration import run_dbos_database_migrations
 
 __all__ = [
     "DBOSConfig",
@@ -37,4 +38,5 @@ __all__ = [
     "Debouncer",
     "DebouncerClient",
     "Serializer",
+    "run_dbos_database_migrations",
 ]

dbos/_app_db.py

@@ -70,6 +70,17 @@ class ApplicationDatabase(ABC):
         schema: Optional[str],
         debug_mode: bool = False,
     ):
+        # Log application database connection information
+        printable_url = sa.make_url(database_url).render_as_string(hide_password=True)
+        dbos_logger.info(
+            f"Initializing DBOS application database with URL: {printable_url}"
+        )
+        if not database_url.startswith("sqlite"):
+            dbos_logger.info(
+                f"DBOS application database engine parameters: {engine_kwargs}"
+            )
+
+        # Configure and initialize the application database
         if database_url.startswith("sqlite"):
             self.schema = None
         else:
@@ -173,81 +184,6 @@ class ApplicationDatabase(ABC):
         }
         return result
 
-    def get_transactions(self, workflow_uuid: str) -> List[StepInfo]:
-        with self.engine.begin() as conn:
-            rows = conn.execute(
-                sa.select(
-                    ApplicationSchema.transaction_outputs.c.function_id,
-                    ApplicationSchema.transaction_outputs.c.function_name,
-                    ApplicationSchema.transaction_outputs.c.output,
-                    ApplicationSchema.transaction_outputs.c.error,
-                ).where(
-                    ApplicationSchema.transaction_outputs.c.workflow_uuid
-                    == workflow_uuid,
-                )
-            ).all()
-        return [
-            StepInfo(
-                function_id=row[0],
-                function_name=row[1],
-                output=(
-                    self.serializer.deserialize(row[2])
-                    if row[2] is not None
-                    else row[2]
-                ),
-                error=(
-                    self.serializer.deserialize(row[3])
-                    if row[3] is not None
-                    else row[3]
-                ),
-                child_workflow_id=None,
-                started_at_epoch_ms=None,
-                completed_at_epoch_ms=None,
-            )
-            for row in rows
-        ]
-
-    def clone_workflow_transactions(
-        self, src_workflow_id: str, forked_workflow_id: str, start_step: int
-    ) -> None:
-        """
-        Copies all steps from dbos.transctions_outputs where function_id < input function_id
-        into a new workflow_uuid. Returns the new workflow_uuid.
-        """
-
-        with self.engine.begin() as conn:
-
-            insert_stmt = sa.insert(ApplicationSchema.transaction_outputs).from_select(
-                [
-                    "workflow_uuid",
-                    "function_id",
-                    "output",
-                    "error",
-                    "txn_id",
-                    "txn_snapshot",
-                    "executor_id",
-                    "function_name",
-                ],
-                sa.select(
-                    sa.literal(forked_workflow_id).label("workflow_uuid"),
-                    ApplicationSchema.transaction_outputs.c.function_id,
-                    ApplicationSchema.transaction_outputs.c.output,
-                    ApplicationSchema.transaction_outputs.c.error,
-                    ApplicationSchema.transaction_outputs.c.txn_id,
-                    ApplicationSchema.transaction_outputs.c.txn_snapshot,
-                    ApplicationSchema.transaction_outputs.c.executor_id,
-                    ApplicationSchema.transaction_outputs.c.function_name,
-                ).where(
-                    (
-                        ApplicationSchema.transaction_outputs.c.workflow_uuid
-                        == src_workflow_id
-                    )
-                    & (ApplicationSchema.transaction_outputs.c.function_id < start_step)
-                ),
-            )
-
-            conn.execute(insert_stmt)
-
     def garbage_collect(
         self, cutoff_epoch_timestamp_ms: int, pending_workflow_ids: list[str]
     ) -> None:
@@ -302,18 +238,24 @@ class PostgresApplicationDatabase(ApplicationDatabase):
             return
         # Check if the database exists
         app_db_url = self.engine.url
-        postgres_db_engine = sa.create_engine(
-            app_db_url.set(database="postgres"),
-            **self._engine_kwargs,
-        )
-        with postgres_db_engine.connect() as conn:
-            conn.execution_options(isolation_level="AUTOCOMMIT")
-            if not conn.execute(
-                sa.text("SELECT 1 FROM pg_database WHERE datname=:db_name"),
-                parameters={"db_name": app_db_url.database},
-            ).scalar():
-                conn.execute(sa.text(f"CREATE DATABASE {app_db_url.database}"))
-        postgres_db_engine.dispose()
+        try:
+            postgres_db_engine = sa.create_engine(
+                app_db_url.set(database="postgres"),
+                **self._engine_kwargs,
+            )
+            with postgres_db_engine.connect() as conn:
+                conn.execution_options(isolation_level="AUTOCOMMIT")
+                if not conn.execute(
+                    sa.text("SELECT 1 FROM pg_database WHERE datname=:db_name"),
+                    parameters={"db_name": app_db_url.database},
+                ).scalar():
+                    conn.execute(sa.text(f"CREATE DATABASE {app_db_url.database}"))
+        except Exception:
+            dbos_logger.warning(
+                f"Could not connect to postgres database to verify existence of {app_db_url.database}. Continuing..."
+            )
+        finally:
+            postgres_db_engine.dispose()
 
         # Create the dbos schema and transaction_outputs table in the application database
         with self.engine.begin() as conn:

dbos/_client.py

@@ -18,6 +18,7 @@ import sqlalchemy as sa
 
 from dbos._app_db import ApplicationDatabase
 from dbos._context import MaxPriority, MinPriority
+from dbos._core import DEFAULT_POLLING_INTERVAL
 from dbos._sys_db import SystemDatabase
 from dbos._utils import generate_uuid
 
@@ -42,7 +43,6 @@ from dbos._workflow_commands import (
     fork_workflow,
     get_workflow,
     list_queued_workflows,
-    list_workflow_steps,
     list_workflows,
 )
 
@@ -85,8 +85,12 @@ class WorkflowHandleClientPolling(Generic[R]):
     def get_workflow_id(self) -> str:
         return self.workflow_id
 
-    def get_result(self) -> R:
-        res: R = self._sys_db.await_workflow_result(self.workflow_id)
+    def get_result(
+        self, *, polling_interval_sec: float = DEFAULT_POLLING_INTERVAL
+    ) -> R:
+        res: R = self._sys_db.await_workflow_result(
+            self.workflow_id, polling_interval_sec
+        )
         return res
 
     def get_status(self) -> WorkflowStatus:
@@ -105,9 +109,11 @@ class WorkflowHandleClientAsyncPolling(Generic[R]):
     def get_workflow_id(self) -> str:
         return self.workflow_id
 
-    async def get_result(self) -> R:
+    async def get_result(
+        self, *, polling_interval_sec: float = DEFAULT_POLLING_INTERVAL
+    ) -> R:
         res: R = await asyncio.to_thread(
-            self._sys_db.await_workflow_result, self.workflow_id
+            self._sys_db.await_workflow_result, self.workflow_id, polling_interval_sec
         )
         return res
 
@@ -472,7 +478,7 @@ class DBOSClient:
         )
 
     def list_workflow_steps(self, workflow_id: str) -> List[StepInfo]:
-        return list_workflow_steps(self._sys_db, self._app_db, workflow_id)
+        return self._sys_db.list_workflow_steps(workflow_id)
 
     async def list_workflow_steps_async(self, workflow_id: str) -> List[StepInfo]:
         return await asyncio.to_thread(self.list_workflow_steps, workflow_id)
@@ -486,7 +492,6 @@ class DBOSClient:
     ) -> "WorkflowHandle[Any]":
         forked_workflow_id = fork_workflow(
             self._sys_db,
-            self._app_db,
             workflow_id,
             start_step,
             application_version=application_version,
@@ -503,7 +508,6 @@ class DBOSClient:
         forked_workflow_id = await asyncio.to_thread(
             fork_workflow,
             self._sys_db,
-            self._app_db,
             workflow_id,
             start_step,
             application_version=application_version,

dbos/_conductor/conductor.py

@@ -16,7 +16,6 @@ from dbos._workflow_commands import (
     get_workflow,
     global_timeout,
     list_queued_workflows,
-    list_workflow_steps,
     list_workflows,
 )
 
@@ -117,6 +116,8 @@ class ConductorWebsocket(threading.Thread):
                                 executor_id=GlobalParams.executor_id,
                                 application_version=GlobalParams.app_version,
                                 hostname=socket.gethostname(),
+                                language="python",
+                                dbos_version=GlobalParams.dbos_version,
                             )
                             websocket.send(info_response.to_json())
                             self.dbos.logger.info("Connected to DBOS conductor")
@@ -339,10 +340,8 @@ class ConductorWebsocket(threading.Thread):
                             list_steps_message = p.ListStepsRequest.from_json(message)
                             step_info = None
                             try:
-                                step_info = list_workflow_steps(
-                                    self.dbos._sys_db,
-                                    self.dbos._app_db,
-                                    list_steps_message.workflow_id,
+                                self.dbos._sys_db.list_workflow_steps(
+                                    list_steps_message.workflow_id
                                 )
                             except Exception as e:
                                 error_message = f"Exception encountered when getting workflow {list_steps_message.workflow_id}: {traceback.format_exc()}"
@@ -397,6 +396,42 @@ class ConductorWebsocket(threading.Thread):
                                 error_message=error_message,
                             )
                             websocket.send(retention_response.to_json())
+                        elif msg_type == p.MessageType.GET_METRICS:
+                            get_metrics_message = p.GetMetricsRequest.from_json(message)
+                            self.dbos.logger.debug(
+                                f"Received metrics request for time range {get_metrics_message.start_time} to {get_metrics_message.end_time}"
+                            )
+                            metrics_data = []
+                            if (
+                                get_metrics_message.metric_class
+                                == "workflow_step_count"
+                            ):
+                                try:
+                                    sys_metrics = self.dbos._sys_db.get_metrics(
+                                        get_metrics_message.start_time,
+                                        get_metrics_message.end_time,
+                                    )
+                                    metrics_data = [
+                                        p.MetricData(
+                                            metric_type=m["metric_type"],
+                                            metric_name=m["metric_name"],
+                                            value=m["value"],
+                                        )
+                                        for m in sys_metrics
+                                    ]
+                                except Exception as e:
+                                    error_message = f"Exception encountered when getting metrics: {traceback.format_exc()}"
+                                    self.dbos.logger.error(error_message)
+                            else:
+                                error_message = f"Unexpected metric class: {get_metrics_message.metric_class}"
+                                self.dbos.logger.warning(error_message)
+                            get_metrics_response = p.GetMetricsResponse(
+                                type=p.MessageType.GET_METRICS,
+                                request_id=base_message.request_id,
+                                metrics=metrics_data,
+                                error_message=error_message,
+                            )
+                            websocket.send(get_metrics_response.to_json())
                         else:
                             self.dbos.logger.warning(
                                 f"Unexpected message type: {msg_type}"

dbos/_conductor/protocol.py

@@ -19,6 +19,7 @@ class MessageType(str, Enum):
     LIST_STEPS = "list_steps"
     FORK_WORKFLOW = "fork_workflow"
     RETENTION = "retention"
+    GET_METRICS = "get_metrics"
 
 
 T = TypeVar("T", bound="BaseMessage")
@@ -63,6 +64,8 @@ class ExecutorInfoResponse(BaseMessage):
     executor_id: str
     application_version: str
     hostname: Optional[str]
+    language: Optional[str]
+    dbos_version: Optional[str]
     error_message: Optional[str] = None
 
 
@@ -339,3 +342,23 @@ class RetentionRequest(BaseMessage):
 class RetentionResponse(BaseMessage):
     success: bool
     error_message: Optional[str] = None
+
+
+@dataclass
+class GetMetricsRequest(BaseMessage):
+    start_time: str  # ISO 8601
+    end_time: str  # ISO 8601
+    metric_class: str
+
+
+@dataclass
+class MetricData:
+    metric_type: str
+    metric_name: str
+    value: int
+
+
+@dataclass
+class GetMetricsResponse(BaseMessage):
+    metrics: List[MetricData]
+    error_message: Optional[str] = None

dbos/_core.py

@@ -91,6 +91,7 @@ F = TypeVar("F", bound=Callable[..., Any])
 
 TEMP_SEND_WF_NAME = "<temp>.temp_send_workflow"
 DEBOUNCER_WORKFLOW_NAME = "_dbos_debouncer_workflow"
+DEFAULT_POLLING_INTERVAL = 1.0
 
 
 class WorkflowHandleFuture(Generic[R]):
@@ -103,7 +104,9 @@ class WorkflowHandleFuture(Generic[R]):
     def get_workflow_id(self) -> str:
         return self.workflow_id
 
-    def get_result(self) -> R:
+    def get_result(
+        self, *, polling_interval_sec: float = DEFAULT_POLLING_INTERVAL
+    ) -> R:
         try:
             r = self.future.result()
         except Exception as e:
@@ -130,9 +133,13 @@ class WorkflowHandlePolling(Generic[R]):
     def get_workflow_id(self) -> str:
         return self.workflow_id
 
-    def get_result(self) -> R:
+    def get_result(
+        self, *, polling_interval_sec: float = DEFAULT_POLLING_INTERVAL
+    ) -> R:
         try:
-            r: R = self.dbos._sys_db.await_workflow_result(self.workflow_id)
+            r: R = self.dbos._sys_db.await_workflow_result(
+                self.workflow_id, polling_interval_sec
+            )
         except Exception as e:
             serialized_e = self.dbos._serializer.serialize(e)
             self.dbos._sys_db.record_get_result(self.workflow_id, None, serialized_e)
@@ -158,7 +165,9 @@ class WorkflowHandleAsyncTask(Generic[R]):
     def get_workflow_id(self) -> str:
         return self.workflow_id
 
-    async def get_result(self) -> R:
+    async def get_result(
+        self, *, polling_interval_sec: float = DEFAULT_POLLING_INTERVAL
+    ) -> R:
         try:
             r = await self.task
         except Exception as e:
@@ -192,10 +201,14 @@ class WorkflowHandleAsyncPolling(Generic[R]):
     def get_workflow_id(self) -> str:
         return self.workflow_id
 
-    async def get_result(self) -> R:
+    async def get_result(
+        self, *, polling_interval_sec: float = DEFAULT_POLLING_INTERVAL
+    ) -> R:
         try:
             r: R = await asyncio.to_thread(
-                self.dbos._sys_db.await_workflow_result, self.workflow_id
+                self.dbos._sys_db.await_workflow_result,
+                self.workflow_id,
+                polling_interval_sec,
             )
         except Exception as e:
             serialized_e = self.dbos._serializer.serialize(e)
@@ -366,7 +379,7 @@ def _get_wf_invoke_func(
             )
             # Directly return the result if the workflow is already completed
             recorded_result: R = dbos._sys_db.await_workflow_result(
-                status["workflow_uuid"]
+                status["workflow_uuid"], polling_interval=DEFAULT_POLLING_INTERVAL
             )
             return recorded_result
         try:
@@ -381,7 +394,9 @@ def _get_wf_invoke_func(
             return output
         except DBOSWorkflowConflictIDError:
             # Await the workflow result
-            r: R = dbos._sys_db.await_workflow_result(status["workflow_uuid"])
+            r: R = dbos._sys_db.await_workflow_result(
+                status["workflow_uuid"], polling_interval=DEFAULT_POLLING_INTERVAL
+            )
             return r
         except DBOSWorkflowCancelledError as error:
             raise DBOSAwaitedWorkflowCancelledError(status["workflow_uuid"])
@@ -579,11 +594,14 @@ def start_workflow(
     ctx = new_wf_ctx
     new_child_workflow_id = ctx.id_assigned_for_next_workflow
     if ctx.has_parent():
-        child_workflow_id = dbos._sys_db.check_child_workflow(
-            ctx.parent_workflow_id, ctx.parent_workflow_fid
+        recorded_result = dbos._sys_db.check_operation_execution(
+            ctx.parent_workflow_id, ctx.parent_workflow_fid, get_dbos_func_name(func)
         )
-        if child_workflow_id is not None:
-            return WorkflowHandlePolling(child_workflow_id, dbos)
+        if recorded_result and recorded_result["error"]:
+            e: Exception = dbos._sys_db.serializer.deserialize(recorded_result["error"])
+            raise e
+        elif recorded_result and recorded_result["child_workflow_id"]:
+            return WorkflowHandlePolling(recorded_result["child_workflow_id"], dbos)
 
     status = _init_workflow(
         dbos,
@@ -675,13 +693,19 @@ async def start_workflow_async(
     ctx = new_wf_ctx
     new_child_workflow_id = ctx.id_assigned_for_next_workflow
     if ctx.has_parent():
-        child_workflow_id = await asyncio.to_thread(
-            dbos._sys_db.check_child_workflow,
+        recorded_result = await asyncio.to_thread(
+            dbos._sys_db.check_operation_execution,
             ctx.parent_workflow_id,
             ctx.parent_workflow_fid,
+            get_dbos_func_name(func),
         )
-        if child_workflow_id is not None:
-            return WorkflowHandleAsyncPolling(child_workflow_id, dbos)
+        if recorded_result and recorded_result["error"]:
+            e: Exception = dbos._sys_db.serializer.deserialize(recorded_result["error"])
+            raise e
+        elif recorded_result and recorded_result["child_workflow_id"]:
+            return WorkflowHandleAsyncPolling(
+                recorded_result["child_workflow_id"], dbos
+            )
 
     status = await asyncio.to_thread(
         _init_workflow,
@@ -788,7 +812,9 @@ def workflow_wrapper(
                 c_wfid: str, dbos: "DBOS"
             ) -> Callable[[Callable[[], R]], R]:
                 def recorded_result_inner(func: Callable[[], R]) -> R:
-                    r: R = dbos._sys_db.await_workflow_result(c_wfid)
+                    r: R = dbos._sys_db.await_workflow_result(
+                        c_wfid, polling_interval=DEFAULT_POLLING_INTERVAL
+                    )
                     return r
 
                 return recorded_result_inner
@@ -798,11 +824,16 @@ def workflow_wrapper(
             workflow_id = ctx.workflow_id
 
             if ctx.has_parent():
-                child_workflow_id = dbos._sys_db.check_child_workflow(
-                    ctx.parent_workflow_id, ctx.parent_workflow_fid
+                r = dbos._sys_db.check_operation_execution(
+                    ctx.parent_workflow_id,
+                    ctx.parent_workflow_fid,
+                    get_dbos_func_name(func),
                 )
-                if child_workflow_id is not None:
-                    return recorded_result(child_workflow_id, dbos)
+                if r and r["error"]:
+                    e: Exception = dbos._sys_db.serializer.deserialize(r["error"])
+                    raise e
+                elif r and r["child_workflow_id"]:
+                    return recorded_result(r["child_workflow_id"], dbos)
 
             status = _init_workflow(
                 dbos,
@@ -889,12 +920,6 @@ def decorate_transaction(
                 )
 
             dbos = dbosreg.dbos
-            ctx = assert_current_dbos_context()
-            status = dbos._sys_db.get_workflow_status(ctx.workflow_id)
-            if status and status["status"] == WorkflowStatusString.CANCELLED.value:
-                raise DBOSWorkflowCancelledError(
-                    f"Workflow {ctx.workflow_id} is cancelled. Aborting transaction {transaction_name}."
-                )
             assert (
                 dbos._app_db
             ), "Transactions can only be used if DBOS is configured with an application_database_url"
@@ -905,6 +930,26 @@ def decorate_transaction(
                 }
                 with EnterDBOSTransaction(session, attributes=attributes):
                     ctx = assert_current_dbos_context()
+                    # Check if the step record for this transaction exists
+                    recorded_step_output = dbos._sys_db.check_operation_execution(
+                        ctx.workflow_id, ctx.function_id, transaction_name
+                    )
+                    if recorded_step_output:
+                        dbos.logger.debug(
+                            f"Replaying transaction, id: {ctx.function_id}, name: {attributes['name']}"
+                        )
+                        if recorded_step_output["error"]:
+                            step_error: Exception = dbos._serializer.deserialize(
+                                recorded_step_output["error"]
+                            )
+                            raise step_error
+                        elif recorded_step_output["output"]:
+                            return dbos._serializer.deserialize(
+                                recorded_step_output["output"]
+                            )
+                        else:
+                            raise Exception("Output and error are both None")
+
                     txn_output: TransactionResultInternal = {
                         "workflow_uuid": ctx.workflow_id,
                         "function_id": ctx.function_id,
@@ -915,6 +960,14 @@ def decorate_transaction(
                         "txn_id": None,
                         "function_name": transaction_name,
                     }
+                    step_output: OperationResultInternal = {
+                        "workflow_uuid": ctx.workflow_id,
+                        "function_id": ctx.function_id,
+                        "function_name": transaction_name,
+                        "output": None,
+                        "error": None,
+                        "started_at_epoch_ms": int(time.time() * 1000),
+                    }
                     retry_wait_seconds = 0.001
                     backoff_factor = 1.5
                     max_retry_wait_seconds = 2.0
@@ -953,8 +1006,18 @@ def decorate_transaction(
                                             )
                                         )
                                         has_recorded_error = True
+                                        step_output["error"] = recorded_output["error"]
+                                        dbos._sys_db.record_operation_result(
+                                            step_output
+                                        )
                                         raise deserialized_error
                                     elif recorded_output["output"]:
+                                        step_output["output"] = recorded_output[
+                                            "output"
+                                        ]
+                                        dbos._sys_db.record_operation_result(
+                                            step_output
+                                        )
                                         return dbos._serializer.deserialize(
                                             recorded_output["output"]
                                         )
@@ -1011,10 +1074,13 @@ def decorate_transaction(
                         finally:
                             # Don't record the error if it was already recorded
                             if txn_error and not has_recorded_error:
-                                txn_output["error"] = dbos._serializer.serialize(
-                                    txn_error
+                                step_output["error"] = txn_output["error"] = (
+                                    dbos._serializer.serialize(txn_error)
                                 )
                                 dbos._app_db.record_transaction_error(txn_output)
+                                dbos._sys_db.record_operation_result(step_output)
+            step_output["output"] = dbos._serializer.serialize(output)
+            dbos._sys_db.record_operation_result(step_output)
             return output
 
         if inspect.iscoroutinefunction(func):
@@ -1283,7 +1349,9 @@ def set_event(dbos: "DBOS", key: str, value: Any) -> None:
                     ctx.workflow_id, ctx.curr_step_function_id, key, value
                 )
         elif cur_ctx.is_step():
-            dbos._sys_db.set_event_from_step(cur_ctx.workflow_id, key, value)
+            dbos._sys_db.set_event_from_step(
+                cur_ctx.workflow_id, cur_ctx.curr_step_function_id, key, value
+            )
         else:
             raise DBOSException(
                 "set_event() must be called from within a workflow or step"