dbos 2.3.0a1__tar.gz → 2.3.0a2__tar.gz

{dbos-2.3.0a1 → dbos-2.3.0a2}/PKG-INFO

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

{dbos-2.3.0a1 → dbos-2.3.0a2}/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-2.3.0a1 → dbos-2.3.0a2}/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.3.0a1 → dbos-2.3.0a2}/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.3.0a1 → dbos-2.3.0a2}/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.3.0a1 → dbos-2.3.0a2}/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.3.0a1 → dbos-2.3.0a2}/pyproject.toml

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

{dbos-2.3.0a1 → dbos-2.3.0a2}/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"},
     }
 
 

{dbos-2.3.0a1 → dbos-2.3.0a2}/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.3.0a1 → dbos-2.3.0a2}/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.3.0a1 → dbos-2.3.0a2}/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.3.0a1 → dbos-2.3.0a2}/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