dbos 0.6.0a3__py3-none-any.whl → 0.6.1__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/cli.py

@@ -12,7 +12,7 @@ from typing import Any
 import tomlkit
 import typer
 from rich import print
-from rich.prompt import Confirm, Prompt
+from rich.prompt import Prompt
 from typing_extensions import Annotated
 
 from dbos import load_config
@@ -30,6 +30,7 @@ def on_windows() -> bool:
 def start() -> None:
     config = load_config()
     start_commands = config["runtimeConfig"]["start"]
+    typer.echo("Executing start commands from 'dbos-config.yaml'")
     for command in start_commands:
         typer.echo(f"Executing: {command}")
 
@@ -129,9 +130,12 @@ def copy_template(src_dir: str, project_name: str, config_mode: bool) -> None:
         "project_name": project_name,
         "package_name": package_name,
         "db_name": db_name,
+        "migration_command": "alembic upgrade head",
     }
 
     if config_mode:
+        ctx["package_name"] = "."
+        ctx["migration_command"] = "echo 'No migrations specified'"
         copy_dbos_template(
             os.path.join(src_dir, "dbos-config.yaml.dbos"),
             os.path.join(dst_dir, "dbos-config.yaml"),
@@ -244,6 +248,7 @@ def migrate() -> None:
             app_db.destroy()
 
     # Next, run any custom migration commands specified in the configuration
+    typer.echo("Executing migration commands from 'dbos-config.yaml'")
     try:
         migrate_commands = (
             config["database"]["migrate"]

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/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/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/templates/hello/dbos-config.yaml.dbos

@@ -15,7 +15,7 @@ database:
   password: ${PGPASSWORD}
   app_db_name: ${db_name}
   migrate:
-    - alembic upgrade head
+    - ${migration_command}
 telemetry:
   logs:
     logLevel: INFO

{dbos-0.6.0a3.dist-info → dbos-0.6.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dbos
-Version: 0.6.0a3
+Version: 0.6.1
 Summary: Ultra-lightweight durable execution in Python
 Author-Email: "DBOS, Inc." <contact@dbos.dev>
 License: MIT
@@ -91,7 +91,7 @@ def step_two():
     print("Step two completed!")
 
 @DBOS.workflow()
-def workflow():
+def dbos_workflow():
     step_one()
     for _ in range(5):
         print("Press Control + \ to stop the app...")
@@ -99,8 +99,8 @@ def workflow():
     step_two()
 
 @app.get("/")
-def endpoint():
-    workflow()
+def fastapi_endpoint():
+    dbos_workflow()
 ```
 
 Save the program into `main.py`, edit `dbos-config.yaml` to configure your Postgres connection settings, and start it with `fastapi run`.

{dbos-0.6.0a3.dist-info → dbos-0.6.1.dist-info}/RECORD

@@ -1,20 +1,22 @@
-dbos-0.6.0a3.dist-info/METADATA,sha256=6YdSITRGDjunFt92A4EPg9RDR__piiogGZ8NsEXYrKM,5000
-dbos-0.6.0a3.dist-info/WHEEL,sha256=rSwsxJWe3vzyR5HCwjWXQruDgschpei4h_giTm0dJVE,90
-dbos-0.6.0a3.dist-info/entry_points.txt,sha256=3PmOPbM4FYxEmggRRdJw0oAsiBzKR8U0yx7bmwUmMOM,39
-dbos-0.6.0a3.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
-dbos/__init__.py,sha256=s3nYPf5yGR4XxiXS_JQbL8WVTjnnwr5EpFoCaCMrFxg,582
+dbos-0.6.1.dist-info/METADATA,sha256=3WgP4NAQT1ZNQbEkLV9hupdBqevptUrwS2k6WpJ6dss,5016
+dbos-0.6.1.dist-info/WHEEL,sha256=rSwsxJWe3vzyR5HCwjWXQruDgschpei4h_giTm0dJVE,90
+dbos-0.6.1.dist-info/entry_points.txt,sha256=3PmOPbM4FYxEmggRRdJw0oAsiBzKR8U0yx7bmwUmMOM,39
+dbos-0.6.1.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/cli.py,sha256=z5dXbbnGWzSC3E1rfS8Lp1_OIImzcDKM7jP-iu_Q4aI,8602
 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=zOTolY-lpDwpSO2ORrtIIRjIdi6JsQc9ovQDmvlc47I,28985
-dbos/dbos_config.py,sha256=SJ_UI4pHVKKjZDjge9Abm0HgcdQ6aDBO8QKBYujEPeE,5307
+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
@@ -36,7 +38,7 @@ dbos/templates/hello/__package/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5
 dbos/templates/hello/__package/main.py,sha256=eI0SS9Nwj-fldtiuSzIlIG6dC91GXXwdRsoHxv6S_WI,2719
 dbos/templates/hello/__package/schema.py,sha256=7Z27JGC8yy7Z44cbVXIREYxtUhU4JVkLCp5Q7UahVQ0,260
 dbos/templates/hello/alembic.ini,sha256=VKBn4Gy8mMuCdY7Hip1jmo3wEUJ1VG1aW7EqY0_n-as,3695
-dbos/templates/hello/dbos-config.yaml.dbos,sha256=8wxCf_MIEFNWqMXj0nAHUwg1U3YaKz4xcUN6g51WkDE,603
+dbos/templates/hello/dbos-config.yaml.dbos,sha256=X-Ocywo8_QUr2-Xj2AiiBYL0tcd0E3Yl7FuP6fy3gRM,603
 dbos/templates/hello/migrations/env.py.dbos,sha256=CsiFOea3ZIsahqkfYtioha0ewmlGR78Ng0wOB2t5LQg,2208
 dbos/templates/hello/migrations/script.py.mako,sha256=MEqL-2qATlST9TAOeYgscMn1uy6HUS9NFvDgl93dMj8,635
 dbos/templates/hello/migrations/versions/2024_07_31_180642_init.py,sha256=U5thFWGqNN4QLrNXT7wUUqftIFDNE5eSdqD8JNW1mec,942
@@ -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.0a3.dist-info/RECORD,,
+dbos-0.6.1.dist-info/RECORD,,