dbos 0.26.1__py3-none-any.whl → 0.27.0__py3-none-any.whl

dbos/_serialization.py

@@ -1,8 +1,10 @@
 import types
-from typing import Any, Dict, Tuple, TypedDict
+from typing import Any, Dict, Optional, Tuple, TypedDict
 
 import jsonpickle  # type: ignore
 
+from ._logger import dbos_logger
+
 
 class WorkflowInputs(TypedDict):
     args: Tuple[Any, ...]
@@ -51,5 +53,54 @@ def deserialize_args(serialized_data: str) -> WorkflowInputs:
 
 def deserialize_exception(serialized_data: str) -> Exception:
     """Deserialize JSON string back to a Python Exception using jsonpickle."""
-    upo: Exception = jsonpickle.decode(serialized_data)
-    return upo
+    exc: Exception = jsonpickle.decode(serialized_data)
+    return exc
+
+
+def safe_deserialize(
+    workflow_id: str,
+    *,
+    serialized_input: Optional[str],
+    serialized_output: Optional[str],
+    serialized_exception: Optional[str],
+) -> tuple[Optional[WorkflowInputs], Optional[Any], Optional[Exception]]:
+    """
+    This function safely deserializes a workflow's recorded input and output/exception.
+    If any of them is not deserializable, it logs a warning and returns a string instead of throwing an exception.
+
+    This function is used in workflow introspection methods (get_workflows and get_queued_workflow)
+    to ensure errors related to nondeserializable objects are observable.
+    """
+    input: Optional[WorkflowInputs]
+    try:
+        input = (
+            deserialize_args(serialized_input) if serialized_input is not None else None
+        )
+    except Exception as e:
+        dbos_logger.warning(
+            f"Warning: input object could not be deserialized for workflow {workflow_id}, returning as string: {e}"
+        )
+        input = serialized_input  # type: ignore
+    output: Optional[Any]
+    try:
+        output = (
+            deserialize(serialized_output) if serialized_output is not None else None
+        )
+    except Exception as e:
+        dbos_logger.warning(
+            f"Warning: output object could not be deserialized for workflow {workflow_id}, returning as string: {e}"
+        )
+        output = serialized_output
+    exception: Optional[Exception]
+    try:
+        exception = (
+            deserialize_exception(serialized_exception)
+            if serialized_exception is not None
+            else None
+        )
+    except Exception as e:
+        dbos_logger.warning(
+            f"Warning: exception object could not be deserialized for workflow {workflow_id}, returning as string: {e}"
+        )
+        exception = serialized_exception  # type: ignore
+    return input, output, exception

dbos/_sys_db.py

@@ -37,6 +37,7 @@ from ._error import (
     DBOSConflictingWorkflowError,
     DBOSDeadLetterQueueError,
     DBOSNonExistentWorkflowError,
+    DBOSQueueDeduplicatedError,
     DBOSUnexpectedStepError,
     DBOSWorkflowCancelledError,
     DBOSWorkflowConflictIDError,
@@ -135,6 +136,10 @@ class WorkflowStatusInternal(TypedDict):
     workflow_deadline_epoch_ms: Optional[int]
 
 
+class EnqueueOptionsInternal(TypedDict):
+    deduplication_id: Optional[str]  # Unique ID for deduplication on a queue
+
+
 class RecordedResult(TypedDict):
     output: Optional[str]  # JSON (jsonpickle)
     error: Optional[str]  # JSON (jsonpickle)
@@ -248,6 +253,7 @@ class SystemDatabase:
                     sa.text("SELECT 1 FROM pg_database WHERE datname=:db_name"),
                     parameters={"db_name": sysdb_name},
                 ).scalar():
+                    dbos_logger.info(f"Creating system database {sysdb_name}")
                     conn.execute(sa.text(f"CREATE DATABASE {sysdb_name}"))
             engine.dispose()
 
@@ -267,12 +273,23 @@ class SystemDatabase:
         if pool_size is None:
             pool_size = 20
 
+        engine_kwargs = database.get("db_engine_kwargs")
+        if engine_kwargs is None:
+            engine_kwargs = {}
+
+        # Respect user-provided values. Otherwise, set defaults.
+        if "pool_size" not in engine_kwargs:
+            engine_kwargs["pool_size"] = pool_size
+        if "max_overflow" not in engine_kwargs:
+            engine_kwargs["max_overflow"] = 0
+        if "pool_timeout" not in engine_kwargs:
+            engine_kwargs["pool_timeout"] = 30
+        if "connect_args" not in engine_kwargs:
+            engine_kwargs["connect_args"] = {"connect_timeout": 10}
+
         self.engine = sa.create_engine(
             system_db_url,
-            pool_size=pool_size,
-            max_overflow=0,
-            pool_timeout=30,
-            connect_args={"connect_timeout": 10},
+            **engine_kwargs,
         )
 
         # Run a schema migration for the system database
@@ -378,7 +395,6 @@ class SystemDatabase:
         cmd = cmd.returning(SystemSchema.workflow_status.c.recovery_attempts, SystemSchema.workflow_status.c.status, SystemSchema.workflow_status.c.workflow_deadline_epoch_ms, SystemSchema.workflow_status.c.name, SystemSchema.workflow_status.c.class_name, SystemSchema.workflow_status.c.config_name, SystemSchema.workflow_status.c.queue_name)  # type: ignore
 
         results = conn.execute(cmd)
-
         row = results.fetchone()
         if row is not None:
             # Check the started workflow matches the expected name, class_name, config_name, and queue_name
@@ -528,15 +544,17 @@ class SystemDatabase:
             # Execute with snapshot isolation in case of concurrent calls on the same workflow
             c.execute(sa.text("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ"))
             # Check the status of the workflow. If it is complete, do nothing.
-            row = c.execute(
+            status_row = c.execute(
                 sa.select(
                     SystemSchema.workflow_status.c.status,
                 ).where(SystemSchema.workflow_status.c.workflow_uuid == workflow_id)
             ).fetchone()
+            if status_row is None:
+                return
+            status = status_row[0]
             if (
-                row is None
-                or row[0] == WorkflowStatusString.SUCCESS.value
-                or row[0] == WorkflowStatusString.ERROR.value
+                status == WorkflowStatusString.SUCCESS.value
+                or status == WorkflowStatusString.ERROR.value
             ):
                 return
             # Remove the workflow from the queues table so resume can safely be called on an ENQUEUED workflow
@@ -576,7 +594,12 @@ class SystemDatabase:
             return max_function_id
 
     def fork_workflow(
-        self, original_workflow_id: str, forked_workflow_id: str, start_step: int = 1
+        self,
+        original_workflow_id: str,
+        forked_workflow_id: str,
+        start_step: int,
+        *,
+        application_version: Optional[str],
     ) -> str:
 
         status = self.get_workflow_status(original_workflow_id)
@@ -596,7 +619,11 @@ class SystemDatabase:
                     name=status["name"],
                     class_name=status["class_name"],
                     config_name=status["config_name"],
-                    application_version=status["app_version"],
+                    application_version=(
+                        application_version
+                        if application_version is not None
+                        else status["app_version"]
+                    ),
                     application_id=status["app_id"],
                     request=status["request"],
                     authenticated_user=status["authenticated_user"],
@@ -874,13 +901,15 @@ class SystemDatabase:
             info.app_version = row[14]
             info.app_id = row[15]
 
-            inputs = _serialization.deserialize_args(row[16])
-            if inputs is not None:
-                info.input = inputs
-            if info.status == WorkflowStatusString.SUCCESS.value:
-                info.output = _serialization.deserialize(row[17])
-            elif info.status == WorkflowStatusString.ERROR.value:
-                info.error = _serialization.deserialize_exception(row[18])
+            inputs, output, exception = _serialization.safe_deserialize(
+                info.workflow_id,
+                serialized_input=row[16],
+                serialized_output=row[17],
+                serialized_exception=row[18],
+            )
+            info.input = inputs
+            info.output = output
+            info.error = exception
 
             infos.append(info)
         return infos
@@ -980,13 +1009,15 @@ class SystemDatabase:
             info.app_version = row[14]
             info.app_id = row[15]
 
-            inputs = _serialization.deserialize_args(row[16])
-            if inputs is not None:
-                info.input = inputs
-            if info.status == WorkflowStatusString.SUCCESS.value:
-                info.output = _serialization.deserialize(row[17])
-            elif info.status == WorkflowStatusString.ERROR.value:
-                info.error = _serialization.deserialize_exception(row[18])
+            inputs, output, exception = _serialization.safe_deserialize(
+                info.workflow_id,
+                serialized_input=row[16],
+                serialized_output=row[17],
+                serialized_exception=row[18],
+            )
+            info.input = inputs
+            info.output = output
+            info.error = exception
 
             infos.append(info)
 
@@ -1586,17 +1617,43 @@ class SystemDatabase:
             )
         return value
 
-    def enqueue(self, workflow_id: str, queue_name: str, conn: sa.Connection) -> None:
+    def enqueue(
+        self,
+        workflow_id: str,
+        queue_name: str,
+        conn: sa.Connection,
+        *,
+        enqueue_options: Optional[EnqueueOptionsInternal],
+    ) -> None:
         if self._debug_mode:
             raise Exception("called enqueue in debug mode")
-        conn.execute(
-            pg.insert(SystemSchema.workflow_queue)
-            .values(
-                workflow_uuid=workflow_id,
-                queue_name=queue_name,
+        try:
+            deduplication_id = (
+                enqueue_options["deduplication_id"]
+                if enqueue_options is not None
+                else None
             )
-            .on_conflict_do_nothing()
-        )
+            query = (
+                pg.insert(SystemSchema.workflow_queue)
+                .values(
+                    workflow_uuid=workflow_id,
+                    queue_name=queue_name,
+                    deduplication_id=deduplication_id,
+                )
+                .on_conflict_do_nothing(
+                    index_elements=SystemSchema.workflow_queue.primary_key.columns
+                )
+            )  # Ignore primary key constraint violation
+            conn.execute(query)
+        except DBAPIError as dbapi_error:
+            # Unique constraint violation for the deduplication ID
+            if dbapi_error.orig.sqlstate == "23505":  # type: ignore
+                assert (
+                    deduplication_id is not None
+                ), f"deduplication_id should not be None. Workflow ID: {workflow_id}, Queue name: {queue_name}."
+                raise DBOSQueueDeduplicatedError(
+                    workflow_id, queue_name, deduplication_id
+                )
 
     def start_queued_workflows(
         self, queue: "Queue", executor_id: str, app_version: str
@@ -1868,6 +1925,7 @@ class SystemDatabase:
         inputs: str,
         *,
         max_recovery_attempts: Optional[int],
+        enqueue_options: Optional[EnqueueOptionsInternal],
     ) -> tuple[WorkflowStatuses, Optional[int]]:
         """
         Synchronously record the status and inputs for workflows in a single transaction
@@ -1883,24 +1941,16 @@ class SystemDatabase:
                 status["queue_name"] is not None
                 and wf_status == WorkflowStatusString.ENQUEUED.value
             ):
-                self.enqueue(status["workflow_uuid"], status["queue_name"], conn)
+                self.enqueue(
+                    status["workflow_uuid"],
+                    status["queue_name"],
+                    conn,
+                    enqueue_options=enqueue_options,
+                )
         return wf_status, workflow_deadline_epoch_ms
 
 
-def reset_system_database(config: ConfigFile) -> None:
-    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
-    )
-    postgres_db_url = sa.URL.create(
-        "postgresql+psycopg",
-        username=config["database"]["username"],
-        password=config["database"]["password"],
-        host=config["database"]["hostname"],
-        port=config["database"]["port"],
-        database="postgres",
-    )
+def reset_system_database(postgres_db_url: sa.URL, sysdb_name: str) -> None:
     try:
         # Connect to postgres default database
         engine = sa.create_engine(postgres_db_url)

dbos/_tracer.py

@@ -3,8 +3,10 @@ from typing import TYPE_CHECKING, Optional
 
 from opentelemetry import trace
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
+from opentelemetry.semconv.resource import ResourceAttributes
 from opentelemetry.trace import Span
 
 from dbos._utils import GlobalParams
@@ -23,7 +25,13 @@ class DBOSTracer:
 
     def config(self, config: ConfigFile) -> None:
         if not isinstance(trace.get_tracer_provider(), TracerProvider):
-            provider = TracerProvider()
+            resource = Resource(
+                attributes={
+                    ResourceAttributes.SERVICE_NAME: config["name"],
+                }
+            )
+
+            provider = TracerProvider(resource=resource)
             if os.environ.get("DBOS__CONSOLE_TRACES", None) is not None:
                 processor = BatchSpanProcessor(ConsoleSpanExporter())
                 provider.add_span_processor(processor)

dbos/_workflow_commands.py

@@ -1,6 +1,7 @@
 import uuid
 from typing import List, Optional
 
+from dbos._context import get_local_dbos_context
 from dbos._error import DBOSException
 
 from ._app_db import ApplicationDatabase
@@ -103,6 +104,8 @@ def fork_workflow(
     app_db: ApplicationDatabase,
     workflow_id: str,
     start_step: int,
+    *,
+    application_version: Optional[str],
 ) -> str:
     def get_max_function_id(workflow_uuid: str) -> int:
         max_transactions = app_db.get_max_function_id(workflow_uuid) or 0
@@ -114,7 +117,17 @@ def fork_workflow(
         raise DBOSException(
             f"Cannot fork workflow {workflow_id} from step {start_step}. The workflow has {max_function_id} steps."
         )
-    forked_workflow_id = str(uuid.uuid4())
+    ctx = get_local_dbos_context()
+    if ctx is not None and len(ctx.id_assigned_for_next_workflow) > 0:
+        forked_workflow_id = ctx.id_assigned_for_next_workflow
+        ctx.id_assigned_for_next_workflow = ""
+    else:
+        forked_workflow_id = str(uuid.uuid4())
     app_db.clone_workflow_transactions(workflow_id, forked_workflow_id, start_step)
-    sys_db.fork_workflow(workflow_id, forked_workflow_id, start_step)
+    sys_db.fork_workflow(
+        workflow_id,
+        forked_workflow_id,
+        start_step,
+        application_version=application_version,
+    )
     return forked_workflow_id