Python Fixes (#493)

- Fix an issue where DBOS would not create a system database with
special characters in its name
- Fix typing for async scheduled and Kafka workflows
- Improve error messages for unserializable workflow inputs and outputs
and step outputs
- Fix an issue where DBOS does not reconnect after losing its connection
to Kafka

Author: kraftp

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,

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 ._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/_scheduler.py

@@ -2,7 +2,7 @@
 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 ._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/_serialization.py

@@ -25,9 +25,13 @@ def deserialize(cls, serialized_data: str) -> Any:
 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/_sys_db_postgres.py

@@ -41,7 +41,7 @@ def run_migrations(self) -> None:
                     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

tests/test_config.py

@@ -209,15 +209,15 @@ 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",
         "pool_timeout": 30,
         "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,15 +255,15 @@ 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",
         "pool_timeout": 30,
         "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"},
     }
 
 

tests/test_dbos.py

@@ -37,6 +37,7 @@
 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 transaction() -> 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@[email protected]"
+    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:

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 workflow() -> 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()