dbos 0.18.0a1__py3-none-any.whl → 0.19.0__py3-none-any.whl

dbos/_context.py

@@ -57,6 +57,7 @@ class DBOSContext:
         self.request: Optional["Request"] = None
 
         self.id_assigned_for_next_workflow: str = ""
+        self.is_within_set_workflow_id_block: bool = False
 
         self.parent_workflow_id: str = ""
         self.parent_workflow_fid: int = -1
@@ -78,6 +79,7 @@ class DBOSContext:
         rv.logger = self.logger
         rv.id_assigned_for_next_workflow = self.id_assigned_for_next_workflow
         self.id_assigned_for_next_workflow = ""
+        rv.is_within_set_workflow_id_block = self.is_within_set_workflow_id_block
         rv.parent_workflow_id = self.workflow_id
         rv.parent_workflow_fid = self.function_id
         rv.in_recovery = self.in_recovery
@@ -95,6 +97,10 @@ class DBOSContext:
         if len(self.id_assigned_for_next_workflow) > 0:
             wfid = self.id_assigned_for_next_workflow
         else:
+            if self.is_within_set_workflow_id_block:
+                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())
         return wfid
 
@@ -286,7 +292,7 @@ class DBOSContextSwap:
 
 class SetWorkflowID:
     """
-    Set the workflow ID to be used for the enclosed workflow invocation.
+    Set the workflow ID to be used for the enclosed workflow invocation. Note: Only the first workflow will be started with the specified workflow ID within a `with SetWorkflowID` block.
 
     Typical Usage
         ```
@@ -311,7 +317,9 @@ class SetWorkflowID:
         if ctx is None:
             self.created_ctx = True
             _set_local_dbos_context(DBOSContext())
-        assert_current_dbos_context().id_assigned_for_next_workflow = self.wfid
+        ctx = assert_current_dbos_context()
+        ctx.id_assigned_for_next_workflow = self.wfid
+        ctx.is_within_set_workflow_id_block = True
         return self
 
     def __exit__(
@@ -321,6 +329,7 @@ class SetWorkflowID:
         traceback: Optional[TracebackType],
     ) -> Literal[False]:
         # Code to clean up the basic context if we created it
+        assert_current_dbos_context().is_within_set_workflow_id_block = False
         if self.created_ctx:
             _clear_local_dbos_context()
         return False  # Did not handle

dbos/_core.py

@@ -84,7 +84,7 @@ if TYPE_CHECKING:
         IsolationLevel,
     )
 
-from sqlalchemy.exc import DBAPIError
+from sqlalchemy.exc import DBAPIError, InvalidRequestError
 
 P = ParamSpec("P")  # A generic type for workflow parameters
 R = TypeVar("R", covariant=True)  # A generic type for workflow return values
@@ -180,21 +180,24 @@ def _init_workflow(
     if class_name is not None:
         inputs = {"args": inputs["args"][1:], "kwargs": inputs["kwargs"]}
 
+    wf_status = status["status"]
     if temp_wf_type != "transaction" or queue is not None:
         # Synchronously record the status and inputs for workflows and single-step workflows
         # We also have to do this for single-step workflows because of the foreign key constraint on the operation outputs table
         # TODO: Make this transactional (and with the queue step below)
-        dbos._sys_db.update_workflow_status(
+        wf_status = dbos._sys_db.update_workflow_status(
             status, False, ctx.in_recovery, max_recovery_attempts=max_recovery_attempts
         )
+        # TODO: Modify the inputs if they were changed by `update_workflow_inputs`
         dbos._sys_db.update_workflow_inputs(wfid, _serialization.serialize_args(inputs))
     else:
         # Buffer the inputs for single-transaction workflows, but don't buffer the status
         dbos._sys_db.buffer_workflow_inputs(wfid, _serialization.serialize_args(inputs))
 
-    if queue is not None:
+    if queue is not None and wf_status == WorkflowStatusString.ENQUEUED.value:
         dbos._sys_db.enqueue(wfid, queue)
 
+    status["status"] = wf_status
     return status
 
 
@@ -413,7 +416,16 @@ def start_workflow(
         max_recovery_attempts=fi.max_recovery_attempts,
     )
 
-    if not execute_workflow:
+    wf_status = status["status"]
+
+    if (
+        not execute_workflow
+        or wf_status == WorkflowStatusString.ERROR.value
+        or wf_status == WorkflowStatusString.SUCCESS.value
+    ):
+        dbos.logger.debug(
+            f"Workflow {new_wf_id} already completed with status {wf_status}. Directly returning a workflow handle."
+        )
         return WorkflowHandlePolling(new_wf_id, dbos)
 
     if fself is not None:
@@ -486,7 +498,7 @@ def workflow_wrapper(
                 temp_wf_type=get_temp_workflow_type(func),
                 max_recovery_attempts=max_recovery_attempts,
             )
-
+            # TODO: maybe modify the parameters if they've been changed by `_init_workflow`
             dbos.logger.debug(
                 f"Running workflow, id: {ctx.workflow_id}, name: {get_dbos_func_name(func)}"
             )
@@ -545,6 +557,7 @@ def decorate_transaction(
                     max_retry_wait_seconds = 2.0
                     while True:
                         has_recorded_error = False
+                        txn_error: Optional[Exception] = None
                         try:
                             with session.begin():
                                 # This must be the first statement in the transaction!
@@ -608,15 +621,24 @@ def decorate_transaction(
                                     max_retry_wait_seconds,
                                 )
                                 continue
+                            txn_error = dbapi_error
+                            raise
+                        except InvalidRequestError as invalid_request_error:
+                            dbos.logger.error(
+                                f"InvalidRequestError in transaction {func.__qualname__} \033[1m Hint: Do not call commit() or rollback() within a DBOS transaction.\033[0m"
+                            )
+                            txn_error = invalid_request_error
                             raise
                         except Exception as error:
+                            txn_error = error
+                            raise
+                        finally:
                             # Don't record the error if it was already recorded
-                            if not has_recorded_error:
+                            if txn_error and not has_recorded_error:
                                 txn_output["error"] = (
-                                    _serialization.serialize_exception(error)
+                                    _serialization.serialize_exception(txn_error)
                                 )
                                 dbos._app_db.record_transaction_error(txn_output)
-                            raise
             return output
 
         if inspect.iscoroutinefunction(func):

dbos/_db_wizard.py

@@ -1,5 +1,7 @@
+import json
+import os
 import time
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Optional, TypedDict
 
 import docker  # type: ignore
 import typer
@@ -15,16 +17,29 @@ from ._cloudutils.databases import choose_database, get_user_db_credentials
 from ._error import DBOSInitializationError
 from ._logger import dbos_logger
 
+DB_CONNECTION_PATH = os.path.join(".dbos", "db_connection")
 
-def db_connect(config: "ConfigFile", config_file_path: str) -> "ConfigFile":
+
+class DatabaseConnection(TypedDict):
+    hostname: Optional[str]
+    port: Optional[int]
+    username: Optional[str]
+    password: Optional[str]
+    local_suffix: Optional[bool]
+
+
+def db_wizard(config: "ConfigFile", config_file_path: str) -> "ConfigFile":
     # 1. Check the connectivity to the database. Return if successful. If cannot connect, continue to the following steps.
     db_connection_error = _check_db_connectivity(config)
     if db_connection_error is None:
         return config
 
     # 2. If the error is due to password authentication or the configuration is non-default, surface the error and exit.
-    if "password authentication failed" in str(db_connection_error) or "28P01" in str(
-        db_connection_error
+    error_str = str(db_connection_error)
+    if (
+        "password authentication failed" in error_str
+        or "28P01" in error_str
+        or "no password supplied" in error_str
     ):
         raise DBOSInitializationError(
             f"Could not connect to Postgres: password authentication failed: {db_connection_error}"
@@ -79,17 +94,20 @@ def db_connect(config: "ConfigFile", config_file_path: str) -> "ConfigFile":
                 f"Could not connect to the database. Exception: {db_connection_error}"
             )
 
-    # 6. Save the config to the config file and return the updated config.
-    # TODO: make the config file prettier
-    with open(config_file_path, "w") as file:
-        file.write(yaml.dump(config))
-
+    # 6. Save the config to the database connection file
+    updated_connection = DatabaseConnection(
+        hostname=config["database"]["hostname"],
+        port=config["database"]["port"],
+        username=config["database"]["username"],
+        password=config["database"]["password"],
+        local_suffix=config["database"]["local_suffix"],
+    )
+    save_db_connection(updated_connection)
     return config
 
 
 def _start_docker_postgres(config: "ConfigFile") -> bool:
     print("Starting a Postgres Docker container...")
-    config["database"]["password"] = "dbos"
     client = docker.from_env()
     pg_data = "/var/lib/postgresql/data"
     container_name = "dbos-db"
@@ -119,7 +137,7 @@ def _start_docker_postgres(config: "ConfigFile") -> bool:
                 continue
             print("[green]Postgres Docker container started successfully![/green]")
             break
-        except Exception as e:
+        except:
             attempts -= 1
             time.sleep(1)
 
@@ -148,7 +166,7 @@ def _check_db_connectivity(config: "ConfigFile") -> Optional[Exception]:
         host=config["database"]["hostname"],
         port=config["database"]["port"],
         database="postgres",
-        query={"connect_timeout": "2"},
+        query={"connect_timeout": "1"},
     )
     postgres_db_engine = create_engine(postgres_db_url)
     try:
@@ -165,3 +183,26 @@ def _check_db_connectivity(config: "ConfigFile") -> Optional[Exception]:
         postgres_db_engine.dispose()
 
     return None
+
+
+def load_db_connection() -> DatabaseConnection:
+    try:
+        with open(DB_CONNECTION_PATH, "r") as f:
+            data = json.load(f)
+            return DatabaseConnection(
+                hostname=data.get("hostname", None),
+                port=data.get("port", None),
+                username=data.get("username", None),
+                password=data.get("password", None),
+                local_suffix=data.get("local_suffix", None),
+            )
+    except:
+        return DatabaseConnection(
+            hostname=None, port=None, username=None, password=None, local_suffix=None
+        )
+
+
+def save_db_connection(connection: DatabaseConnection) -> None:
+    os.makedirs(".dbos", exist_ok=True)
+    with open(DB_CONNECTION_PATH, "w") as f:
+        json.dump(connection, f)

dbos/_dbos.py

@@ -83,7 +83,7 @@ from ._context import (
 )
 from ._dbos_config import ConfigFile, load_config, set_env_vars
 from ._error import DBOSException, DBOSNonExistentWorkflowError
-from ._logger import add_otlp_to_all_loggers, dbos_logger, init_logger
+from ._logger import add_otlp_to_all_loggers, dbos_logger
 from ._sys_db import SystemDatabase
 
 # Most DBOS functions are just any callable F, so decorators / wrappers work on F

dbos/_dbos_config.py

@@ -6,12 +6,15 @@ from typing import Any, Dict, List, Optional, TypedDict, cast
 
 import yaml
 from jsonschema import ValidationError, validate
+from rich import print
 from sqlalchemy import URL
 
-from ._db_wizard import db_connect
+from ._db_wizard import db_wizard, load_db_connection
 from ._error import DBOSInitializationError
 from ._logger import config_logger, dbos_logger, init_logger
 
+DBOS_CONFIG_PATH = "dbos-config.yaml"
+
 
 class RuntimeConfig(TypedDict, total=False):
     start: List[str]
@@ -23,7 +26,7 @@ class DatabaseConfig(TypedDict, total=False):
     hostname: str
     port: int
     username: str
-    password: Optional[str]
+    password: str
     connectionTimeoutMillis: Optional[int]
     app_db_name: str
     sys_db_name: Optional[str]
@@ -93,7 +96,7 @@ def _substitute_env_vars(content: str) -> str:
     return re.sub(regex, replace_func, content)
 
 
-def get_dbos_database_url(config_file_path: str = "dbos-config.yaml") -> str:
+def get_dbos_database_url(config_file_path: str = DBOS_CONFIG_PATH) -> str:
     """
     Retrieve application database URL from configuration `.yaml` file.
 
@@ -119,7 +122,9 @@ def get_dbos_database_url(config_file_path: str = "dbos-config.yaml") -> str:
     return db_url.render_as_string(hide_password=False)
 
 
-def load_config(config_file_path: str = "dbos-config.yaml") -> ConfigFile:
+def load_config(
+    config_file_path: str = DBOS_CONFIG_PATH, *, use_db_wizard: bool = True
+) -> ConfigFile:
     """
     Load the DBOS `ConfigFile` from the specified path (typically `dbos-config.yaml`).
 
@@ -151,6 +156,9 @@ def load_config(config_file_path: str = "dbos-config.yaml") -> ConfigFile:
     except ValidationError as e:
         raise DBOSInitializationError(f"Validation error: {e}")
 
+    if "database" not in data:
+        data["database"] = {}
+
     if "name" not in data:
         raise DBOSInitializationError(
             f"dbos-config.yaml must specify an application name"
@@ -169,8 +177,6 @@ def load_config(config_file_path: str = "dbos-config.yaml") -> ConfigFile:
     if "runtimeConfig" not in data or "start" not in data["runtimeConfig"]:
         raise DBOSInitializationError(f"dbos-config.yaml must specify a start command")
 
-    data = cast(ConfigFile, data)
-
     if not _is_valid_app_name(data["name"]):
         raise DBOSInitializationError(
             f'Invalid app name {data["name"]}.  App names must be between 3 and 30 characters long and contain only lowercase letters, numbers, dashes, and underscores.'
@@ -179,10 +185,49 @@ def load_config(config_file_path: str = "dbos-config.yaml") -> ConfigFile:
     if "app_db_name" not in data["database"]:
         data["database"]["app_db_name"] = _app_name_to_db_name(data["name"])
 
+    # Load the DB connection file. Use its values for missing fields from dbos-config.yaml. Use defaults otherwise.
+    data = cast(ConfigFile, data)
+    db_connection = load_db_connection()
+    if data["database"].get("hostname"):
+        print(
+            "[bold blue]Loading database connection parameters from dbos-config.yaml[/bold blue]"
+        )
+    elif db_connection.get("hostname"):
+        print(
+            "[bold blue]Loading database connection parameters from .dbos/db_connection[/bold blue]"
+        )
+    else:
+        print(
+            "[bold blue]Using default database connection parameters (localhost)[/bold blue]"
+        )
+
+    data["database"]["hostname"] = (
+        data["database"].get("hostname") or db_connection.get("hostname") or "localhost"
+    )
+    data["database"]["port"] = (
+        data["database"].get("port") or db_connection.get("port") or 5432
+    )
+    data["database"]["username"] = (
+        data["database"].get("username") or db_connection.get("username") or "postgres"
+    )
+    data["database"]["password"] = (
+        data["database"].get("password")
+        or db_connection.get("password")
+        or os.environ.get("PGPASSWORD")
+        or "dbos"
+    )
+    data["database"]["local_suffix"] = (
+        data["database"].get("local_suffix")
+        or db_connection.get("local_suffix")
+        or False
+    )
+
+    # Configure the DBOS logger
     config_logger(data)
 
     # Check the connectivity to the database and make sure it's properly configured
-    data = db_connect(data, config_file_path)
+    if use_db_wizard:
+        data = db_wizard(data, config_file_path)
 
     if "local_suffix" in data["database"] and data["database"]["local_suffix"]:
         data["database"]["app_db_name"] = f"{data['database']['app_db_name']}_local"

dbos/_error.py

@@ -35,6 +35,7 @@ class DBOSErrorCode(Enum):
     DeadLetterQueueError = 6
     MaxStepRetriesExceeded = 7
     NotAuthorized = 8
+    ConflictingWorkflowError = 9
 
 
 class DBOSWorkflowConflictIDError(DBOSException):
@@ -47,6 +48,16 @@ class DBOSWorkflowConflictIDError(DBOSException):
         )
 
 
+class DBOSConflictingWorkflowError(DBOSException):
+    """Exception raised different workflows started with the same workflow ID."""
+
+    def __init__(self, workflow_id: str, message: Optional[str] = None):
+        super().__init__(
+            f"Conflicting workflow invocation with the same ID ({workflow_id}): {message}",
+            dbos_error_code=DBOSErrorCode.ConflictingWorkflowError.value,
+        )
+
+
 class DBOSRecoveryError(DBOSException):
     """Exception raised when a workflow recovery fails."""
 

dbos/_kafka.py

@@ -1,3 +1,4 @@
+import re
 import threading
 from typing import TYPE_CHECKING, Any, Callable, NoReturn
 
@@ -19,6 +20,14 @@ _kafka_queue: Queue
 _in_order_kafka_queues: dict[str, Queue] = {}
 
 
+def safe_group_name(method_name: str, topics: list[str]) -> str:
+    safe_group_id = "-".join(
+        re.sub(r"[^a-zA-Z0-9\-]", "", str(r)) for r in [method_name, *topics]
+    )
+
+    return f"dbos-kafka-group-{safe_group_id}"[:255]
+
+
 def _kafka_consumer_loop(
     func: _KafkaConsumerWorkflow,
     config: dict[str, Any],
@@ -34,6 +43,12 @@ def _kafka_consumer_loop(
     if "auto.offset.reset" not in config:
         config["auto.offset.reset"] = "earliest"
 
+    if config.get("group.id") is None:
+        config["group.id"] = safe_group_name(func.__qualname__, topics)
+        dbos_logger.warning(
+            f"Consumer group ID not found. Using generated group.id {config['group.id']}"
+        )
+
     consumer = Consumer(config)
     try:
         consumer.subscribe(topics)
@@ -71,8 +86,9 @@ def _kafka_consumer_loop(
                     topic=cmsg.topic(),
                     value=cmsg.value(),
                 )
+                groupID = config.get("group.id")
                 with SetWorkflowID(
-                    f"kafka-unique-id-{msg.topic}-{msg.partition}-{msg.offset}"
+                    f"kafka-unique-id-{msg.topic}-{msg.partition}-{groupID}-{msg.offset}"
                 ):
                     if in_order:
                         assert msg.topic is not None

dbos/_migrations/versions/04ca4f231047_workflow_queues_executor_id.py

@@ -0,0 +1,34 @@
+"""workflow_queues_executor_id
+
+Revision ID: 04ca4f231047
+Revises: d76646551a6c
+Create Date: 2025-01-15 15:05:08.043190
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "04ca4f231047"
+down_revision: Union[str, None] = "d76646551a6c"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "workflow_queue",
+        sa.Column(
+            "executor_id",
+            sa.Text(),
+            nullable=True,
+        ),
+        schema="dbos",
+    )
+
+
+def downgrade() -> None:
+    op.drop_column("workflow_queue", "executor_id", schema="dbos")

dbos/_queue.py

@@ -2,6 +2,9 @@ import threading
 import traceback
 from typing import TYPE_CHECKING, Optional, TypedDict
 
+from psycopg import errors
+from sqlalchemy.exc import OperationalError
+
 from ._core import P, R, execute_workflow_by_id, start_workflow
 
 if TYPE_CHECKING:
@@ -33,9 +36,20 @@ class Queue:
         name: str,
         concurrency: Optional[int] = None,
         limiter: Optional[QueueRateLimit] = None,
+        *,  # Disable positional arguments from here on
+        worker_concurrency: Optional[int] = None,
     ) -> None:
+        if (
+            worker_concurrency is not None
+            and concurrency is not None
+            and worker_concurrency > concurrency
+        ):
+            raise ValueError(
+                "worker_concurrency must be less than or equal to concurrency"
+            )
         self.name = name
         self.concurrency = concurrency
+        self.worker_concurrency = worker_concurrency
         self.limiter = limiter
         from ._dbos import _get_or_create_dbos_registry
 
@@ -60,6 +74,12 @@ def queue_thread(stop_event: threading.Event, dbos: "DBOS") -> None:
                 wf_ids = dbos._sys_db.start_queued_workflows(queue, dbos._executor_id)
                 for id in wf_ids:
                     execute_workflow_by_id(dbos, id)
+            except OperationalError as e:
+                # Ignore serialization error
+                if not isinstance(e.orig, errors.SerializationFailure):
+                    dbos.logger.warning(
+                        f"Exception encountered in queue thread: {traceback.format_exc()}"
+                    )
             except Exception:
                 dbos.logger.warning(
                     f"Exception encountered in queue thread: {traceback.format_exc()}"

dbos/_schemas/system_database.py

@@ -154,6 +154,7 @@ class SystemSchema:
             nullable=False,
             primary_key=True,
         ),
+        Column("executor_id", Text),
         Column("queue_name", Text, nullable=False),
         Column(
             "created_at_epoch_ms",