dbos 1.1.0a3__py3-none-any.whl → 1.2.0__py3-none-any.whl

dbos/_admin_server.py

@@ -5,8 +5,9 @@ import re
 import threading
 from functools import partial
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
-from typing import TYPE_CHECKING, Any, List, TypedDict
+from typing import TYPE_CHECKING, Any, List, Optional, TypedDict
 
+from ._context import SetWorkflowID
 from ._error import DBOSException
 from ._logger import dbos_logger
 from ._recovery import recover_pending_workflows
@@ -141,7 +142,11 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
                 try:
                     data = json.loads(post_data.decode("utf-8"))
                     start_step: int = data.get("start_step", 1)
-                    self._handle_fork(workflow_id, start_step)
+                    new_workflow_id: Optional[str] = data.get("new_workflow_id")
+                    application_version: Optional[str] = data.get("application_version")
+                    self._handle_fork(
+                        workflow_id, start_step, new_workflow_id, application_version
+                    )
                 except (json.JSONDecodeError, AttributeError) as e:
                     self.send_response(500)
                     self.send_header("Content-Type", "application/json")
@@ -191,9 +196,24 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
             self.end_headers()
             self.wfile.write(response_body)
 
-    def _handle_fork(self, workflow_id: str, start_step: int) -> None:
+    def _handle_fork(
+        self,
+        workflow_id: str,
+        start_step: int,
+        new_workflow_id: Optional[str],
+        application_version: Optional[str],
+    ) -> None:
         try:
-            handle = self.dbos.fork_workflow(workflow_id, start_step)
+            print(f"Forking workflow {workflow_id} from step {start_step}")
+            if new_workflow_id is not None:
+                with SetWorkflowID(new_workflow_id):
+                    handle = self.dbos.fork_workflow(
+                        workflow_id, start_step, application_version=application_version
+                    )
+            else:
+                handle = self.dbos.fork_workflow(
+                    workflow_id, start_step, application_version=application_version
+                )
             response_body = json.dumps(
                 {
                     "workflow_id": handle.workflow_id,

dbos/_app_db.py

@@ -216,21 +216,6 @@ class ApplicationDatabase:
             for row in rows
         ]
 
-    def get_max_function_id(self, workflow_uuid: str) -> Optional[int]:
-        with self.engine.begin() as conn:
-            max_function_id_row = conn.execute(
-                sa.select(
-                    sa.func.max(ApplicationSchema.transaction_outputs.c.function_id)
-                ).where(
-                    ApplicationSchema.transaction_outputs.c.workflow_uuid
-                    == workflow_uuid
-                )
-            ).fetchone()
-
-            max_function_id = max_function_id_row[0] if max_function_id_row else None
-
-            return max_function_id
-
     def clone_workflow_transactions(
         self, src_workflow_id: str, forked_workflow_id: str, start_step: int
     ) -> None:

dbos/_client.py

@@ -3,8 +3,6 @@ import sys
 import uuid
 from typing import Any, Generic, List, Optional, TypedDict, TypeVar
 
-from sqlalchemy import URL
-
 from dbos._app_db import ApplicationDatabase
 from dbos._context import MaxPriority, MinPriority
 
@@ -15,6 +13,7 @@ else:
 
 from dbos import _serialization
 from dbos._dbos import WorkflowHandle, WorkflowHandleAsync
+from dbos._dbos_config import is_valid_database_url
 from dbos._error import DBOSException, DBOSNonExistentWorkflowError
 from dbos._registrations import DEFAULT_MAX_RECOVERY_ATTEMPTS
 from dbos._serialization import WorkflowInputs
@@ -99,6 +98,7 @@ class WorkflowHandleClientAsyncPolling(Generic[R]):
 
 class DBOSClient:
     def __init__(self, database_url: str, *, system_database: Optional[str] = None):
+        assert is_valid_database_url(database_url)
         # We only create database connections but do not run migrations
         self._sys_db = SystemDatabase(
             database_url=database_url,
@@ -109,6 +109,7 @@ class DBOSClient:
             },
             sys_db_name=system_database,
         )
+        self._sys_db.check_connection()
         self._app_db = ApplicationDatabase(
             database_url=database_url,
             engine_kwargs={
@@ -231,7 +232,7 @@ class DBOSClient:
             "workflow_deadline_epoch_ms": None,
         }
         with self._sys_db.engine.begin() as conn:
-            self._sys_db.insert_workflow_status(
+            self._sys_db._insert_workflow_status(
                 status, conn, max_recovery_attempts=None
             )
         self._sys_db.send(status["workflow_uuid"], 0, destination_id, message, topic)

dbos/_context.py

@@ -392,6 +392,7 @@ class SetWorkflowTimeout:
             else None
         )
         self.saved_workflow_timeout: Optional[int] = None
+        self.saved_workflow_deadline_epoch_ms: Optional[int] = None
 
     def __enter__(self) -> SetWorkflowTimeout:
         # Code to create a basic context
@@ -402,6 +403,8 @@ class SetWorkflowTimeout:
         ctx = assert_current_dbos_context()
         self.saved_workflow_timeout = ctx.workflow_timeout_ms
         ctx.workflow_timeout_ms = self.workflow_timeout_ms
+        self.saved_workflow_deadline_epoch_ms = ctx.workflow_deadline_epoch_ms
+        ctx.workflow_deadline_epoch_ms = None
         return self
 
     def __exit__(
@@ -411,6 +414,9 @@ class SetWorkflowTimeout:
         traceback: Optional[TracebackType],
     ) -> Literal[False]:
         assert_current_dbos_context().workflow_timeout_ms = self.saved_workflow_timeout
+        assert_current_dbos_context().workflow_deadline_epoch_ms = (
+            self.saved_workflow_deadline_epoch_ms
+        )
         # Code to clean up the basic context if we created it
         if self.created_ctx:
             _clear_local_dbos_context()

dbos/_core.py

@@ -20,8 +20,10 @@ from typing import (
     cast,
 )
 
+import psycopg
+
 from dbos._outcome import Immediate, NoResult, Outcome, Pending
-from dbos._utils import GlobalParams
+from dbos._utils import GlobalParams, retriable_postgres_exception
 
 from ._app_db import ApplicationDatabase, TransactionResultInternal
 
@@ -602,7 +604,6 @@ async def start_workflow_async(
     *args: P.args,
     **kwargs: P.kwargs,
 ) -> "WorkflowHandleAsync[R]":
-
     # If the function has a class, add the class object as its first argument
     fself: Optional[object] = None
     if hasattr(func, "__self__"):
@@ -932,12 +933,18 @@ def decorate_transaction(
                                 )
                                 break
                         except DBAPIError as dbapi_error:
-                            if dbapi_error.orig.sqlstate == "40001":  # type: ignore
+                            driver_error = cast(
+                                Optional[psycopg.OperationalError], dbapi_error.orig
+                            )
+                            if retriable_postgres_exception(dbapi_error) or (
+                                driver_error is not None
+                                and driver_error.sqlstate == "40001"
+                            ):
                                 # Retry on serialization failure
                                 span = ctx.get_current_span()
                                 if span:
                                     span.add_event(
-                                        "Transaction Serialization Failure",
+                                        "Transaction Failure",
                                         {"retry_wait_seconds": retry_wait_seconds},
                                     )
                                 time.sleep(retry_wait_seconds)

dbos/_dbos.py

@@ -90,10 +90,8 @@ from ._context import (
 from ._dbos_config import (
     ConfigFile,
     DBOSConfig,
-    check_config_consistency,
     overwrite_config,
     process_config,
-    set_env_vars,
     translate_dbos_config_to_config_file,
 )
 from ._error import (
@@ -299,6 +297,7 @@ class DBOS:
 
         self._launched: bool = False
         self._debug_mode: bool = False
+        self._configured_threadpool: bool = False
         self._sys_db_field: Optional[SystemDatabase] = None
         self._app_db_field: Optional[ApplicationDatabase] = None
         self._registry: DBOSRegistry = _get_or_create_dbos_registry()
@@ -323,14 +322,12 @@ class DBOS:
         unvalidated_config = translate_dbos_config_to_config_file(config)
         if os.environ.get("DBOS__CLOUD") == "true":
             unvalidated_config = overwrite_config(unvalidated_config)
-        check_config_consistency(name=unvalidated_config["name"])
 
         if unvalidated_config is not None:
             self._config: ConfigFile = process_config(data=unvalidated_config)
         else:
             raise ValueError("No valid configuration was loaded.")
 
-        set_env_vars(self._config)
         config_logger(self._config)
         dbos_tracer.config(self._config)
         dbos_logger.info(f"Initializing DBOS (v{GlobalParams.dbos_version})")
@@ -719,6 +716,7 @@ class DBOS:
         **kwargs: P.kwargs,
     ) -> WorkflowHandleAsync[R]:
         """Invoke a workflow function on the event loop, returning a handle to the ongoing execution."""
+        await cls._configure_asyncio_thread_pool()
         return await start_workflow_async(
             _get_dbos_instance(), func, None, True, *args, **kwargs
         )
@@ -736,6 +734,7 @@ class DBOS:
     async def get_workflow_status_async(
         cls, workflow_id: str
     ) -> Optional[WorkflowStatus]:
+        await cls._configure_asyncio_thread_pool()
         """Return the status of a workflow execution."""
         return await asyncio.to_thread(cls.get_workflow_status, workflow_id)
 
@@ -757,6 +756,7 @@ class DBOS:
     ) -> WorkflowHandleAsync[R]:
         """Return a `WorkflowHandle` for a workflow execution."""
         dbos = _get_dbos_instance()
+        await cls._configure_asyncio_thread_pool()
         if existing_workflow:
             stat = await dbos.get_workflow_status_async(workflow_id)
             if stat is None:
@@ -775,6 +775,7 @@ class DBOS:
         cls, destination_id: str, message: Any, topic: Optional[str] = None
     ) -> None:
         """Send a message to a workflow execution."""
+        await cls._configure_asyncio_thread_pool()
         await asyncio.to_thread(lambda: DBOS.send(destination_id, message, topic))
 
     @classmethod
@@ -797,6 +798,7 @@ class DBOS:
         This function is to be called from within a workflow.
         `recv_async` will return the message sent on `topic`, asyncronously waiting if necessary.
         """
+        await cls._configure_asyncio_thread_pool()
         return await asyncio.to_thread(lambda: DBOS.recv(topic, timeout_seconds))
 
     @classmethod
@@ -835,6 +837,7 @@ class DBOS:
         It is important to use `DBOS.sleep` or `DBOS.sleep_async` (as opposed to any other sleep) within workflows,
         as the DBOS sleep methods are durable and completed sleeps will be skipped during recovery.
         """
+        await cls._configure_asyncio_thread_pool()
         await asyncio.to_thread(lambda: DBOS.sleep(seconds))
 
     @classmethod
@@ -869,6 +872,7 @@ class DBOS:
             value(Any): A serializable value to associate with the key
 
         """
+        await cls._configure_asyncio_thread_pool()
         await asyncio.to_thread(lambda: DBOS.set_event(key, value))
 
     @classmethod
@@ -901,6 +905,7 @@ class DBOS:
             timeout_seconds(float): The amount of time to wait, in case `set_event` has not yet been called byt the workflow
 
         """
+        await cls._configure_asyncio_thread_pool()
         return await asyncio.to_thread(
             lambda: DBOS.get_event(workflow_id, key, timeout_seconds)
         )
@@ -929,6 +934,19 @@ class DBOS:
             fn, "DBOS.cancelWorkflow"
         )
 
+    @classmethod
+    async def _configure_asyncio_thread_pool(cls) -> None:
+        """
+        Configure the thread pool for asyncio.to_thread.
+
+        This function is called before the first call to asyncio.to_thread.
+        """
+        if _get_dbos_instance()._configured_threadpool:
+            return
+        loop = asyncio.get_running_loop()
+        loop.set_default_executor(_get_dbos_instance()._executor)
+        _get_dbos_instance()._configured_threadpool = True
+
     @classmethod
     def resume_workflow(cls, workflow_id: str) -> WorkflowHandle[Any]:
         """Resume a workflow by ID."""

dbos/_dbos_config.py

@@ -329,17 +329,9 @@ def process_config(
     if data.get("database_url") is not None and data["database_url"] != "":
         # Parse the db string and check required fields
         assert data["database_url"] is not None
+        assert is_valid_database_url(data["database_url"])
+
         url = make_url(data["database_url"])
-        required_fields = [
-            ("username", "Username must be specified in the connection URL"),
-            ("password", "Password must be specified in the connection URL"),
-            ("host", "Host must be specified in the connection URL"),
-            ("database", "Database name must be specified in the connection URL"),
-        ]
-        for field_name, error_message in required_fields:
-            field_value = getattr(url, field_name, None)
-            if not field_value:
-                raise DBOSInitializationError(error_message)
 
         if not data["database"].get("sys_db_name"):
             assert url.database is not None
@@ -385,6 +377,9 @@ def process_config(
     if not silent and logs["logLevel"] == "INFO" or logs["logLevel"] == "DEBUG":
         log_url = make_url(data["database_url"]).render_as_string(hide_password=True)
         print(f"[bold blue]Using database connection string: {log_url}[/bold blue]")
+        print(
+            f"[bold blue]Database engine parameters: {data['database']['db_engine_kwargs']}[/bold blue]"
+        )
 
     # Return data as ConfigFile type
     return data
@@ -407,6 +402,7 @@ def configure_db_engine_parameters(
         "pool_timeout": 30,
         "max_overflow": 0,
         "pool_size": 20,
+        "pool_pre_ping": True,
     }
     # If user-provided kwargs are present, use them instead
     user_kwargs = data.get("db_engine_kwargs")
@@ -431,6 +427,21 @@ def configure_db_engine_parameters(
     data["sys_db_engine_kwargs"] = system_engine_kwargs
 
 
+def is_valid_database_url(database_url: str) -> bool:
+    url = make_url(database_url)
+    required_fields = [
+        ("username", "Username must be specified in the connection URL"),
+        ("password", "Password must be specified in the connection URL"),
+        ("host", "Host must be specified in the connection URL"),
+        ("database", "Database name must be specified in the connection URL"),
+    ]
+    for field_name, error_message in required_fields:
+        field_value = getattr(url, field_name, None)
+        if not field_value:
+            raise DBOSInitializationError(error_message)
+    return True
+
+
 def _is_valid_app_name(name: str) -> bool:
     name_len = len(name)
     if name_len < 3 or name_len > 30:
@@ -444,12 +455,6 @@ def _app_name_to_db_name(app_name: str) -> str:
     return name if not name[0].isdigit() else f"_{name}"
 
 
-def set_env_vars(config: ConfigFile) -> None:
-    for env, value in config.get("env", {}).items():
-        if value is not None:
-            os.environ[env] = str(value)
-
-
 def overwrite_config(provided_config: ConfigFile) -> ConfigFile:
     # Load the DBOS configuration file and force the use of:
     # 1. The database url provided by DBOS_DATABASE_URL
@@ -529,26 +534,3 @@ def overwrite_config(provided_config: ConfigFile) -> ConfigFile:
         del provided_config["env"]
 
     return provided_config
-
-
-def check_config_consistency(
-    *,
-    name: str,
-    config_file_path: str = DBOS_CONFIG_PATH,
-) -> None:
-    # First load the config file and check whether it is present
-    try:
-        config = load_config(config_file_path, silent=True, run_process_config=False)
-    except FileNotFoundError:
-        dbos_logger.debug(
-            f"No configuration file {config_file_path} found. Skipping consistency check with provided config."
-        )
-        return
-    except Exception as e:
-        raise e
-
-    # Check the name
-    if name != config["name"]:
-        raise DBOSInitializationError(
-            f"Provided app name '{name}' does not match the app name '{config['name']}' in {config_file_path}."
-        )

dbos/_error.py

@@ -62,6 +62,7 @@ class DBOSErrorCode(Enum):
     WorkflowCancelled = 10
     UnexpectedStep = 11
     QueueDeduplicated = 12
+    AwaitedWorkflowCancelled = 13
     ConflictingRegistrationError = 25
 
 
@@ -206,6 +207,19 @@ class DBOSQueueDeduplicatedError(DBOSException):
         )
 
 
+class DBOSAwaitedWorkflowCancelledError(DBOSException):
+    def __init__(self, workflow_id: str):
+        self.workflow_id = workflow_id
+        super().__init__(
+            f"Awaited workflow {workflow_id} was cancelled",
+            dbos_error_code=DBOSErrorCode.AwaitedWorkflowCancelled.value,
+        )
+
+    def __reduce__(self) -> Any:
+        # Tell jsonpickle how to reconstruct this object
+        return (self.__class__, (self.workflow_id,))
+
+
 #######################################
 ## BaseException
 #######################################

dbos/_event_loop.py

@@ -1,5 +1,6 @@
 import asyncio
 import threading
+from concurrent.futures import ThreadPoolExecutor
 from typing import Any, Coroutine, Optional, TypeVar
 
 
@@ -33,15 +34,17 @@ class BackgroundEventLoop:
 
     def _run_event_loop(self) -> None:
         self._loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(self._loop)
+        with ThreadPoolExecutor(max_workers=64) as thread_pool:
+            self._loop.set_default_executor(thread_pool)
+            asyncio.set_event_loop(self._loop)
 
-        self._running = True
-        self._ready.set()  # Signal that the loop is ready
+            self._running = True
+            self._ready.set()  # Signal that the loop is ready
 
-        try:
-            self._loop.run_forever()
-        finally:
-            self._loop.close()
+            try:
+                self._loop.run_forever()
+            finally:
+                self._loop.close()
 
     async def _shutdown(self) -> None:
         if self._loop is None:

dbos/_queue.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Callable, Coroutine, Optional, TypedDict
 from psycopg import errors
 from sqlalchemy.exc import OperationalError
 
+from dbos._logger import dbos_logger
 from dbos._utils import GlobalParams
 
 from ._core import P, R, execute_workflow_by_id, start_workflow, start_workflow_async
@@ -56,6 +57,8 @@ class Queue:
         from ._dbos import _get_or_create_dbos_registry
 
         registry = _get_or_create_dbos_registry()
+        if self.name in registry.queue_info_map:
+            dbos_logger.warning(f"Queue {name} has already been declared")
         registry.queue_info_map[self.name] = self
 
     def enqueue(
@@ -95,12 +98,8 @@ def queue_thread(stop_event: threading.Event, dbos: "DBOS") -> None:
                 if not isinstance(
                     e.orig, (errors.SerializationFailure, errors.LockNotAvailable)
                 ):
-                    dbos.logger.warning(
-                        f"Exception encountered in queue thread: {traceback.format_exc()}"
-                    )
-            except Exception:
+                    dbos.logger.warning(f"Exception encountered in queue thread: {e}")
+            except Exception as e:
                 if not stop_event.is_set():
                     # Only print the error if the thread is not stopping
-                    dbos.logger.warning(
-                        f"Exception encountered in queue thread: {traceback.format_exc()}"
-                    )
+                    dbos.logger.warning(f"Exception encountered in queue thread: {e}")