dbos 0.6.0a4__tar.gz → 0.7.0a0__tar.gz

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

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

{dbos-0.6.0a4 → 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.0a4 → 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,

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.0a4 → dbos-0.7.0a0}/pyproject.toml

@@ -23,7 +23,7 @@ dependencies = [
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "0.6.0a4"
+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.0a4 → 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.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