dbos 1.2.0a2__py3-none-any.whl → 1.2.0a5__py3-none-any.whl

dbos/_client.py

@@ -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/_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
 
@@ -931,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,7 +90,6 @@ from ._context import (
 from ._dbos_config import (
     ConfigFile,
     DBOSConfig,
-    check_config_consistency,
     overwrite_config,
     process_config,
     set_env_vars,
@@ -324,7 +323,6 @@ 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)

dbos/_dbos_config.py

@@ -407,6 +407,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")
@@ -529,26 +530,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/_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}")

dbos/_sys_db.py

@@ -1,7 +1,9 @@
 import datetime
+import functools
 import json
 import logging
 import os
+import random
 import re
 import threading
 import time
@@ -17,6 +19,7 @@ from typing import (
     Sequence,
     TypedDict,
     TypeVar,
+    cast,
 )
 
 import psycopg
@@ -27,7 +30,7 @@ from alembic.config import Config
 from sqlalchemy.exc import DBAPIError
 from sqlalchemy.sql import func
 
-from dbos._utils import INTERNAL_QUEUE_NAME
+from dbos._utils import INTERNAL_QUEUE_NAME, retriable_postgres_exception
 
 from . import _serialization
 from ._context import get_local_dbos_context
@@ -268,6 +271,51 @@ class ThreadSafeConditionDict:
                 dbos_logger.warning(f"Key {key} not found in condition dictionary.")
 
 
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def db_retry(
+    initial_backoff: float = 1.0, max_backoff: float = 60.0
+) -> Callable[[F], F]:
+    """
+    If a workflow encounters a database connection issue while performing an operation,
+    block the workflow and retry the operation until it reconnects and succeeds.
+
+    In other words, if DBOS loses its database connection, everything pauses until the connection is recovered,
+    trading off availability for correctness.
+    """
+
+    def decorator(func: F) -> F:
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            retries: int = 0
+            backoff: float = initial_backoff
+            while True:
+                try:
+                    return func(*args, **kwargs)
+                except DBAPIError as e:
+
+                    # Determine if this is a retriable exception
+                    if not retriable_postgres_exception(e):
+                        raise
+
+                    retries += 1
+                    # Calculate backoff with jitter
+                    actual_backoff: float = backoff * (0.5 + random.random())
+                    dbos_logger.warning(
+                        f"Database connection failed: {str(e)}. "
+                        f"Retrying in {actual_backoff:.2f}s (attempt {retries})"
+                    )
+                    # Sleep with backoff
+                    time.sleep(actual_backoff)
+                    # Increase backoff for next attempt (exponential)
+                    backoff = min(backoff * 2, max_backoff)
+
+        return cast(F, wrapper)
+
+    return decorator
+
+
 class SystemDatabase:
 
     def __init__(
@@ -365,7 +413,7 @@ class SystemDatabase:
             self.notification_conn.close()
         self.engine.dispose()
 
-    def insert_workflow_status(
+    def _insert_workflow_status(
         self,
         status: WorkflowStatusInternal,
         conn: sa.Connection,
@@ -474,53 +522,46 @@ class SystemDatabase:
 
         return wf_status, workflow_deadline_epoch_ms
 
+    @db_retry()
     def update_workflow_status(
         self,
         status: WorkflowStatusInternal,
-        *,
-        conn: Optional[sa.Connection] = None,
     ) -> None:
         if self._debug_mode:
             raise Exception("called update_workflow_status in debug mode")
         wf_status: WorkflowStatuses = status["status"]
-
-        cmd = (
-            pg.insert(SystemSchema.workflow_status)
-            .values(
-                workflow_uuid=status["workflow_uuid"],
-                status=status["status"],
-                name=status["name"],
-                class_name=status["class_name"],
-                config_name=status["config_name"],
-                output=status["output"],
-                error=status["error"],
-                executor_id=status["executor_id"],
-                application_version=status["app_version"],
-                application_id=status["app_id"],
-                authenticated_user=status["authenticated_user"],
-                authenticated_roles=status["authenticated_roles"],
-                assumed_role=status["assumed_role"],
-                queue_name=status["queue_name"],
-                recovery_attempts=(
-                    1 if wf_status != WorkflowStatusString.ENQUEUED.value else 0
-                ),
-            )
-            .on_conflict_do_update(
-                index_elements=["workflow_uuid"],
-                set_=dict(
+        with self.engine.begin() as c:
+            c.execute(
+                pg.insert(SystemSchema.workflow_status)
+                .values(
+                    workflow_uuid=status["workflow_uuid"],
                     status=status["status"],
+                    name=status["name"],
+                    class_name=status["class_name"],
+                    config_name=status["config_name"],
                     output=status["output"],
                     error=status["error"],
-                    updated_at=func.extract("epoch", func.now()) * 1000,
-                ),
+                    executor_id=status["executor_id"],
+                    application_version=status["app_version"],
+                    application_id=status["app_id"],
+                    authenticated_user=status["authenticated_user"],
+                    authenticated_roles=status["authenticated_roles"],
+                    assumed_role=status["assumed_role"],
+                    queue_name=status["queue_name"],
+                    recovery_attempts=(
+                        1 if wf_status != WorkflowStatusString.ENQUEUED.value else 0
+                    ),
+                )
+                .on_conflict_do_update(
+                    index_elements=["workflow_uuid"],
+                    set_=dict(
+                        status=status["status"],
+                        output=status["output"],
+                        error=status["error"],
+                        updated_at=func.extract("epoch", func.now()) * 1000,
+                    ),
+                )
             )
-        )
-
-        if conn is not None:
-            conn.execute(cmd)
-        else:
-            with self.engine.begin() as c:
-                c.execute(cmd)
 
     def cancel_workflow(
         self,
@@ -686,6 +727,7 @@ class SystemDatabase:
             )
         return forked_workflow_id
 
+    @db_retry()
     def get_workflow_status(
         self, workflow_uuid: str
     ) -> Optional[WorkflowStatusInternal]:
@@ -735,6 +777,7 @@ class SystemDatabase:
             }
             return status
 
+    @db_retry()
     def await_workflow_result(self, workflow_id: str) -> Any:
         while True:
             with self.engine.begin() as c:
@@ -761,7 +804,7 @@ class SystemDatabase:
                     pass  # CB: I guess we're assuming the WF will show up eventually.
             time.sleep(1)
 
-    def update_workflow_inputs(
+    def _update_workflow_inputs(
         self, workflow_uuid: str, inputs: str, conn: sa.Connection
     ) -> None:
         if self._debug_mode:
@@ -791,6 +834,7 @@ class SystemDatabase:
 
         return
 
+    @db_retry()
     def get_workflow_inputs(
         self, workflow_uuid: str
     ) -> Optional[_serialization.WorkflowInputs]:
@@ -1084,8 +1128,8 @@ class SystemDatabase:
                 for row in rows
             ]
 
-    def record_operation_result(
-        self, result: OperationResultInternal, conn: Optional[sa.Connection] = None
+    def _record_operation_result_txn(
+        self, result: OperationResultInternal, conn: sa.Connection
     ) -> None:
         if self._debug_mode:
             raise Exception("called record_operation_result in debug mode")
@@ -1100,16 +1144,18 @@ class SystemDatabase:
             error=error,
         )
         try:
-            if conn is not None:
-                conn.execute(sql)
-            else:
-                with self.engine.begin() as c:
-                    c.execute(sql)
+            conn.execute(sql)
         except DBAPIError as dbapi_error:
             if dbapi_error.orig.sqlstate == "23505":  # type: ignore
                 raise DBOSWorkflowConflictIDError(result["workflow_uuid"])
             raise
 
+    @db_retry()
+    def record_operation_result(self, result: OperationResultInternal) -> None:
+        with self.engine.begin() as c:
+            self._record_operation_result_txn(result, c)
+
+    @db_retry()
     def record_get_result(
         self, result_workflow_id: str, output: Optional[str], error: Optional[str]
     ) -> None:
@@ -1135,6 +1181,7 @@ class SystemDatabase:
         with self.engine.begin() as c:
             c.execute(sql)
 
+    @db_retry()
     def record_child_workflow(
         self,
         parentUUID: str,
@@ -1159,13 +1206,12 @@ class SystemDatabase:
                 raise DBOSWorkflowConflictIDError(parentUUID)
             raise
 
-    def check_operation_execution(
+    def _check_operation_execution_txn(
         self,
         workflow_id: str,
         function_id: int,
         function_name: str,
-        *,
-        conn: Optional[sa.Connection] = None,
+        conn: sa.Connection,
     ) -> Optional[RecordedResult]:
         # First query: Retrieve the workflow status
         workflow_status_sql = sa.select(
@@ -1183,13 +1229,8 @@ class SystemDatabase:
         )
 
         # Execute both queries
-        if conn is not None:
-            workflow_status_rows = conn.execute(workflow_status_sql).all()
-            operation_output_rows = conn.execute(operation_output_sql).all()
-        else:
-            with self.engine.begin() as c:
-                workflow_status_rows = c.execute(workflow_status_sql).all()
-                operation_output_rows = c.execute(operation_output_sql).all()
+        workflow_status_rows = conn.execute(workflow_status_sql).all()
+        operation_output_rows = conn.execute(operation_output_sql).all()
 
         # Check if the workflow exists
         assert (
@@ -1231,6 +1272,16 @@ class SystemDatabase:
         }
         return result
 
+    @db_retry()
+    def check_operation_execution(
+        self, workflow_id: str, function_id: int, function_name: str
+    ) -> Optional[RecordedResult]:
+        with self.engine.begin() as c:
+            return self._check_operation_execution_txn(
+                workflow_id, function_id, function_name, c
+            )
+
+    @db_retry()
     def check_child_workflow(
         self, workflow_uuid: str, function_id: int
     ) -> Optional[str]:
@@ -1248,6 +1299,7 @@ class SystemDatabase:
             return None
         return str(row[0])
 
+    @db_retry()
     def send(
         self,
         workflow_uuid: str,
@@ -1259,7 +1311,7 @@ class SystemDatabase:
         function_name = "DBOS.send"
         topic = topic if topic is not None else _dbos_null_topic
         with self.engine.begin() as c:
-            recorded_output = self.check_operation_execution(
+            recorded_output = self._check_operation_execution_txn(
                 workflow_uuid, function_id, function_name, conn=c
             )
             if self._debug_mode and recorded_output is None:
@@ -1297,8 +1349,9 @@ class SystemDatabase:
                 "output": None,
                 "error": None,
             }
-            self.record_operation_result(output, conn=c)
+            self._record_operation_result_txn(output, conn=c)
 
+    @db_retry()
     def recv(
         self,
         workflow_uuid: str,
@@ -1391,7 +1444,7 @@ class SystemDatabase:
             message: Any = None
             if len(rows) > 0:
                 message = _serialization.deserialize(rows[0][0])
-            self.record_operation_result(
+            self._record_operation_result_txn(
                 {
                     "workflow_uuid": workflow_uuid,
                     "function_id": function_id,
@@ -1455,13 +1508,14 @@ class SystemDatabase:
                             dbos_logger.error(f"Unknown channel: {channel}")
             except Exception as e:
                 if self._run_background_processes:
-                    dbos_logger.error(f"Notification listener error: {e}")
+                    dbos_logger.warning(f"Notification listener error: {e}")
                     time.sleep(1)
                     # Then the loop will try to reconnect and restart the listener
             finally:
                 if self.notification_conn is not None:
                     self.notification_conn.close()
 
+    @db_retry()
     def sleep(
         self,
         workflow_uuid: str,
@@ -1501,6 +1555,7 @@ class SystemDatabase:
             time.sleep(duration)
         return duration
 
+    @db_retry()
     def set_event(
         self,
         workflow_uuid: str,
@@ -1510,7 +1565,7 @@ class SystemDatabase:
     ) -> None:
         function_name = "DBOS.setEvent"
         with self.engine.begin() as c:
-            recorded_output = self.check_operation_execution(
+            recorded_output = self._check_operation_execution_txn(
                 workflow_uuid, function_id, function_name, conn=c
             )
             if self._debug_mode and recorded_output is None:
@@ -1542,8 +1597,9 @@ class SystemDatabase:
                 "output": None,
                 "error": None,
             }
-            self.record_operation_result(output, conn=c)
+            self._record_operation_result_txn(output, conn=c)
 
+    @db_retry()
     def get_event(
         self,
         target_uuid: str,
@@ -1634,7 +1690,7 @@ class SystemDatabase:
             )
         return value
 
-    def enqueue(
+    def _enqueue(
         self,
         workflow_id: str,
         queue_name: str,
@@ -1710,13 +1766,8 @@ class SystemDatabase:
                 if num_recent_queries >= queue.limiter["limit"]:
                     return []
 
-            # Dequeue functions eligible for this worker and ordered by the time at which they were enqueued.
-            # If there is a global or local concurrency limit N, select only the N oldest enqueued
-            # functions, else select all of them.
-
-            # First lets figure out how many tasks are eligible for dequeue.
-            # This means figuring out how many unstarted tasks are within the local and global concurrency limits
-            running_tasks_query = (
+            # Count how many workflows on this queue are currently PENDING both locally and globally.
+            pending_tasks_query = (
                 sa.select(
                     SystemSchema.workflow_status.c.executor_id,
                     sa.func.count().label("task_count"),
@@ -1730,41 +1781,37 @@ class SystemDatabase:
                 )
                 .where(SystemSchema.workflow_queue.c.queue_name == queue.name)
                 .where(
-                    SystemSchema.workflow_queue.c.started_at_epoch_ms.isnot(
-                        None
-                    )  # Task is started
-                )
-                .where(
-                    SystemSchema.workflow_queue.c.completed_at_epoch_ms.is_(
-                        None
-                    )  # Task is not completed.
+                    SystemSchema.workflow_status.c.status
+                    == WorkflowStatusString.PENDING.value
                 )
                 .group_by(SystemSchema.workflow_status.c.executor_id)
             )
-            running_tasks_result = c.execute(running_tasks_query).fetchall()
-            running_tasks_result_dict = {row[0]: row[1] for row in running_tasks_result}
-            running_tasks_for_this_worker = running_tasks_result_dict.get(
-                executor_id, 0
-            )  # Get count for current executor
+            pending_workflows = c.execute(pending_tasks_query).fetchall()
+            pending_workflows_dict = {row[0]: row[1] for row in pending_workflows}
+            local_pending_workflows = pending_workflows_dict.get(executor_id, 0)
 
+            # Compute max_tasks, the number of workflows that can be dequeued given local and global concurrency limits,
             max_tasks = float("inf")
             if queue.worker_concurrency is not None:
-                max_tasks = max(
-                    0, queue.worker_concurrency - running_tasks_for_this_worker
-                )
+                # Print a warning if the local concurrency limit is violated
+                if local_pending_workflows > queue.worker_concurrency:
+                    dbos_logger.warning(
+                        f"The number of local pending workflows ({local_pending_workflows}) on queue {queue.name} exceeds the local concurrency limit ({queue.worker_concurrency})"
+                    )
+                max_tasks = max(0, queue.worker_concurrency - local_pending_workflows)
+
             if queue.concurrency is not None:
-                total_running_tasks = sum(running_tasks_result_dict.values())
-                # Queue global concurrency limit should always be >= running_tasks_count
-                # This should never happen but a check + warning doesn't hurt
-                if total_running_tasks > queue.concurrency:
+                global_pending_workflows = sum(pending_workflows_dict.values())
+                # Print a warning if the global concurrency limit is violated
+                if global_pending_workflows > queue.concurrency:
                     dbos_logger.warning(
-                        f"Total running tasks ({total_running_tasks}) exceeds the global concurrency limit ({queue.concurrency})"
+                        f"The total number of pending workflows ({global_pending_workflows}) on queue {queue.name} exceeds the global concurrency limit ({queue.concurrency})"
                     )
-                available_tasks = max(0, queue.concurrency - total_running_tasks)
+                available_tasks = max(0, queue.concurrency - global_pending_workflows)
                 max_tasks = min(max_tasks, available_tasks)
 
             # Retrieve the first max_tasks workflows in the queue.
-            # Only retrieve workflows of the appropriate version (or without version set)
+            # Only retrieve workflows of the local version (or without version set)
             query = (
                 sa.select(
                     SystemSchema.workflow_queue.c.workflow_uuid,
@@ -1777,8 +1824,10 @@ class SystemDatabase:
                     )
                 )
                 .where(SystemSchema.workflow_queue.c.queue_name == queue.name)
-                .where(SystemSchema.workflow_queue.c.started_at_epoch_ms == None)
-                .where(SystemSchema.workflow_queue.c.completed_at_epoch_ms == None)
+                .where(
+                    SystemSchema.workflow_status.c.status
+                    == WorkflowStatusString.ENQUEUED.value
+                )
                 .where(
                     sa.or_(
                         SystemSchema.workflow_status.c.application_version
@@ -1807,20 +1856,16 @@ class SystemDatabase:
             ret_ids: list[str] = []
 
             for id in dequeued_ids:
-                # If we have a limiter, stop starting functions when the number
-                # of functions started this period exceeds the limit.
+                # If we have a limiter, stop dequeueing workflows when the number
+                # of workflows started this period exceeds the limit.
                 if queue.limiter is not None:
                     if len(ret_ids) + num_recent_queries >= queue.limiter["limit"]:
                         break
 
-                # To start a function, first set its status to PENDING and update its executor ID
-                res = c.execute(
+                # To start a workflow, first set its status to PENDING and update its executor ID
+                c.execute(
                     SystemSchema.workflow_status.update()
                     .where(SystemSchema.workflow_status.c.workflow_uuid == id)
-                    .where(
-                        SystemSchema.workflow_status.c.status
-                        == WorkflowStatusString.ENQUEUED.value
-                    )
                     .values(
                         status=WorkflowStatusString.PENDING.value,
                         application_version=app_version,
@@ -1843,16 +1888,15 @@ class SystemDatabase:
                         ),
                     )
                 )
-                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)
+                # Then give it a start time
+                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
+            # If we have a limiter, garbage-collect all completed workflows started
             # before the period. If there's no limiter, there's no need--they were
             # deleted on completion.
             if queue.limiter is not None:
@@ -1869,6 +1913,7 @@ class SystemDatabase:
             # Return the IDs of all functions we started
             return ret_ids
 
+    @db_retry()
     def remove_from_queue(self, workflow_id: str, queue: "Queue") -> None:
         if self._debug_mode:
             raise Exception("called remove_from_queue in debug mode")
@@ -1957,6 +2002,7 @@ class SystemDatabase:
             )
         return result
 
+    @db_retry()
     def init_workflow(
         self,
         status: WorkflowStatusInternal,
@@ -1969,17 +2015,17 @@ class SystemDatabase:
         Synchronously record the status and inputs for workflows in a single transaction
         """
         with self.engine.begin() as conn:
-            wf_status, workflow_deadline_epoch_ms = self.insert_workflow_status(
+            wf_status, workflow_deadline_epoch_ms = self._insert_workflow_status(
                 status, conn, max_recovery_attempts=max_recovery_attempts
             )
             # TODO: Modify the inputs if they were changed by `update_workflow_inputs`
-            self.update_workflow_inputs(status["workflow_uuid"], inputs, conn)
+            self._update_workflow_inputs(status["workflow_uuid"], inputs, conn)
 
             if (
                 status["queue_name"] is not None
                 and wf_status == WorkflowStatusString.ENQUEUED.value
             ):
-                self.enqueue(
+                self._enqueue(
                     status["workflow_uuid"],
                     status["queue_name"],
                     conn,
@@ -1987,6 +2033,14 @@ class SystemDatabase:
                 )
         return wf_status, workflow_deadline_epoch_ms
 
+    def check_connection(self) -> None:
+        try:
+            with self.engine.begin() as conn:
+                conn.execute(sa.text("SELECT 1")).fetchall()
+        except Exception as e:
+            dbos_logger.error(f"Error connecting to the DBOS system database: {e}")
+            raise
+
 
 def reset_system_database(postgres_db_url: sa.URL, sysdb_name: str) -> None:
     try:

dbos/_utils.py

@@ -1,6 +1,9 @@
 import importlib.metadata
 import os
 
+import psycopg
+from sqlalchemy.exc import DBAPIError
+
 INTERNAL_QUEUE_NAME = "_dbos_internal_queue"
 
 request_id_header = "x-request-id"
@@ -15,3 +18,33 @@ class GlobalParams:
     except importlib.metadata.PackageNotFoundError:
         # If package is not installed or during development
         dbos_version = "unknown"
+
+
+def retriable_postgres_exception(e: DBAPIError) -> bool:
+    if e.connection_invalidated:
+        return True
+    if isinstance(e.orig, psycopg.OperationalError):
+        driver_error: psycopg.OperationalError = e.orig
+        pgcode = driver_error.sqlstate or ""
+        # Failure to establish connection
+        if "connection failed" in str(driver_error):
+            return True
+        # Error within database transaction
+        elif "server closed the connection unexpectedly" in str(driver_error):
+            return True
+        # Connection timeout
+        if isinstance(driver_error, psycopg.errors.ConnectionTimeout):
+            return True
+        # Insufficient resources
+        elif pgcode.startswith("53"):
+            return True
+        # Connection exception
+        elif pgcode.startswith("08"):
+            return True
+        # Operator intervention
+        elif pgcode.startswith("57"):
+            return True
+        else:
+            return False
+    else:
+        return False

{dbos-1.2.0a2.dist-info → dbos-1.2.0a5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dbos
-Version: 1.2.0a2
+Version: 1.2.0a5
 Summary: Ultra-lightweight durable execution in Python
 Author-Email: "DBOS, Inc." <contact@dbos.dev>
 License: MIT

{dbos-1.2.0a2.dist-info → dbos-1.2.0a5.dist-info}/RECORD

@@ -1,20 +1,20 @@
-dbos-1.2.0a2.dist-info/METADATA,sha256=at-2zS4N-BoXxmKlOeZt7HRA5shIG9pu3pb98t8VFNs,13267
-dbos-1.2.0a2.dist-info/WHEEL,sha256=tSfRZzRHthuv7vxpI4aehrdN9scLjk-dCJkPLzkHxGg,90
-dbos-1.2.0a2.dist-info/entry_points.txt,sha256=_QOQ3tVfEjtjBlr1jS4sHqHya9lI2aIEIWkz8dqYp14,58
-dbos-1.2.0a2.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
+dbos-1.2.0a5.dist-info/METADATA,sha256=PeZezLDhF3k-2FR5-9qCP0_BVWuEVSvo41QucKSXD8o,13267
+dbos-1.2.0a5.dist-info/WHEEL,sha256=tSfRZzRHthuv7vxpI4aehrdN9scLjk-dCJkPLzkHxGg,90
+dbos-1.2.0a5.dist-info/entry_points.txt,sha256=_QOQ3tVfEjtjBlr1jS4sHqHya9lI2aIEIWkz8dqYp14,58
+dbos-1.2.0a5.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
 dbos/__init__.py,sha256=NssPCubaBxdiKarOWa-wViz1hdJSkmBGcpLX_gQ4NeA,891
 dbos/__main__.py,sha256=G7Exn-MhGrVJVDbgNlpzhfh8WMX_72t3_oJaFT9Lmt8,653
 dbos/_admin_server.py,sha256=TWXi4drrzKFpKkUmEJpJkQBZxAtOalnhtYicEn2nDK0,10618
 dbos/_app_db.py,sha256=0PKqpxJ3EbIaak3Wl0lNl3hXvhBfz4EEHaCw1bUOvIM,9937
 dbos/_classproperty.py,sha256=f0X-_BySzn3yFDRKB2JpCbLYQ9tLwt1XftfshvY7CBs,626
-dbos/_client.py,sha256=-nK2GjS9D0qnD2DkRDs7gKxNECwYlsvW6hFCjADlnv0,14186
+dbos/_client.py,sha256=mGDuQRcSdkyEHf1s0rJuqHQiWbqIBt85qijNJSYmBik,14227
 dbos/_conductor/conductor.py,sha256=o0IaZjwnZ2TOyHeP2H4iSX6UnXLXQ4uODvWAKD9hHMs,21703
 dbos/_conductor/protocol.py,sha256=wgOFZxmS81bv0WCB9dAyg0s6QzldpzVKQDoSPeaX0Ws,6967
 dbos/_context.py,sha256=5ajoWAmToAfzzmMLylnJZoL4Ny9rBwZWuG05sXadMIA,24798
-dbos/_core.py,sha256=7ukQH_KClBaMFy0sVTSR5tWylW-RqI9qaReBY-LDKrk,48316
+dbos/_core.py,sha256=m2i9lsHjNKTi8BQyiSOUBrAVH5OvMoBswNZPRpMVIC0,48662
 dbos/_croniter.py,sha256=XHAyUyibs_59sJQfSNWkP7rqQY6_XrlfuuCxk4jYqek,47559
-dbos/_dbos.py,sha256=MuMYbtqUyk2uihCH8aMVDeHmn_P8X8-udqeNT1RLesY,47365
-dbos/_dbos_config.py,sha256=IufNrIC-M2xSNTXyT_KXlEdfB3j03pPLv_nE0fEq4_U,20955
+dbos/_dbos.py,sha256=1EhH7r6v2vwW3Z74nK6_Zw8InE1jSXedEsztz0I4ggA,47269
+dbos/_dbos_config.py,sha256=JYtEbhjcCxLUhktMgqIEBz7i5nk1Ryg0vqSJHXqdGOo,20264
 dbos/_debug.py,sha256=MNlQVZ6TscGCRQeEEL0VE8Uignvr6dPeDDDefS3xgIE,1823
 dbos/_docker_pg_helper.py,sha256=tLJXWqZ4S-ExcaPnxg_i6cVxL6ZxrYlZjaGsklY-s2I,6115
 dbos/_error.py,sha256=q0OQJZTbR8FFHV9hEpAGpz9oWBT5L509zUhmyff7FJw,8500
@@ -38,7 +38,7 @@ dbos/_migrations/versions/d76646551a6c_workflow_queue.py,sha256=G942nophZ2uC2vc4
 dbos/_migrations/versions/eab0cc1d9a14_job_queue.py,sha256=uvhFOtqbBreCePhAxZfIT0qCAI7BiZTou9wt6QnbY7c,1412
 dbos/_migrations/versions/f4b9b32ba814_functionname_childid_op_outputs.py,sha256=m90Lc5YH0ZISSq1MyxND6oq3RZrZKrIqEsZtwJ1jWxA,1049
 dbos/_outcome.py,sha256=EXxBg4jXCVJsByDQ1VOCIedmbeq_03S6d-p1vqQrLFU,6810
-dbos/_queue.py,sha256=aKCGahWBGJOLOv5PCOOId96Va3YQ4ICuHWXy-eQXohE,3526
+dbos/_queue.py,sha256=oDQcydDwYM68U5KQKN6iZiSC-4LXye6KFmSJ7ohG048,3558
 dbos/_recovery.py,sha256=jVMexjfCCNopzyn8gVQzJCmGJaP9G3C1EFaoCQ_Nh7g,2564
 dbos/_registrations.py,sha256=CZt1ElqDjCT7hz6iyT-1av76Yu-iuwu_c9lozO87wvM,7303
 dbos/_roles.py,sha256=iOsgmIAf1XVzxs3gYWdGRe1B880YfOw5fpU7Jwx8_A8,2271
@@ -47,7 +47,7 @@ dbos/_schemas/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/_schemas/application_database.py,sha256=SypAS9l9EsaBHFn9FR8jmnqt01M74d9AF1AMa4m2hhI,1040
 dbos/_schemas/system_database.py,sha256=3Z0L72bOgHnusK1hBaETWU9RfiLBP0QnS-fdu41i0yY,5835
 dbos/_serialization.py,sha256=bWuwhXSQcGmiazvhJHA5gwhrRWxtmFmcCFQSDJnqqkU,3666
-dbos/_sys_db.py,sha256=IMmRbeIcrsOFJfVcBhMkDWiA3_SvxeKbOGipFiplHPM,83735
+dbos/_sys_db.py,sha256=T02hZbe-4tpsK4hGVlatft06ybu86SJ4w6-anaf55KQ,85528
 dbos/_templates/dbos-db-starter/README.md,sha256=GhxhBj42wjTt1fWEtwNriHbJuKb66Vzu89G4pxNHw2g,930
 dbos/_templates/dbos-db-starter/__package/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/_templates/dbos-db-starter/__package/main.py.dbos,sha256=aQnBPSSQpkB8ERfhf7gB7P9tsU6OPKhZscfeh0yiaD8,2702
@@ -59,7 +59,7 @@ dbos/_templates/dbos-db-starter/migrations/script.py.mako,sha256=MEqL-2qATlST9TA
 dbos/_templates/dbos-db-starter/migrations/versions/2024_07_31_180642_init.py,sha256=MpS7LGaJS0CpvsjhfDkp9EJqvMvVCjRPfUp4c0aE2ys,941
 dbos/_templates/dbos-db-starter/start_postgres_docker.py,sha256=lQVLlYO5YkhGPEgPqwGc7Y8uDKse9HsWv5fynJEFJHM,1681
 dbos/_tracer.py,sha256=yN6GRDKu_1p-EqtQLNarMocPfga2ZuqpzStzzSPYhzo,2732
-dbos/_utils.py,sha256=UbpMYRBSyvJqdXeWAnfSw8xXM1R1mfnyl1oTunhEjJM,513
+dbos/_utils.py,sha256=uywq1QrjMwy17btjxW4bES49povlQwYwYbvKwMT6C2U,1575
 dbos/_workflow_commands.py,sha256=UCpHWvCEXjVZtf5FNanFvtJpgUJDSI1EFBqQP0x_2A0,3346
 dbos/cli/_github_init.py,sha256=Y_bDF9gfO2jB1id4FV5h1oIxEJRWyqVjhb7bNEa5nQ0,3224
 dbos/cli/_template_init.py,sha256=7JBcpMqP1r2mfCnvWatu33z8ctEGHJarlZYKgB83cXE,2972
@@ -67,4 +67,4 @@ dbos/cli/cli.py,sha256=HinoCGrAUTiSeq7AAoCFfhdiE0uDw7vLMuDMN1_YTLI,20705
 dbos/dbos-config.schema.json,sha256=CjaspeYmOkx6Ip_pcxtmfXJTn_YGdSx_0pcPBF7KZmo,6060
 dbos/py.typed,sha256=QfzXT1Ktfk3Rj84akygc7_42z0lRpCq0Ilh8OXI6Zas,44
 version/__init__.py,sha256=L4sNxecRuqdtSFdpUGX3TtBi9KL3k7YsZVIvv-fv9-A,1678
-dbos-1.2.0a2.dist-info/RECORD,,
+dbos-1.2.0a5.dist-info/RECORD,,