dbos 2.2.0a3__tar.gz → 2.3.0__tar.gz

{dbos-2.2.0a3 → dbos-2.3.0}/PKG-INFO

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

{dbos-2.2.0a3 → dbos-2.3.0}/dbos/_client.py

@@ -1,4 +1,5 @@
 import asyncio
+import json
 import time
 import uuid
 from typing import (
@@ -63,6 +64,8 @@ class EnqueueOptions(_EnqueueOptionsRequired, total=False):
     priority: int
     max_recovery_attempts: int
     queue_partition_key: str
+    authenticated_user: str
+    authenticated_roles: list[str]
 
 
 def validate_enqueue_options(options: EnqueueOptions) -> None:
@@ -189,6 +192,13 @@ class DBOSClient:
             "queue_partition_key": options.get("queue_partition_key"),
         }
 
+        authenticated_user = options.get("authenticated_user")
+        authenticated_roles = (
+            json.dumps(options.get("authenticated_roles"))
+            if options.get("authenticated_roles")
+            else None
+        )
+
         inputs: WorkflowInputs = {
             "args": args,
             "kwargs": kwargs,
@@ -202,9 +212,9 @@ class DBOSClient:
             "queue_name": queue_name,
             "app_version": enqueue_options_internal["app_version"],
             "config_name": None,
-            "authenticated_user": None,
+            "authenticated_user": authenticated_user,
             "assumed_role": None,
-            "authenticated_roles": None,
+            "authenticated_roles": authenticated_roles,
             "output": None,
             "error": None,
             "created_at": None,

{dbos-2.2.0a3 → dbos-2.3.0}/dbos/_core.py

@@ -93,14 +93,6 @@ TEMP_SEND_WF_NAME = "<temp>.temp_send_workflow"
 DEBOUNCER_WORKFLOW_NAME = "_dbos_debouncer_workflow"
 
 
-def check_is_in_coroutine() -> bool:
-    try:
-        asyncio.get_running_loop()
-        return True
-    except RuntimeError:
-        return False
-
-
 class WorkflowHandleFuture(Generic[R]):
 
     def __init__(self, workflow_id: str, future: Future[R], dbos: "DBOS"):
@@ -856,11 +848,6 @@ def workflow_wrapper(
             dbos._sys_db.record_get_result(workflow_id, serialized_r, None)
             return r
 
-        if check_is_in_coroutine() and not inspect.iscoroutinefunction(func):
-            dbos_logger.warning(
-                f"Sync workflow ({get_dbos_func_name(func)}) shouldn't be invoked from within another async function. Define it as async or use asyncio.to_thread instead."
-            )
-
         outcome = (
             wfOutcome.wrap(init_wf, dbos=dbos)
             .also(DBOSAssumeRole(rr))
@@ -1046,10 +1033,6 @@ def decorate_transaction(
                 assert (
                     ctx.is_workflow()
                 ), "Transactions must be called from within workflows"
-                if check_is_in_coroutine():
-                    dbos_logger.warning(
-                        f"Transaction function ({get_dbos_func_name(func)}) shouldn't be invoked from within another async function. Use asyncio.to_thread instead."
-                    )
                 with DBOSAssumeRole(rr):
                     return invoke_tx(*args, **kwargs)
             else:
@@ -1194,10 +1177,6 @@ def decorate_step(
 
         @wraps(func)
         def wrapper(*args: Any, **kwargs: Any) -> Any:
-            if check_is_in_coroutine() and not inspect.iscoroutinefunction(func):
-                dbos_logger.warning(
-                    f"Sync step ({get_dbos_func_name(func)}) shouldn't be invoked from within another async function. Define it as async or use asyncio.to_thread instead."
-                )
             # If the step is called from a workflow, run it as a step.
             # Otherwise, run it as a normal function.
             ctx = get_local_dbos_context()

{dbos-2.2.0a3 → dbos-2.3.0}/dbos/_dbos_config.py

@@ -444,6 +444,7 @@ def configure_db_engine_parameters(
 
     # Configure user database engine parameters
     app_engine_kwargs: dict[str, Any] = {
+        "connect_args": {"application_name": "dbos_transact"},
         "pool_timeout": 30,
         "max_overflow": 0,
         "pool_size": 20,
@@ -477,8 +478,6 @@ def is_valid_database_url(database_url: str) -> bool:
         return True
     url = make_url(database_url)
     required_fields = [
-        ("username", "Username 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:

{dbos-2.2.0a3 → dbos-2.3.0}/dbos/_kafka.py

@@ -1,6 +1,6 @@
 import re
 import threading
-from typing import TYPE_CHECKING, Any, Callable, NoReturn
+from typing import TYPE_CHECKING, Any, Callable, Coroutine, NoReturn
 
 from confluent_kafka import Consumer, KafkaError, KafkaException
 
@@ -15,7 +15,9 @@ from ._kafka_message import KafkaMessage
 from ._logger import dbos_logger
 from ._registrations import get_dbos_func_name
 
-_KafkaConsumerWorkflow = Callable[[KafkaMessage], None]
+_KafkaConsumerWorkflow = (
+    Callable[[KafkaMessage], None] | Callable[[KafkaMessage], Coroutine[Any, Any, None]]
+)
 
 _kafka_queue: Queue
 _in_order_kafka_queues: dict[str, Queue] = {}
@@ -37,8 +39,8 @@ def _kafka_consumer_loop(
     in_order: bool,
 ) -> None:
 
-    def on_error(err: KafkaError) -> NoReturn:
-        raise KafkaException(err)
+    def on_error(err: KafkaError) -> None:
+        dbos_logger.error(f"Exception in Kafka consumer: {err}")
 
     config["error_cb"] = on_error
     if "auto.offset.reset" not in config:

{dbos-2.2.0a3 → dbos-2.3.0}/dbos/_logger.py

@@ -68,30 +68,37 @@ def config_logger(config: "ConfigFile") -> None:
     )
     disable_otlp = config.get("telemetry", {}).get("disable_otlp", False)  # type: ignore
 
-    if not disable_otlp and otlp_logs_endpoints:
+    if not disable_otlp:
 
-        from opentelemetry._logs import set_logger_provider
+        from opentelemetry._logs import get_logger_provider, set_logger_provider
         from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
         from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
         from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
         from opentelemetry.sdk.resources import Resource
         from opentelemetry.semconv.attributes.service_attributes import SERVICE_NAME
 
-        log_provider = LoggerProvider(
-            Resource.create(
-                attributes={
-                    SERVICE_NAME: config["name"],
-                }
-            )
-        )
-        set_logger_provider(log_provider)
-        for e in otlp_logs_endpoints:
-            log_provider.add_log_record_processor(
-                BatchLogRecordProcessor(
-                    OTLPLogExporter(endpoint=e),
-                    export_timeout_millis=5000,
+        # Only set up OTLP provider and exporter if endpoints are provided
+        log_provider = get_logger_provider()
+        if otlp_logs_endpoints is not None:
+            if not isinstance(log_provider, LoggerProvider):
+                log_provider = LoggerProvider(
+                    Resource.create(
+                        attributes={
+                            SERVICE_NAME: config["name"],
+                        }
+                    )
+                )
+                set_logger_provider(log_provider)
+
+            for e in otlp_logs_endpoints:
+                log_provider.add_log_record_processor(
+                    BatchLogRecordProcessor(
+                        OTLPLogExporter(endpoint=e),
+                        export_timeout_millis=5000,
+                    )
                 )
-            )
+
+        # Even if no endpoints are provided, we still need a LoggerProvider to create the LoggingHandler
         global _otlp_handler
         _otlp_handler = LoggingHandler(logger_provider=log_provider)
 

{dbos-2.2.0a3 → dbos-2.3.0}/dbos/_scheduler.py

@@ -2,7 +2,7 @@ import random
 import threading
 import traceback
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING, Any, Callable, Coroutine
 
 from ._logger import dbos_logger
 from ._queue import Queue
@@ -14,7 +14,10 @@ from ._context import SetWorkflowID
 from ._croniter import croniter  # type: ignore
 from ._registrations import get_dbos_func_name
 
-ScheduledWorkflow = Callable[[datetime, datetime], None]
+ScheduledWorkflow = (
+    Callable[[datetime, datetime], None]
+    | Callable[[datetime, datetime], Coroutine[Any, Any, None]]
+)
 
 
 def scheduler_loop(

{dbos-2.2.0a3 → dbos-2.3.0}/dbos/_serialization.py

@@ -25,9 +25,13 @@ class Serializer(ABC):
 class DefaultSerializer(Serializer):
 
     def serialize(self, data: Any) -> str:
-        pickled_data: bytes = pickle.dumps(data)
-        encoded_data: str = base64.b64encode(pickled_data).decode("utf-8")
-        return encoded_data
+        try:
+            pickled_data: bytes = pickle.dumps(data)
+            encoded_data: str = base64.b64encode(pickled_data).decode("utf-8")
+            return encoded_data
+        except Exception as e:
+            dbos_logger.error(f"Error serializing object: {data}", exc_info=e)
+            raise
 
     def deserialize(cls, serialized_data: str) -> Any:
         pickled_data: bytes = base64.b64decode(serialized_data)

{dbos-2.2.0a3 → dbos-2.3.0}/dbos/_sys_db_postgres.py

@@ -41,7 +41,7 @@ class PostgresSystemDatabase(SystemDatabase):
                     parameters={"db_name": sysdb_name},
                 ).scalar():
                     dbos_logger.info(f"Creating system database {sysdb_name}")
-                    conn.execute(sa.text(f"CREATE DATABASE {sysdb_name}"))
+                    conn.execute(sa.text(f'CREATE DATABASE "{sysdb_name}"'))
             engine.dispose()
         else:
             # If we were provided an engine, validate it can connect

{dbos-2.2.0a3 → dbos-2.3.0}/dbos/_tracer.py

@@ -25,6 +25,10 @@ class DBOSTracer:
     def config(self, config: ConfigFile) -> None:
         self.otlp_attributes = config.get("telemetry", {}).get("otlp_attributes", {})  # type: ignore
         self.disable_otlp = config.get("telemetry", {}).get("disable_otlp", False)  # type: ignore
+        otlp_traces_endpoints = (
+            config.get("telemetry", {}).get("OTLPExporter", {}).get("tracesEndpoint")  # type: ignore
+        )
+
         if not self.disable_otlp:
             from opentelemetry import trace
             from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
@@ -38,25 +42,26 @@ class DBOSTracer:
             )
             from opentelemetry.semconv.attributes.service_attributes import SERVICE_NAME
 
-            if not isinstance(trace.get_tracer_provider(), TracerProvider):
-                resource = Resource(
-                    attributes={
-                        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)
-                otlp_traces_endpoints = (
-                    config.get("telemetry", {}).get("OTLPExporter", {}).get("tracesEndpoint")  # type: ignore
-                )
-                if otlp_traces_endpoints:
-                    for e in otlp_traces_endpoints:
-                        processor = BatchSpanProcessor(OTLPSpanExporter(endpoint=e))
-                        provider.add_span_processor(processor)
-                trace.set_tracer_provider(provider)
+            tracer_provider = trace.get_tracer_provider()
+
+            # Only set up OTLP provider and exporter if endpoints are provided
+            if otlp_traces_endpoints is not None:
+                if not isinstance(tracer_provider, TracerProvider):
+                    resource = Resource(
+                        attributes={
+                            SERVICE_NAME: config["name"],
+                        }
+                    )
+
+                    tracer_provider = TracerProvider(resource=resource)
+                    if os.environ.get("DBOS__CONSOLE_TRACES", None) is not None:
+                        processor = BatchSpanProcessor(ConsoleSpanExporter())
+                        tracer_provider.add_span_processor(processor)
+                    trace.set_tracer_provider(tracer_provider)
+
+                for e in otlp_traces_endpoints:
+                    processor = BatchSpanProcessor(OTLPSpanExporter(endpoint=e))
+                    tracer_provider.add_span_processor(processor)
 
     def set_provider(self, provider: "Optional[TracerProvider]") -> None:
         self.provider = provider

{dbos-2.2.0a3 → dbos-2.3.0}/dbos/cli/cli.py

@@ -140,26 +140,12 @@ def start() -> None:
             Forward kill signals to children.
 
             When we receive a signal, send it to the entire process group of the child.
-            If that doesn't work, SIGKILL them then exit.
             """
             # Send the signal to the child's entire process group
             if process.poll() is None:
                 os.killpg(os.getpgid(process.pid), signum)
 
-            # Give some time for the child to terminate
-            for _ in range(10):  # Wait up to 1 second
-                if process.poll() is not None:
-                    break
-                time.sleep(0.1)
-
-            # If the child is still running, force kill it
-            if process.poll() is None:
-                try:
-                    os.killpg(os.getpgid(process.pid), signal.SIGKILL)
-                except Exception:
-                    pass
-
-            # Exit immediately
+            # Exit 
             os._exit(process.returncode if process.returncode is not None else 1)
 
         # Configure the single handler only on Unix-like systems.

{dbos-2.2.0a3 → dbos-2.3.0}/pyproject.toml

@@ -34,7 +34,7 @@ classifiers = [
     "Topic :: Software Development :: Libraries :: Python Modules",
     "Framework :: AsyncIO",
 ]
-version = "2.2.0a3"
+version = "2.3.0"
 
 [project.license]
 text = "MIT"

{dbos-2.2.0a3 → dbos-2.3.0}/tests/test_client.py

@@ -593,3 +593,35 @@ def test_enqueue_with_priority(dbos: DBOS, client: DBOSClient) -> None:
 def test_client_bad_url() -> None:
     with pytest.raises(DBAPIError) as exc_info:
         DBOSClient("postgresql://postgres:fakepassword@localhost:5433/fake_database")
+
+
+def test_client_auth(dbos: DBOS, client: DBOSClient) -> None:
+    run_client_collateral()
+
+    johnDoe: Person = {"first": "John", "last": "Doe", "age": 30}
+    wfid = str(uuid.uuid4())
+
+    user = "testuser"
+    roles = ["role1", "role2"]
+
+    options: EnqueueOptions = {
+        "queue_name": "test_queue",
+        "workflow_name": "enqueue_test",
+        "workflow_id": wfid,
+        "authenticated_user": user,
+        "authenticated_roles": roles,
+    }
+
+    handle: WorkflowHandle[str] = client.enqueue(options, 42, "test", johnDoe)
+    result = handle.get_result()
+    assert result == '42-test-{"first": "John", "last": "Doe", "age": 30}'
+
+    list_results = client.list_workflows()
+    assert len(list_results) == 1
+    assert list_results[0].workflow_id == wfid
+    assert list_results[0].status == "SUCCESS"
+    assert list_results[0].output == result
+    assert list_results[0].input is not None
+    assert list_results[0].authenticated_user == user
+    assert list_results[0].authenticated_roles == roles
+    assert list_results[0].assumed_role is None

{dbos-2.2.0a3 → dbos-2.3.0}/tests/test_config.py

@@ -209,7 +209,7 @@ def test_process_config_full():
         "max_overflow": 0,
         "pool_size": 20,
         "pool_pre_ping": True,
-        "connect_args": {"connect_timeout": 1},
+        "connect_args": {"connect_timeout": 1, "application_name": "dbos_transact"},
     }
     assert configFile["database"]["sys_db_engine_kwargs"] == {
         "key": "value",
@@ -217,7 +217,7 @@ def test_process_config_full():
         "max_overflow": 0,
         "pool_size": 27,
         "pool_pre_ping": True,
-        "connect_args": {"connect_timeout": 1},
+        "connect_args": {"connect_timeout": 1, "application_name": "dbos_transact"},
     }
     assert configFile["runtimeConfig"]["start"] == ["python3 main.py"]
     assert configFile["runtimeConfig"]["admin_port"] == 8001
@@ -255,7 +255,7 @@ def test_process_config_system_database():
         "max_overflow": 0,
         "pool_size": 20,
         "pool_pre_ping": True,
-        "connect_args": {"connect_timeout": 1},
+        "connect_args": {"connect_timeout": 1, "application_name": "dbos_transact"},
     }
     assert configFile["database"]["sys_db_engine_kwargs"] == {
         "key": "value",
@@ -263,7 +263,7 @@ def test_process_config_system_database():
         "max_overflow": 0,
         "pool_size": 27,
         "pool_pre_ping": True,
-        "connect_args": {"connect_timeout": 1},
+        "connect_args": {"connect_timeout": 1, "application_name": "dbos_transact"},
     }
 
 
@@ -397,14 +397,14 @@ def test_configure_db_engine_parameters_defaults():
         "max_overflow": 0,
         "pool_size": 20,
         "pool_pre_ping": True,
-        "connect_args": {"connect_timeout": 10},
+        "connect_args": {"connect_timeout": 10, "application_name": "dbos_transact"},
     }
     assert data["sys_db_engine_kwargs"] == {
         "pool_timeout": 30,
         "max_overflow": 0,
         "pool_size": 20,
         "pool_pre_ping": True,
-        "connect_args": {"connect_timeout": 10},
+        "connect_args": {"connect_timeout": 10, "application_name": "dbos_transact"},
     }
 
 
@@ -419,14 +419,14 @@ def test_configure_db_engine_parameters_custom_sys_db_pool_sizes():
         "max_overflow": 0,
         "pool_size": 20,
         "pool_pre_ping": True,
-        "connect_args": {"connect_timeout": 10},
+        "connect_args": {"connect_timeout": 10, "application_name": "dbos_transact"},
     }
     assert data["sys_db_engine_kwargs"] == {
         "pool_timeout": 30,
         "max_overflow": 0,
         "pool_size": 35,
         "pool_pre_ping": True,
-        "connect_args": {"connect_timeout": 10},
+        "connect_args": {"connect_timeout": 10, "application_name": "dbos_transact"},
     }
 
 
@@ -440,7 +440,11 @@ def test_configure_db_engine_parameters_user_kwargs_override():
             "pool_pre_ping": True,
             "custom_param": "value",
             "pool_size": 50,
-            "connect_args": {"connect_timeout": 30, "key": "value"},
+            "connect_args": {
+                "connect_timeout": 30,
+                "key": "value",
+                "application_name": "dbos_transact",
+            },
         },
     }
 
@@ -453,7 +457,11 @@ def test_configure_db_engine_parameters_user_kwargs_override():
         "pool_pre_ping": True,
         "custom_param": "value",
         "pool_size": 50,
-        "connect_args": {"connect_timeout": 30, "key": "value"},
+        "connect_args": {
+            "connect_timeout": 30,
+            "key": "value",
+            "application_name": "dbos_transact",
+        },
     }
 
     # System engine kwargs should use system pool size but same user overrides
@@ -463,7 +471,11 @@ def test_configure_db_engine_parameters_user_kwargs_override():
         "pool_pre_ping": True,
         "custom_param": "value",
         "pool_size": 35,
-        "connect_args": {"connect_timeout": 30, "key": "value"},
+        "connect_args": {
+            "connect_timeout": 30,
+            "key": "value",
+            "application_name": "dbos_transact",
+        },
     }
 
 
@@ -487,7 +499,7 @@ def test_configure_db_engine_parameters_user_kwargs_and_db_url_connect_timeout()
         "pool_pre_ping": True,
         "custom_param": "value",
         "pool_size": 50,
-        "connect_args": {"connect_timeout": 22},
+        "connect_args": {"connect_timeout": 22, "application_name": "dbos_transact"},
     }
 
     # System engine kwargs should use system pool size but same user overrides
@@ -497,7 +509,7 @@ def test_configure_db_engine_parameters_user_kwargs_and_db_url_connect_timeout()
         "pool_pre_ping": True,
         "custom_param": "value",
         "pool_size": 50,
-        "connect_args": {"connect_timeout": 22},
+        "connect_args": {"connect_timeout": 22, "application_name": "dbos_transact"},
     }
 
 
@@ -556,7 +568,7 @@ def test_configure_db_engine_parameters_user_kwargs_mixed_params():
         "pool_pre_ping": True,
         "custom_param": "value",
         "pool_size": 50,
-        "connect_args": {"connect_timeout": 10},
+        "connect_args": {"connect_timeout": 10, "application_name": "dbos_transact"},
     }
 
     # System engine kwargs should use system pool size but same user overrides
@@ -566,7 +578,7 @@ def test_configure_db_engine_parameters_user_kwargs_mixed_params():
         "pool_pre_ping": True,
         "custom_param": "value",
         "pool_size": 50,
-        "connect_args": {"connect_timeout": 10},
+        "connect_args": {"connect_timeout": 10, "application_name": "dbos_transact"},
     }
 
 
@@ -581,14 +593,14 @@ def test_configure_db_engine_parameters_empty_user_kwargs():
         "max_overflow": 0,
         "pool_size": 20,
         "pool_pre_ping": True,
-        "connect_args": {"connect_timeout": 10},
+        "connect_args": {"connect_timeout": 10, "application_name": "dbos_transact"},
     }
     assert data["sys_db_engine_kwargs"] == {
         "pool_timeout": 30,
         "max_overflow": 0,
         "pool_size": 20,
         "pool_pre_ping": True,
-        "connect_args": {"connect_timeout": 10},
+        "connect_args": {"connect_timeout": 10, "application_name": "dbos_transact"},
     }
 
 
@@ -598,24 +610,6 @@ def test_configure_db_engine_parameters_empty_user_kwargs():
 
 
 def test_process_config_with_wrong_db_url():
-    # Missing username
-    config: ConfigFile = {
-        "name": "some-app",
-        "database_url": "postgres://:password@h:1234/dbname",
-    }
-    with pytest.raises(DBOSInitializationError) as exc_info:
-        process_config(data=config)
-    assert "Username must be specified in the connection URL" in str(exc_info.value)
-
-    # Missing host
-    config: ConfigFile = {
-        "name": "some-app",
-        "database_url": "postgres://user:password@:1234/dbname",
-    }
-    with pytest.raises(DBOSInitializationError) as exc_info:
-        process_config(data=config)
-    assert "Host must be specified in the connection URL" in str(exc_info.value)
-
     # Missing dbname
     config: ConfigFile = {
         "name": "some-app",

{dbos-2.2.0a3 → dbos-2.3.0}/tests/test_dbos.py

@@ -37,6 +37,7 @@ from dbos._error import (
 from dbos._schemas.system_database import SystemSchema
 from dbos._sys_db import GetWorkflowsInput
 from dbos._utils import GlobalParams
+from tests.conftest import using_sqlite
 
 
 def test_simple_workflow(dbos: DBOS) -> None:
@@ -1796,6 +1797,65 @@ def test_without_appdb(config: DBOSConfig, cleanup_test_databases: None) -> None
         assert s["function_name"] == step.__qualname__
 
 
+def test_custom_database(
+    config: DBOSConfig, db_engine: sa.Engine, cleanup_test_databases: None
+) -> None:
+    DBOS.destroy(destroy_registry=True)
+    assert config["system_database_url"]
+    custom_database = "F8nny_dAtaB@s3@-n@m3.sqlite"
+    url = sa.make_url(config["system_database_url"])
+    url = url.set(database=custom_database)
+    config["system_database_url"] = url.render_as_string(hide_password=False)
+    # Destroy the database if it exists
+    if using_sqlite():
+        parsed_url = sa.make_url(config["system_database_url"])
+        db_path = parsed_url.database
+        assert db_path is not None
+        if os.path.exists(db_path):
+            os.remove(db_path)
+    else:
+        with db_engine.connect() as connection:
+            connection.execution_options(isolation_level="AUTOCOMMIT")
+            connection.execute(
+                sa.text(f'DROP DATABASE IF EXISTS "{custom_database}" WITH (FORCE)')
+            )
+    DBOS(config=config)
+    DBOS.launch()
+
+    key = "key"
+    val = "val"
+
+    @DBOS.transaction()
+    def transaction() -> None:
+        return
+
+    @DBOS.workflow()
+    def recv_workflow() -> Any:
+        transaction()
+        DBOS.set_event(key, val)
+        return DBOS.recv()
+
+    handle = DBOS.start_workflow(recv_workflow)
+    assert DBOS.get_event(handle.workflow_id, key) == val
+    DBOS.send(handle.workflow_id, val)
+    assert handle.get_result() == val
+    assert len(DBOS.list_workflows()) == 2
+    steps = DBOS.list_workflow_steps(handle.workflow_id)
+    assert len(steps) == 4
+    assert "transaction" in steps[0]["function_name"]
+    DBOS.destroy(destroy_registry=True)
+
+    # Test custom database with client
+    client = DBOSClient(
+        system_database_url=config["system_database_url"],
+        application_database_url=config["application_database_url"],
+    )
+    assert len(client.list_workflows()) == 2
+    steps = client.list_workflow_steps(handle.workflow_id)
+    assert len(steps) == 4
+    assert "transaction" in steps[0]["function_name"]
+
+
 def test_custom_schema(
     config: DBOSConfig, cleanup_test_databases: None, skip_with_sqlite: None
 ) -> None:

{dbos-2.2.0a3 → dbos-2.3.0}/tests/test_failures.py

@@ -1,7 +1,7 @@
 import threading
 import time
 import uuid
-from typing import cast
+from typing import Any, Generator, cast
 
 import pytest
 import sqlalchemy as sa
@@ -553,3 +553,16 @@ def test_unregistered_workflow(dbos: DBOS, config: DBOSConfig) -> None:
 
     with pytest.raises(DBOSWorkflowFunctionNotFoundError):
         DBOS._recover_pending_workflows()
+
+
+def test_nonserializable_return(dbos: DBOS) -> None:
+    @DBOS.step()
+    def step() -> Generator[str, Any, None]:
+        yield "val"
+
+    @DBOS.workflow()
+    def workflow() -> None:
+        step()
+
+    with pytest.raises(TypeError):
+        workflow()

{dbos-2.2.0a3 → dbos-2.3.0}/tests/test_kafka.py

@@ -12,24 +12,25 @@ from dbos import DBOS, KafkaMessage
 # Without it, they're automatically skipped.
 # Here's a docker-compose script you can use to set up local Kafka:
 
-# version: "3.7"
 # services:
 #   broker:
-#       image: bitnami/kafka:latest
+#       image: apache/kafka:latest
 #       hostname: broker
 #       container_name: broker
 #       ports:
 #         - '9092:9092'
 #       environment:
-#         KAFKA_CFG_NODE_ID: 1
-#         KAFKA_CFG_LISTENER_SECURITY_PROTOCOL_MAP: 'CONTROLLER:PLAINTEXT,PLAINTEXT:PLAINTEXT,PLAINTEXT_HOST:PLAINTEXT'
-#         KAFKA_CFG_ADVERTISED_LISTENERS: 'PLAINTEXT_HOST://localhost:9092,PLAINTEXT://broker:19092'
-#         KAFKA_CFG_PROCESS_ROLES: 'broker,controller'
-#         KAFKA_CFG_CONTROLLER_QUORUM_VOTERS: '1@broker:29093'
-#         KAFKA_CFG_LISTENERS: 'CONTROLLER://:29093,PLAINTEXT_HOST://:9092,PLAINTEXT://:19092'
-#         KAFKA_CFG_INTER_BROKER_LISTENER_NAME: 'PLAINTEXT'
-#         KAFKA_CFG_CONTROLLER_LISTENER_NAMES: 'CONTROLLER'
-
+#           KAFKA_NODE_ID: 1
+#           KAFKA_LISTENERS: PLAINTEXT://0.0.0.0:9092,CONTROLLER://0.0.0.0:9093
+#           KAFKA_ADVERTISED_LISTENERS: PLAINTEXT://127.0.0.1:9092
+#           KAFKA_PROCESS_ROLES: broker,controller
+#           KAFKA_CONTROLLER_QUORUM_VOTERS: 1@localhost:9093
+#           KAFKA_CONTROLLER_LISTENER_NAMES: CONTROLLER
+#           KAFKA_LISTENER_SECURITY_PROTOCOL_MAP: PLAINTEXT:PLAINTEXT,CONTROLLER:PLAINTEXT
+#           KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR: 1
+#           KAFKA_TRANSACTION_STATE_LOG_REPLICATION_FACTOR: 1
+#           KAFKA_TRANSACTION_STATE_LOG_MIN_ISR: 1
+#           CLUSTER_ID: MkU3OEVBNTcwNTJENDM2Qk
 
 NUM_EVENTS = 3
 
@@ -81,12 +82,44 @@ def test_kafka(dbos: DBOS) -> None:
         assert b"test message key" in msg.key  # type: ignore
         assert b"test message value" in msg.value  # type: ignore
         print(msg)
-        if kafka_count == 3:
+        if kafka_count == NUM_EVENTS:
+            event.set()
+
+    wait = event.wait(timeout=10)
+    assert wait
+    assert kafka_count == NUM_EVENTS
+
+
+def test_kafka_async(dbos: DBOS) -> None:
+    event = threading.Event()
+    kafka_count = 0
+    server = "localhost:9092"
+    topic = f"dbos-kafka-{random.randrange(1_000_000_000)}"
+
+    if not send_test_messages(server, topic):
+        pytest.skip("Kafka not available")
+
+    @DBOS.kafka_consumer(
+        {
+            "bootstrap.servers": server,
+            "group.id": "dbos-test",
+            "auto.offset.reset": "earliest",
+        },
+        [topic],
+    )
+    @DBOS.workflow()
+    async def test_kafka_workflow(msg: KafkaMessage) -> None:
+        nonlocal kafka_count
+        kafka_count += 1
+        assert b"test message key" in msg.key  # type: ignore
+        assert b"test message value" in msg.value  # type: ignore
+        print(msg)
+        if kafka_count == NUM_EVENTS:
             event.set()
 
     wait = event.wait(timeout=10)
     assert wait
-    assert kafka_count == 3
+    assert kafka_count == NUM_EVENTS
 
 
 def test_kafka_in_order(dbos: DBOS) -> None:
@@ -114,12 +147,12 @@ def test_kafka_in_order(dbos: DBOS) -> None:
         kafka_count += 1
         assert f"test message key {kafka_count - 1}".encode() == msg.key
         print(msg)
-        if kafka_count == 3:
+        if kafka_count == NUM_EVENTS:
             event.set()
 
     wait = event.wait(timeout=15)
     assert wait
-    assert kafka_count == 3
+    assert kafka_count == NUM_EVENTS
     time.sleep(2)  # Wait for things to clean up
 
 
@@ -150,9 +183,9 @@ def test_kafka_no_groupid(dbos: DBOS) -> None:
         assert b"test message key" in msg.key  # type: ignore
         assert b"test message value" in msg.value  # type: ignore
         print(msg)
-        if kafka_count == 6:
+        if kafka_count == NUM_EVENTS * 2:
             event.set()
 
     wait = event.wait(timeout=10)
     assert wait
-    assert kafka_count == 6
+    assert kafka_count == NUM_EVENTS * 2

{dbos-2.2.0a3 → dbos-2.3.0}/tests/test_queue.py

@@ -1476,7 +1476,7 @@ def test_unsetting_timeout(dbos: DBOS) -> None:
             queue.enqueue(child)
 
     child_one, child_two = str(uuid.uuid4()), str(uuid.uuid4())
-    with SetWorkflowTimeout(1.0):
+    with SetWorkflowTimeout(2.0):
         queue.enqueue(parent, child_one, child_two).get_result()
 
     # Verify child one, which has a propagated timeout, is cancelled

{dbos-2.2.0a3 → dbos-2.3.0}/tests/test_scheduler.py

@@ -105,6 +105,19 @@ def test_scheduled_workflow(dbos: DBOS) -> None:
     assert wf_counter > 1 and wf_counter <= 5
 
 
+def test_async_scheduled_workflow(dbos: DBOS) -> None:
+    wf_counter: int = 0
+
+    @DBOS.scheduled("* * * * * *")
+    @DBOS.workflow()
+    async def test_workflow(scheduled: datetime, actual: datetime) -> None:
+        nonlocal wf_counter
+        wf_counter += 1
+
+    time.sleep(5)
+    assert wf_counter > 1 and wf_counter <= 5
+
+
 def test_appdb_downtime(dbos: DBOS, skip_with_sqlite: None) -> None:
     wf_counter: int = 0
 

{dbos-2.2.0a3 → dbos-2.3.0}/tests/test_spans.py

@@ -5,7 +5,6 @@ import pytest
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 from inline_snapshot import snapshot
-from opentelemetry._logs import set_logger_provider
 from opentelemetry.sdk import trace as tracesdk
 from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
 from opentelemetry.sdk._logs.export import BatchLogRecordProcessor, InMemoryLogExporter
@@ -17,6 +16,7 @@ from dbos import DBOS, DBOSConfig
 from dbos._logger import dbos_logger
 from dbos._tracer import dbos_tracer
 from dbos._utils import GlobalParams
+from tests.conftest import default_config
 
 
 @dataclass
@@ -65,7 +65,6 @@ def test_spans(config: DBOSConfig) -> None:
     log_processor = BatchLogRecordProcessor(log_exporter)
     log_provider = LoggerProvider()
     log_provider.add_log_record_processor(log_processor)
-    set_logger_provider(log_provider)
     dbos_logger.addHandler(LoggingHandler(logger_provider=log_provider))
 
     test_workflow()
@@ -192,7 +191,6 @@ async def test_spans_async(dbos: DBOS) -> None:
     log_processor = BatchLogRecordProcessor(log_exporter)
     log_provider = LoggerProvider()
     log_provider.add_log_record_processor(log_processor)
-    set_logger_provider(log_provider)
     dbos_logger.addHandler(LoggingHandler(logger_provider=log_provider))
 
     await test_workflow()
@@ -303,7 +301,6 @@ def test_wf_fastapi(dbos_fastapi: Tuple[DBOS, FastAPI]) -> None:
     log_processor = BatchLogRecordProcessor(log_exporter)
     log_provider = LoggerProvider()
     log_provider.add_log_record_processor(log_processor)
-    set_logger_provider(log_provider)
     dbos_logger.addHandler(LoggingHandler(logger_provider=log_provider))
 
     client = TestClient(app)
@@ -378,7 +375,6 @@ def test_disable_otlp_no_spans(config: DBOSConfig) -> None:
     log_processor = BatchLogRecordProcessor(log_exporter)
     log_provider = LoggerProvider()
     log_provider.add_log_record_processor(log_processor)
-    set_logger_provider(log_provider)
     dbos_logger.addHandler(LoggingHandler(logger_provider=log_provider))
 
     test_workflow()