dbos 0.6.0a4__py3-none-any.whl → 0.7.0a0__py3-none-any.whl

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/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/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/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.dist-info → dbos-0.7.0a0.dist-info}/METADATA

@@ -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.dist-info → dbos-0.7.0a0.dist-info}/RECORD

@@ -1,20 +1,22 @@
-dbos-0.6.0a4.dist-info/METADATA,sha256=UT1U_MYOXQ6YWzSgkSQHNYZGq8pRqMM2ago9AkTarhA,5000
-dbos-0.6.0a4.dist-info/WHEEL,sha256=rSwsxJWe3vzyR5HCwjWXQruDgschpei4h_giTm0dJVE,90
-dbos-0.6.0a4.dist-info/entry_points.txt,sha256=3PmOPbM4FYxEmggRRdJw0oAsiBzKR8U0yx7bmwUmMOM,39
-dbos-0.6.0a4.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
-dbos/__init__.py,sha256=s3nYPf5yGR4XxiXS_JQbL8WVTjnnwr5EpFoCaCMrFxg,582
+dbos-0.7.0a0.dist-info/METADATA,sha256=x4O-g52DZCDhnpYf5SKAfG872Ee83HodkUzffOAhKhA,5000
+dbos-0.7.0a0.dist-info/WHEEL,sha256=rSwsxJWe3vzyR5HCwjWXQruDgschpei4h_giTm0dJVE,90
+dbos-0.7.0a0.dist-info/entry_points.txt,sha256=3PmOPbM4FYxEmggRRdJw0oAsiBzKR8U0yx7bmwUmMOM,39
+dbos-0.7.0a0.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
+dbos/__init__.py,sha256=heuB3bqRXlVdfea9sKHBIVKSqFP6UuwhQecQfV4hyas,642
 dbos/admin_sever.py,sha256=Qg5T3YRrbPW05PR_99yAaxgo1ugQrAp_uTeTqSfjm_k,3397
 dbos/application_database.py,sha256=1K3kE96BgGi_QWOd2heXluyNTwFAwlUVuAR6JKKUqf0,5659
 dbos/cli.py,sha256=YARlQiWHUwFni-fEOr0k5P_-pqPS4xkywj_B0oTMXn0,8318
 dbos/context.py,sha256=NVMGyvAa2RIiBVspvDz-8MBk_BQyGyYdPdorgO-GSng,16407
 dbos/core.py,sha256=jeqO8DABPAUrFlJXOfRfFDSnA8BGwiPnMa1JNbGuYLs,28584
 dbos/dbos-config.schema.json,sha256=azpfmoDZg7WfSy3kvIsk9iEiKB_-VZt03VEOoXJAkqE,5331
-dbos/dbos.py,sha256=Tf18wfvTrVx3Xo3ur05kLcNXaF75coJNiubwcCWAFgY,29076
+dbos/dbos.py,sha256=wzL51K7bY3J8grHXi0k0F0N09vFVBYWKwI0DGxyN4MY,29730
 dbos/dbos_config.py,sha256=ih_TD_1zTKhPKxk8TPdEIp3ihu82R06SGKg-s4rHxws,5344
 dbos/decorators.py,sha256=lbPefsLK6Cya4cb7TrOcLglOpGT3pc6qjZdsQKlfZLg,629
 dbos/error.py,sha256=DDhB0VHmoZE_CP51ICdFMZSL2gmVS3Dm0aPNWncci94,3876
 dbos/fastapi.py,sha256=s7LnwwYVpJm_QZZwBW5um8NV2Q2Qx85uVZqGcKlSZAo,1881
 dbos/flask.py,sha256=azr4geMEGuuTBCyxIZmgDmmP-6s_pTIF-lGyp9Q4IB8,2430
+dbos/kafka.py,sha256=FtngQHBu2TKfyDF7GFsKJAawFQJiOFxgKEUlNNxrdrw,3055
+dbos/kafka_message.py,sha256=NYvOXNG3Qn7bghn1pv3fg4Pbs86ILZGcK4IB-MLUNu0,409
 dbos/logger.py,sha256=D-aFSZUCHBP34J1IZ5YNkTrJW-rDiH3py_v9jLU4Yrk,3565
 dbos/migrations/env.py,sha256=38SIGVbmn_VV2x2u1aHLcPOoWgZ84eCymf3g_NljmbU,1626
 dbos/migrations/script.py.mako,sha256=MEqL-2qATlST9TAOeYgscMn1uy6HUS9NFvDgl93dMj8,635
@@ -44,4 +46,4 @@ dbos/templates/hello/start_postgres_docker.py,sha256=lQVLlYO5YkhGPEgPqwGc7Y8uDKs
 dbos/tracer.py,sha256=GaXDhdKKF_IQp5SAMipGXiDVwteRKjNbrXyYCH1mor0,2520
 dbos/utils.py,sha256=hWj9iWDrby2cVEhb0pG-IdnrxLqP64NhkaWUXiLc8bA,402
 version/__init__.py,sha256=L4sNxecRuqdtSFdpUGX3TtBi9KL3k7YsZVIvv-fv9-A,1678
-dbos-0.6.0a4.dist-info/RECORD,,
+dbos-0.7.0a0.dist-info/RECORD,,