dbos 0.24.1__py3-none-any.whl → 0.25.0__py3-none-any.whl

dbos/_sys_db.py

@@ -28,11 +28,11 @@ from sqlalchemy.sql import func
 from dbos._utils import GlobalParams
 
 from . import _serialization
-from ._dbos_config import ConfigFile
+from ._context import get_local_dbos_context
+from ._dbos_config import ConfigFile, DatabaseConfig
 from ._error import (
     DBOSConflictingWorkflowError,
     DBOSDeadLetterQueueError,
-    DBOSException,
     DBOSNonExistentWorkflowError,
     DBOSWorkflowConflictIDError,
 )
@@ -89,6 +89,7 @@ class RecordedResult(TypedDict):
 class OperationResultInternal(TypedDict):
     workflow_uuid: str
     function_id: int
+    function_name: str
     output: Optional[str]  # JSON (jsonpickle)
     error: Optional[str]  # JSON (jsonpickle)
 
@@ -114,7 +115,7 @@ class GetWorkflowsInput:
         self.authenticated_user: Optional[str] = None  # The user who ran the workflow.
         self.start_time: Optional[str] = None  # Timestamp in ISO 8601 format
         self.end_time: Optional[str] = None  # Timestamp in ISO 8601 format
-        self.status: Optional[WorkflowStatuses] = None
+        self.status: Optional[str] = None
         self.application_version: Optional[str] = (
             None  # The application version that ran this workflow. = None
         )
@@ -151,30 +152,39 @@ class GetPendingWorkflowsOutput:
         self.queue_name: Optional[str] = queue_name
 
 
+class StepInfo(TypedDict):
+    # The unique ID of the step in the workflow
+    function_id: int
+    # The (fully qualified) name of the step
+    function_name: str
+    # The step's output, if any
+    output: Optional[Any]
+    # The error the step threw, if any
+    error: Optional[Exception]
+    # If the step starts or retrieves the result of a workflow, its ID
+    child_workflow_id: Optional[str]
+
+
 _dbos_null_topic = "__null__topic__"
-_buffer_flush_batch_size = 100
-_buffer_flush_interval_secs = 1.0
 
 
 class SystemDatabase:
 
-    def __init__(self, config: ConfigFile, *, debug_mode: bool = False):
-        self.config = config
-
+    def __init__(self, database: DatabaseConfig, *, debug_mode: bool = False):
         sysdb_name = (
-            config["database"]["sys_db_name"]
-            if "sys_db_name" in config["database"] and config["database"]["sys_db_name"]
-            else config["database"]["app_db_name"] + SystemSchema.sysdb_suffix
+            database["sys_db_name"]
+            if "sys_db_name" in database and database["sys_db_name"]
+            else database["app_db_name"] + SystemSchema.sysdb_suffix
         )
 
         if not debug_mode:
             # If the system database does not already exist, create it
             postgres_db_url = sa.URL.create(
                 "postgresql+psycopg",
-                username=config["database"]["username"],
-                password=config["database"]["password"],
-                host=config["database"]["hostname"],
-                port=config["database"]["port"],
+                username=database["username"],
+                password=database["password"],
+                host=database["hostname"],
+                port=database["port"],
                 database="postgres",
                 # fills the "application_name" column in pg_stat_activity
                 query={"application_name": f"dbos_transact_{GlobalParams.executor_id}"},
@@ -191,19 +201,23 @@ class SystemDatabase:
 
         system_db_url = sa.URL.create(
             "postgresql+psycopg",
-            username=config["database"]["username"],
-            password=config["database"]["password"],
-            host=config["database"]["hostname"],
-            port=config["database"]["port"],
+            username=database["username"],
+            password=database["password"],
+            host=database["hostname"],
+            port=database["port"],
             database=sysdb_name,
             # fills the "application_name" column in pg_stat_activity
             query={"application_name": f"dbos_transact_{GlobalParams.executor_id}"},
         )
 
         # Create a connection pool for the system database
+        pool_size = database.get("sys_db_pool_size")
+        if pool_size is None:
+            pool_size = 20
+
         self.engine = sa.create_engine(
             system_db_url,
-            pool_size=config["database"]["sys_db_pool_size"],
+            pool_size=pool_size,
             max_overflow=0,
             pool_timeout=30,
             connect_args={"connect_timeout": 10},
@@ -250,32 +264,17 @@ class SystemDatabase:
         self.notifications_map: Dict[str, threading.Condition] = {}
         self.workflow_events_map: Dict[str, threading.Condition] = {}
 
-        # Initialize the workflow status and inputs buffers
-        self._workflow_status_buffer: Dict[str, WorkflowStatusInternal] = {}
-        self._workflow_inputs_buffer: Dict[str, str] = {}
-        # Two sets for tracking which single-transaction workflows have been exported to the status table
-        self._exported_temp_txn_wf_status: Set[str] = set()
-        self._temp_txn_wf_ids: Set[str] = set()
-        self._is_flushing_status_buffer = False
-
         # Now we can run background processes
         self._run_background_processes = True
         self._debug_mode = debug_mode
 
     # Destroy the pool when finished
     def destroy(self) -> None:
-        self.wait_for_buffer_flush()
         self._run_background_processes = False
         if self.notification_conn is not None:
             self.notification_conn.close()
         self.engine.dispose()
 
-    def wait_for_buffer_flush(self) -> None:
-        # Wait until the buffers are flushed.
-        while self._is_flushing_status_buffer or not self._is_buffers_empty:
-            dbos_logger.debug("Waiting for system buffers to be exported")
-            time.sleep(1)
-
     def insert_workflow_status(
         self,
         status: WorkflowStatusInternal,
@@ -426,10 +425,6 @@ class SystemDatabase:
             with self.engine.begin() as c:
                 c.execute(cmd)
 
-        # If this is a single-transaction workflow, record that its status has been exported
-        if status["workflow_uuid"] in self._temp_txn_wf_ids:
-            self._exported_temp_txn_wf_status.add(status["workflow_uuid"])
-
     def cancel_workflow(
         self,
         workflow_id: str,
@@ -531,31 +526,6 @@ class SystemDatabase:
             }
             return status
 
-    def get_workflow_status_within_wf(
-        self, workflow_uuid: str, calling_wf: str, calling_wf_fn: int
-    ) -> Optional[WorkflowStatusInternal]:
-        res = self.check_operation_execution(calling_wf, calling_wf_fn)
-        if res is not None:
-            if res["output"]:
-                resstat: WorkflowStatusInternal = _serialization.deserialize(
-                    res["output"]
-                )
-                return resstat
-            else:
-                raise DBOSException(
-                    "Workflow status record not found. This should not happen! \033[1m Hint: Check if your workflow is deterministic.\033[0m"
-                )
-        stat = self.get_workflow_status(workflow_uuid)
-        self.record_operation_result(
-            {
-                "workflow_uuid": calling_wf,
-                "function_id": calling_wf_fn,
-                "output": _serialization.serialize(stat),
-                "error": None,
-            }
-        )
-        return stat
-
     def await_workflow_result_internal(self, workflow_uuid: str) -> dict[str, Any]:
         polling_interval_secs: float = 1.000
 
@@ -632,10 +602,7 @@ class SystemDatabase:
                     f"Workflow {workflow_uuid} has been called multiple times with different inputs"
                 )
             # TODO: actually changing the input
-        if workflow_uuid in self._temp_txn_wf_ids:
-            # Clean up the single-transaction tracking sets
-            self._exported_temp_txn_wf_status.discard(workflow_uuid)
-            self._temp_txn_wf_ids.discard(workflow_uuid)
+
         return
 
     def get_workflow_inputs(
@@ -771,6 +738,36 @@ class SystemDatabase:
                 for row in rows
             ]
 
+    def get_workflow_steps(self, workflow_id: str) -> List[StepInfo]:
+        with self.engine.begin() as c:
+            rows = c.execute(
+                sa.select(
+                    SystemSchema.operation_outputs.c.function_id,
+                    SystemSchema.operation_outputs.c.function_name,
+                    SystemSchema.operation_outputs.c.output,
+                    SystemSchema.operation_outputs.c.error,
+                    SystemSchema.operation_outputs.c.child_workflow_id,
+                ).where(SystemSchema.operation_outputs.c.workflow_uuid == workflow_id)
+            ).fetchall()
+            return [
+                StepInfo(
+                    function_id=row[0],
+                    function_name=row[1],
+                    output=(
+                        _serialization.deserialize(row[2])
+                        if row[2] is not None
+                        else row[2]
+                    ),
+                    error=(
+                        _serialization.deserialize_exception(row[3])
+                        if row[3] is not None
+                        else row[3]
+                    ),
+                    child_workflow_id=row[4],
+                )
+                for row in rows
+            ]
+
     def record_operation_result(
         self, result: OperationResultInternal, conn: Optional[sa.Connection] = None
     ) -> None:
@@ -782,6 +779,7 @@ class SystemDatabase:
         sql = pg.insert(SystemSchema.operation_outputs).values(
             workflow_uuid=result["workflow_uuid"],
             function_id=result["function_id"],
+            function_name=result["function_name"],
             output=output,
             error=error,
         )
@@ -796,6 +794,55 @@ class SystemDatabase:
                 raise DBOSWorkflowConflictIDError(result["workflow_uuid"])
             raise
 
+    def record_get_result(
+        self, result_workflow_id: str, output: Optional[str], error: Optional[str]
+    ) -> None:
+        ctx = get_local_dbos_context()
+        # Only record get_result called in workflow functions
+        if ctx is None or not ctx.is_workflow():
+            return
+        ctx.function_id += 1  # Record the get_result as a step
+        # Because there's no corresponding check, we do nothing on conflict
+        # and do not raise a DBOSWorkflowConflictIDError
+        sql = (
+            pg.insert(SystemSchema.operation_outputs)
+            .values(
+                workflow_uuid=ctx.workflow_id,
+                function_id=ctx.function_id,
+                function_name="DBOS.getResult",
+                output=output,
+                error=error,
+                child_workflow_id=result_workflow_id,
+            )
+            .on_conflict_do_nothing()
+        )
+        with self.engine.begin() as c:
+            c.execute(sql)
+
+    def record_child_workflow(
+        self,
+        parentUUID: str,
+        childUUID: str,
+        functionID: int,
+        functionName: str,
+    ) -> None:
+        if self._debug_mode:
+            raise Exception("called record_child_workflow in debug mode")
+
+        sql = pg.insert(SystemSchema.operation_outputs).values(
+            workflow_uuid=parentUUID,
+            function_id=functionID,
+            function_name=functionName,
+            child_workflow_id=childUUID,
+        )
+        try:
+            with self.engine.begin() as c:
+                c.execute(sql)
+        except DBAPIError as dbapi_error:
+            if dbapi_error.orig.sqlstate == "23505":  # type: ignore
+                raise DBOSWorkflowConflictIDError(parentUUID)
+            raise
+
     def check_operation_execution(
         self, workflow_uuid: str, function_id: int, conn: Optional[sa.Connection] = None
     ) -> Optional[RecordedResult]:
@@ -822,6 +869,23 @@ class SystemDatabase:
         }
         return result
 
+    def check_child_workflow(
+        self, workflow_uuid: str, function_id: int
+    ) -> Optional[str]:
+        sql = sa.select(SystemSchema.operation_outputs.c.child_workflow_id).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
+        return str(row[0])
+
     def send(
         self,
         workflow_uuid: str,
@@ -866,6 +930,7 @@ class SystemDatabase:
             output: OperationResultInternal = {
                 "workflow_uuid": workflow_uuid,
                 "function_id": function_id,
+                "function_name": "DBOS.send",
                 "output": None,
                 "error": None,
             }
@@ -959,6 +1024,7 @@ class SystemDatabase:
                 {
                     "workflow_uuid": workflow_uuid,
                     "function_id": function_id,
+                    "function_name": "DBOS.recv",
                     "output": _serialization.serialize(
                         message
                     ),  # None will be serialized to 'null'
@@ -1049,6 +1115,7 @@ class SystemDatabase:
                     {
                         "workflow_uuid": workflow_uuid,
                         "function_id": function_id,
+                        "function_name": "DBOS.sleep",
                         "output": _serialization.serialize(end_time),
                         "error": None,
                     }
@@ -1096,6 +1163,7 @@ class SystemDatabase:
             output: OperationResultInternal = {
                 "workflow_uuid": workflow_uuid,
                 "function_id": function_id,
+                "function_name": "DBOS.setEvent",
                 "output": None,
                 "error": None,
             }
@@ -1176,6 +1244,7 @@ class SystemDatabase:
                 {
                     "workflow_uuid": caller_ctx["workflow_uuid"],
                     "function_id": caller_ctx["function_id"],
+                    "function_name": "DBOS.getEvent",
                     "output": _serialization.serialize(
                         value
                     ),  # None will be serialized to 'null'
@@ -1184,106 +1253,6 @@ class SystemDatabase:
             )
         return value
 
-    def _flush_workflow_status_buffer(self) -> None:
-        if self._debug_mode:
-            raise Exception("called _flush_workflow_status_buffer in debug mode")
-
-        """Export the workflow status buffer to the database, up to the batch size."""
-        if len(self._workflow_status_buffer) == 0:
-            return
-
-        # Record the exported status so far, and add them back on errors.
-        exported_status: Dict[str, WorkflowStatusInternal] = {}
-        with self.engine.begin() as c:
-            exported = 0
-            status_iter = iter(list(self._workflow_status_buffer))
-            wf_id: Optional[str] = None
-            while (
-                exported < _buffer_flush_batch_size
-                and (wf_id := next(status_iter, None)) is not None
-            ):
-                # Pop the first key in the buffer (FIFO)
-                status = self._workflow_status_buffer.pop(wf_id, None)
-                if status is None:
-                    continue
-                exported_status[wf_id] = status
-                try:
-                    self.update_workflow_status(status, conn=c)
-                    exported += 1
-                except Exception as e:
-                    dbos_logger.error(f"Error while flushing status buffer: {e}")
-                    c.rollback()
-                    # Add the exported status back to the buffer, so they can be retried next time
-                    self._workflow_status_buffer.update(exported_status)
-                    break
-
-    def _flush_workflow_inputs_buffer(self) -> None:
-        if self._debug_mode:
-            raise Exception("called _flush_workflow_inputs_buffer in debug mode")
-
-        """Export the workflow inputs buffer to the database, up to the batch size."""
-        if len(self._workflow_inputs_buffer) == 0:
-            return
-
-        # Record exported inputs so far, and add them back on errors.
-        exported_inputs: Dict[str, str] = {}
-        with self.engine.begin() as c:
-            exported = 0
-            input_iter = iter(list(self._workflow_inputs_buffer))
-            wf_id: Optional[str] = None
-            while (
-                exported < _buffer_flush_batch_size
-                and (wf_id := next(input_iter, None)) is not None
-            ):
-                if wf_id not in self._exported_temp_txn_wf_status:
-                    # Skip exporting inputs if the status has not been exported yet
-                    continue
-                inputs = self._workflow_inputs_buffer.pop(wf_id, None)
-                if inputs is None:
-                    continue
-                exported_inputs[wf_id] = inputs
-                try:
-                    self.update_workflow_inputs(wf_id, inputs, conn=c)
-                    exported += 1
-                except Exception as e:
-                    dbos_logger.error(f"Error while flushing inputs buffer: {e}")
-                    c.rollback()
-                    # Add the exported inputs back to the buffer, so they can be retried next time
-                    self._workflow_inputs_buffer.update(exported_inputs)
-                    break
-
-    def flush_workflow_buffers(self) -> None:
-        """Flush the workflow status and inputs buffers periodically, via a background thread."""
-        while self._run_background_processes:
-            try:
-                self._is_flushing_status_buffer = True
-                # Must flush the status buffer first, as the inputs table has a foreign key constraint on the status table.
-                self._flush_workflow_status_buffer()
-                self._flush_workflow_inputs_buffer()
-                self._is_flushing_status_buffer = False
-                if self._is_buffers_empty:
-                    # Only sleep if both buffers are empty
-                    time.sleep(_buffer_flush_interval_secs)
-            except Exception as e:
-                dbos_logger.error(f"Error while flushing buffers: {e}")
-                time.sleep(_buffer_flush_interval_secs)
-                # Will retry next time
-
-    def buffer_workflow_status(self, status: WorkflowStatusInternal) -> None:
-        self._workflow_status_buffer[status["workflow_uuid"]] = status
-
-    def buffer_workflow_inputs(self, workflow_id: str, inputs: str) -> None:
-        # inputs is a serialized WorkflowInputs string
-        self._workflow_inputs_buffer[workflow_id] = inputs
-        self._temp_txn_wf_ids.add(workflow_id)
-
-    @property
-    def _is_buffers_empty(self) -> bool:
-        return (
-            len(self._workflow_status_buffer) == 0
-            and len(self._workflow_inputs_buffer) == 0
-        )
-
     def enqueue(self, workflow_id: str, queue_name: str) -> None:
         if self._debug_mode:
             raise Exception("called enqueue in debug mode")
@@ -1297,7 +1266,9 @@ class SystemDatabase:
                 .on_conflict_do_nothing()
             )
 
-    def start_queued_workflows(self, queue: "Queue", executor_id: str) -> List[str]:
+    def start_queued_workflows(
+        self, queue: "Queue", executor_id: str, app_version: str
+    ) -> List[str]:
         if self._debug_mode:
             return []
 
@@ -1412,26 +1383,36 @@ class SystemDatabase:
                         break
 
                 # To start a function, first set its status to PENDING and update its executor ID
-                c.execute(
+                res = c.execute(
                     SystemSchema.workflow_status.update()
                     .where(SystemSchema.workflow_status.c.workflow_uuid == id)
                     .where(
                         SystemSchema.workflow_status.c.status
                         == WorkflowStatusString.ENQUEUED.value
                     )
+                    .where(
+                        sa.or_(
+                            SystemSchema.workflow_status.c.application_version
+                            == app_version,
+                            SystemSchema.workflow_status.c.application_version.is_(
+                                None
+                            ),
+                        )
+                    )
                     .values(
                         status=WorkflowStatusString.PENDING.value,
+                        application_version=app_version,
                         executor_id=executor_id,
                     )
                 )
-
-                # Then give it a start time and assign the executor ID
-                c.execute(
-                    SystemSchema.workflow_queue.update()
-                    .where(SystemSchema.workflow_queue.c.workflow_uuid == id)
-                    .values(started_at_epoch_ms=start_time_ms)
-                )
-                ret_ids.append(id)
+                if res.rowcount > 0:
+                    # Then give it a start time and assign the executor ID
+                    c.execute(
+                        SystemSchema.workflow_queue.update()
+                        .where(SystemSchema.workflow_queue.c.workflow_uuid == id)
+                        .values(started_at_epoch_ms=start_time_ms)
+                    )
+                    ret_ids.append(id)
 
             # If we have a limiter, garbage-collect all completed functions started
             # before the period. If there's no limiter, there's no need--they were

dbos/_templates/dbos-db-starter/__package/main.py

@@ -6,6 +6,7 @@
 
 # First, let's do imports, create a FastAPI app, and initialize DBOS.
 
+import uvicorn
 from fastapi import FastAPI
 from fastapi.responses import HTMLResponse
 
@@ -37,7 +38,7 @@ def example_transaction(name: str) -> str:
     return greeting
 
 
-# Finally, let's use FastAPI to serve an HTML + CSS readme
+# Now, let's use FastAPI to serve an HTML + CSS readme
 # from the root path.
 
 
@@ -66,14 +67,8 @@ def readme() -> HTMLResponse:
     return HTMLResponse(readme)
 
 
-# To deploy this app to DBOS Cloud:
-# - "npm i -g @dbos-inc/dbos-cloud@latest" to install the Cloud CLI (requires Node)
-# - "dbos-cloud app deploy" to deploy your app
-# - Deploy outputs a URL--visit it to see your app!
+# Finally, we'll launch DBOS then start the FastAPI server.
 
-
-# To run this app locally:
-# - Make sure you have a Postgres database to connect to
-# - "dbos migrate" to set up your database tables
-# - "dbos start" to start the app
-# - Visit localhost:8000 to see your app!
+if __name__ == "__main__":
+    DBOS.launch()
+    uvicorn.run(app, host="0.0.0.0", port=8000)

dbos/_templates/dbos-db-starter/dbos-config.yaml.dbos

@@ -7,8 +7,6 @@ name: ${project_name}
 language: python
 runtimeConfig:
   start:
-    - "fastapi run ${package_name}/main.py"
+    - "${start_command}"
 database_url: ${DBOS_DATABASE_URL}
-database:
-  migrate:
-    - ${migration_command}
+${migration_section}