dbos 0.6.0a3__tar.gz → 0.7.0a0__tar.gz

{dbos-0.6.0a3 → dbos-0.7.0a0}/PKG-INFO

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

{dbos-0.6.0a3 → dbos-0.7.0a0}/dbos/__init__.py

@@ -2,6 +2,7 @@ from . import error as error
 from .context import DBOSContextEnsure, SetWorkflowID
 from .dbos import DBOS, DBOSConfiguredInstance, WorkflowHandle, WorkflowStatus
 from .dbos_config import ConfigFile, get_dbos_database_url, load_config
+from .kafka_message import KafkaMessage
 from .system_database import GetWorkflowsInput, WorkflowStatusString
 
 __all__ = [
@@ -10,6 +11,7 @@ __all__ = [
     "DBOSConfiguredInstance",
     "DBOSContextEnsure",
     "GetWorkflowsInput",
+    "KafkaMessage",
     "SetWorkflowID",
     "WorkflowHandle",
     "WorkflowStatus",

{dbos-0.6.0a3 → dbos-0.7.0a0}/dbos/dbos.py

@@ -54,11 +54,14 @@ from .tracer import dbos_tracer
 
 if TYPE_CHECKING:
     from fastapi import FastAPI
+    from dbos.kafka import KafkaConsumerWorkflow
     from .request import Request
     from flask import Flask
 
 from sqlalchemy.orm import Session
 
+from dbos.request import Request
+
 if sys.version_info < (3, 10):
     from typing_extensions import ParamSpec, TypeAlias
 else:
@@ -506,6 +509,20 @@ class DBOS:
 
         return scheduled(_get_or_create_dbos_registry(), cron)
 
+    @classmethod
+    def kafka_consumer(
+        cls, config: dict[str, Any], topics: list[str]
+    ) -> Callable[[KafkaConsumerWorkflow], KafkaConsumerWorkflow]:
+        """Decorate a function to be used as a Kafka consumer."""
+        try:
+            from dbos.kafka import kafka_consumer
+
+            return kafka_consumer(_get_or_create_dbos_registry(), config, topics)
+        except ModuleNotFoundError as e:
+            raise DBOSException(
+                f"{e.name} dependency not found. Please install {e.name} via your package manager."
+            ) from e
+
     @classmethod
     def start_workflow(
         cls,
@@ -714,6 +731,12 @@ class DBOS:
         ctx = assert_current_dbos_context()
         return ctx.authenticated_roles
 
+    @classproperty
+    def assumed_role(cls) -> Optional[str]:
+        """Return the role currently assumed by the authenticated user, if any, associated with the current context."""
+        ctx = assert_current_dbos_context()
+        return ctx.assumed_role
+
     @classmethod
     def set_authentication(
         cls, authenticated_user: Optional[str], authenticated_roles: Optional[List[str]]
@@ -800,13 +823,9 @@ class DBOSConfiguredInstance:
 
     """
 
-    def __init__(self, config_name: str, dbos: Optional[DBOS] = None) -> None:
+    def __init__(self, config_name: str) -> None:
         self.config_name = config_name
-        if dbos is not None:
-            assert isinstance(dbos, DBOS)
-            dbos._registry.register_instance(self)
-        else:
-            DBOS.register_instance(self)
+        DBOS.register_instance(self)
 
 
 # Apps that import DBOS probably don't exit.  If they do, let's see if

{dbos-0.6.0a3 → dbos-0.7.0a0}/dbos/dbos_config.py

@@ -2,7 +2,7 @@ import json
 import os
 import re
 from importlib import resources
-from typing import Dict, List, Optional, TypedDict
+from typing import Any, Dict, List, Optional, TypedDict
 
 import yaml
 from jsonschema import ValidationError, validate
@@ -69,6 +69,7 @@ class ConfigFile(TypedDict, total=False):
     database: DatabaseConfig
     telemetry: Optional[TelemetryConfig]
     env: Dict[str, str]
+    application: Dict[str, Any]
 
 
 def substitute_env_vars(content: str) -> str:

dbos-0.7.0a0/dbos/kafka.py

@@ -0,0 +1,94 @@
+import threading
+import traceback
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Generator, NoReturn, Optional, Union
+
+from confluent_kafka import Consumer, KafkaError, KafkaException
+from confluent_kafka import Message as CTypeMessage
+
+if TYPE_CHECKING:
+    from dbos.dbos import _DBOSRegistry
+
+from .context import SetWorkflowID
+from .kafka_message import KafkaMessage
+from .logger import dbos_logger
+
+KafkaConsumerWorkflow = Callable[[KafkaMessage], None]
+
+
+def _kafka_consumer_loop(
+    func: KafkaConsumerWorkflow,
+    config: dict[str, Any],
+    topics: list[str],
+    stop_event: threading.Event,
+) -> None:
+
+    def on_error(err: KafkaError) -> NoReturn:
+        raise KafkaException(err)
+
+    config["error_cb"] = on_error
+    if "auto.offset.reset" not in config:
+        config["auto.offset.reset"] = "earliest"
+
+    consumer = Consumer(config)
+    try:
+        consumer.subscribe(topics)
+        while not stop_event.is_set():
+            cmsg = consumer.poll(1.0)
+
+            if stop_event.is_set():
+                return
+
+            if cmsg is None:
+                continue
+
+            err = cmsg.error()
+            if err is not None:
+                dbos_logger.error(
+                    f"Kafka error {err.code()} ({err.name()}): {err.str()}"
+                )
+                # fatal errors require an updated consumer instance
+                if err.code() == KafkaError._FATAL or err.fatal():
+                    original_consumer = consumer
+                    try:
+                        consumer = Consumer(config)
+                        consumer.subscribe(topics)
+                    finally:
+                        original_consumer.close()
+            else:
+                msg = KafkaMessage(
+                    headers=cmsg.headers(),
+                    key=cmsg.key(),
+                    latency=cmsg.latency(),
+                    leader_epoch=cmsg.leader_epoch(),
+                    offset=cmsg.offset(),
+                    partition=cmsg.partition(),
+                    timestamp=cmsg.timestamp(),
+                    topic=cmsg.topic(),
+                    value=cmsg.value(),
+                )
+                with SetWorkflowID(
+                    f"kafka-unique-id-{msg.topic}-{msg.partition}-{msg.offset}"
+                ):
+                    try:
+                        func(msg)
+                    except Exception as e:
+                        dbos_logger.error(
+                            f"Exception encountered in Kafka consumer: {traceback.format_exc()}"
+                        )
+
+    finally:
+        consumer.close()
+
+
+def kafka_consumer(
+    dbosreg: "_DBOSRegistry", config: dict[str, Any], topics: list[str]
+) -> Callable[[KafkaConsumerWorkflow], KafkaConsumerWorkflow]:
+    def decorator(func: KafkaConsumerWorkflow) -> KafkaConsumerWorkflow:
+        stop_event = threading.Event()
+        dbosreg.register_poller(
+            stop_event, _kafka_consumer_loop, func, config, topics, stop_event
+        )
+        return func
+
+    return decorator

dbos-0.7.0a0/dbos/kafka_message.py

@@ -0,0 +1,15 @@
+from dataclasses import dataclass
+from typing import Optional, Union
+
+
+@dataclass
+class KafkaMessage:
+    headers: Optional[list[tuple[str, Union[str, bytes]]]]
+    key: Optional[Union[str, bytes]]
+    latency: Optional[float]
+    leader_epoch: Optional[int]
+    offset: Optional[int]
+    partition: Optional[int]
+    timestamp: tuple[int, int]
+    topic: Optional[str]
+    value: Optional[Union[str, bytes]]

{dbos-0.6.0a3 → dbos-0.7.0a0}/pyproject.toml

@@ -23,7 +23,7 @@ dependencies = [
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "0.6.0a3"
+version = "0.7.0a0"
 
 [project.license]
 text = "MIT"
@@ -60,12 +60,13 @@ dev = [
     "httpx>=0.27.0",
     "types-psycopg2>=2.9.21.20240417",
     "pytz>=2024.1",
-    "pdm>=2.18.0",
-    "pdm-backend>=2.3.3",
     "GitPython>=3.1.43",
+    "confluent-kafka>=2.5.3",
+    "types-confluent-kafka>=1.2.2",
     "flask>=3.0.3",
     "pytest-order>=1.3.0",
     "pyjwt>=2.9.0",
+    "pdm-backend>=2.3.3",
 ]
 
 [tool.black]

{dbos-0.6.0a3 → dbos-0.7.0a0}/tests/classdefs.py

@@ -61,7 +61,7 @@ class DBOSTestClass(DBOSConfiguredInstance):
     @DBOS.workflow()
     @DBOS.required_roles(["admin"])
     def test_func_admin(self, var: str) -> str:
-        assert assert_current_dbos_context().assumed_role == "admin"
+        assert DBOS.assumed_role == "admin"
         return self.config_name + ":" + var
 
 

{dbos-0.6.0a3 → dbos-0.7.0a0}/tests/conftest.py

@@ -124,7 +124,7 @@ def dbos_fastapi(
         warnings.filterwarnings(
             "ignore",
             category=DeprecationWarning,
-            message="\s*on_event is deprecated, use lifespan event handlers instead\.",
+            message=r"\s*on_event is deprecated, use lifespan event handlers instead\.",
         )
         dbos = DBOS(fastapi=app, config=config)
 

{dbos-0.6.0a3 → dbos-0.7.0a0}/tests/test_classdecorators.py

@@ -49,7 +49,7 @@ def test_required_roles_class(dbos: DBOS) -> None:
     @DBOS.default_required_roles(["user"])
     class DBOSTestClassRR(DBOSConfiguredInstance):
         def __init__(self) -> None:
-            super().__init__("myconfig", dbos)
+            super().__init__("myconfig")
 
         @DBOS.workflow()
         def test_func_user(self, var: str) -> str:
@@ -279,7 +279,7 @@ def test_simple_workflow_inst(dbos: DBOS) -> None:
     @DBOS.dbos_class()
     class DBOSTestClassInst(DBOSConfiguredInstance):
         def __init__(self) -> None:
-            super().__init__("bob", dbos)
+            super().__init__("bob")
             self.txn_counter: int = 0
             self.wf_counter: int = 0
             self.step_counter: int = 0
@@ -327,7 +327,7 @@ def test_simple_workflow_inst(dbos: DBOS) -> None:
 def test_forgotten_decorator(dbos: DBOS) -> None:
     class DBOSTestRegErr(DBOSConfiguredInstance):
         def __init__(self) -> None:
-            super().__init__("bob", dbos)
+            super().__init__("bob")
             self.txn_counter: int = 0
             self.wf_counter: int = 0
             self.step_counter: int = 0
@@ -362,7 +362,7 @@ def test_duplicate_reg(dbos: DBOS) -> None:
     @DBOS.dbos_class()
     class DBOSTestRegDup(DBOSConfiguredInstance):
         def __init__(self) -> None:
-            super().__init__("bob", dbos)
+            super().__init__("bob")
 
     # Duplicate class registration
     with pytest.raises(Exception) as exc_info:
@@ -370,7 +370,7 @@ def test_duplicate_reg(dbos: DBOS) -> None:
         @DBOS.dbos_class()
         class DBOSTestRegDup(DBOSConfiguredInstance):  # type: ignore
             def __init__(self) -> None:
-                super().__init__("bob", dbos)
+                super().__init__("bob")
 
     assert "Duplicate type registration for class 'DBOSTestRegDup'" == str(
         exc_info.value
@@ -419,7 +419,7 @@ def test_inst_recovery(dbos: DBOS) -> None:
     @DBOS.dbos_class()
     class DBOSTestInstRec(DBOSConfiguredInstance):
         def __init__(self) -> None:
-            super().__init__("bob", dbos)
+            super().__init__("bob")
 
         @DBOS.workflow()
         def check_inst(self, arg1: str) -> str:

dbos-0.7.0a0/tests/test_kafka.py

@@ -0,0 +1,80 @@
+import random
+import threading
+import uuid
+from typing import Any, List, NoReturn
+
+import pytest
+from confluent_kafka import KafkaError, Producer
+
+from dbos import DBOS, KafkaMessage
+
+# These tests require local Kafka to run.
+# 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
+#       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'
+
+
+def send_test_messages(server: str, topic: str) -> bool:
+
+    try:
+
+        def on_error(err: KafkaError) -> NoReturn:
+            raise Exception(err)
+
+        producer = Producer({"bootstrap.servers": server, "error_cb": on_error})
+
+        producer.produce(topic, key=f"test message key", value=f"test message value")
+
+        producer.poll(10)
+        producer.flush(10)
+        return True
+    except Exception as e:
+        return False
+    finally:
+        pass
+
+
+def test_kafka(dbos: DBOS) -> None:
+    event = threading.Event()
+    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")
+
+    messages: List[KafkaMessage] = []
+
+    @DBOS.kafka_consumer(
+        {
+            "bootstrap.servers": server,
+            "group.id": "dbos-test",
+            "auto.offset.reset": "earliest",
+        },
+        [topic],
+    )
+    @DBOS.workflow()
+    def test_kafka_workflow(msg: KafkaMessage) -> None:
+        print(msg)
+        messages.append(msg)
+        event.set()
+
+    wait = event.wait(timeout=10)
+    assert wait
+    assert len(messages) > 0

{dbos-0.6.0a3 → dbos-0.7.0a0}/tests/test_singleton.py

@@ -146,6 +146,12 @@ database:
 runtimeConfig:
   start:
     - python3 main.py
+application:
+  service_url: 'https://service.org'
+  service_config:
+    port: 80
+    user: "user"
+    password: "password"
 """
 
 
@@ -164,6 +170,12 @@ def test_config_before_singleton(cleanup_test_databases: None) -> None:
 
         x = DBOS.config.get("language")
         assert x == "python"
+        y = DBOS.config["language"]
+        assert y == "python"
+        url = DBOS.config["application"]["service_url"]
+        assert url == "https://service.org"
+        port = DBOS.config["application"]["service_config"]["port"]
+        assert port == 80
 
         # This is OK, it meant load_config anyway
         dbos: DBOS = DBOS()