dbos 2.4.0a3__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

@@ -1,7 +1,6 @@
 import asyncio
 import json
 import time
-import uuid
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -19,7 +18,9 @@ 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
 
 if TYPE_CHECKING:
     from dbos._dbos import WorkflowHandle, WorkflowHandleAsync
@@ -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
 
@@ -158,6 +164,7 @@ class DBOSClient:
             engine=system_database_engine,
             schema=dbos_system_schema,
             serializer=serializer,
+            executor_id=None,
         )
         self._sys_db.check_connection()
         if application_database_url:
@@ -187,7 +194,7 @@ class DBOSClient:
             max_recovery_attempts = DEFAULT_MAX_RECOVERY_ATTEMPTS
         workflow_id = options.get("workflow_id")
         if workflow_id is None:
-            workflow_id = str(uuid.uuid4())
+            workflow_id = generate_uuid()
         workflow_timeout = options.get("workflow_timeout", None)
         enqueue_options_internal: EnqueueOptionsInternal = {
             "deduplication_id": options.get("deduplication_id"),
@@ -280,7 +287,7 @@ class DBOSClient:
         topic: Optional[str] = None,
         idempotency_key: Optional[str] = None,
     ) -> None:
-        idempotency_key = idempotency_key if idempotency_key else str(uuid.uuid4())
+        idempotency_key = idempotency_key if idempotency_key else generate_uuid()
         status: WorkflowStatusInternal = {
             "workflow_uuid": f"{destination_id}-{idempotency_key}",
             "status": WorkflowStatusString.SUCCESS.value,
@@ -471,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)
@@ -485,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,
@@ -502,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

@@ -2,7 +2,6 @@ import socket
 import threading
 import time
 import traceback
-import uuid
 from importlib.metadata import version
 from typing import TYPE_CHECKING, Optional
 
@@ -11,13 +10,12 @@ from websockets.sync.client import connect
 from websockets.sync.connection import Connection
 
 from dbos._context import SetWorkflowID
-from dbos._utils import GlobalParams
+from dbos._utils import GlobalParams, generate_uuid
 from dbos._workflow_commands import (
     garbage_collect,
     get_workflow,
     global_timeout,
     list_queued_workflows,
-    list_workflow_steps,
     list_workflows,
 )
 
@@ -118,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")
@@ -192,7 +192,7 @@ class ConductorWebsocket(threading.Thread):
                             fork_message = p.ForkWorkflowRequest.from_json(message)
                             new_workflow_id = fork_message.body["new_workflow_id"]
                             if new_workflow_id is None:
-                                new_workflow_id = str(uuid.uuid4())
+                                new_workflow_id = generate_uuid()
                             workflow_id = fork_message.body["workflow_id"]
                             start_step = fork_message.body["start_step"]
                             app_version = fork_message.body["application_version"]
@@ -223,23 +223,21 @@ class ConductorWebsocket(threading.Thread):
                             body = list_workflows_message.body
                             infos = []
                             try:
-                                load_input = body.get("load_input", False)
-                                load_output = body.get("load_output", False)
                                 infos = list_workflows(
                                     self.dbos._sys_db,
-                                    workflow_ids=body["workflow_uuids"],
-                                    user=body["authenticated_user"],
-                                    start_time=body["start_time"],
-                                    end_time=body["end_time"],
-                                    status=body["status"],
-                                    app_version=body["application_version"],
-                                    forked_from=body["forked_from"],
-                                    name=body["workflow_name"],
-                                    limit=body["limit"],
-                                    offset=body["offset"],
-                                    sort_desc=body["sort_desc"],
-                                    load_input=load_input,
-                                    load_output=load_output,
+                                    workflow_ids=body.get("workflow_uuids", None),
+                                    user=body.get("authenticated_user", None),
+                                    start_time=body.get("start_time", None),
+                                    end_time=body.get("end_time", None),
+                                    status=body.get("status", None),
+                                    app_version=body.get("application_version", None),
+                                    forked_from=body.get("forked_from", None),
+                                    name=body.get("workflow_name", None),
+                                    limit=body.get("limit", None),
+                                    offset=body.get("offset", None),
+                                    sort_desc=body.get("sort_desc", False),
+                                    load_input=body.get("load_input", False),
+                                    load_output=body.get("load_output", False),
                                 )
                             except Exception as e:
                                 error_message = f"Exception encountered when listing workflows: {traceback.format_exc()}"
@@ -262,19 +260,18 @@ class ConductorWebsocket(threading.Thread):
                             q_body = list_queued_workflows_message.body
                             infos = []
                             try:
-                                q_load_input = q_body.get("load_input", False)
                                 infos = list_queued_workflows(
                                     self.dbos._sys_db,
-                                    start_time=q_body["start_time"],
-                                    end_time=q_body["end_time"],
-                                    status=q_body["status"],
-                                    forked_from=q_body["forked_from"],
-                                    name=q_body["workflow_name"],
-                                    limit=q_body["limit"],
-                                    offset=q_body["offset"],
-                                    queue_name=q_body["queue_name"],
-                                    sort_desc=q_body["sort_desc"],
-                                    load_input=q_load_input,
+                                    start_time=q_body.get("start_time", None),
+                                    end_time=q_body.get("end_time", None),
+                                    status=q_body.get("status", None),
+                                    forked_from=q_body.get("forked_from", None),
+                                    name=q_body.get("workflow_name", None),
+                                    limit=q_body.get("limit", None),
+                                    offset=q_body.get("offset", None),
+                                    queue_name=q_body.get("queue_name", None),
+                                    sort_desc=q_body.get("sort_desc", False),
+                                    load_input=q_body.get("load_input", False),
                                 )
                             except Exception as e:
                                 error_message = f"Exception encountered when listing queued workflows: {traceback.format_exc()}"
@@ -343,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()}"
@@ -401,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/_context.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import json
 import os
-import uuid
 from contextlib import AbstractContextManager
 from contextvars import ContextVar
 from dataclasses import dataclass
@@ -15,7 +14,7 @@ if TYPE_CHECKING:
 
 from sqlalchemy.orm import Session
 
-from dbos._utils import GlobalParams
+from dbos._utils import GlobalParams, generate_uuid
 
 from ._logger import dbos_logger
 from ._tracer import dbos_tracer
@@ -151,7 +150,7 @@ class DBOSContext:
                 self.logger.warning(
                     f"Multiple workflows started in the same SetWorkflowID block. Only the first workflow is assigned the specified workflow ID; subsequent workflows will use a generated workflow ID."
                 )
-            wfid = str(uuid.uuid4())
+            wfid = generate_uuid()
         return wfid
 
     def start_workflow(