dbos 0.5.0a12__tar.gz → 0.6.0a0__tar.gz

{dbos-0.5.0a12 → dbos-0.6.0a0}/PKG-INFO

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

dbos-0.6.0a0/dbos/__init__.py

@@ -0,0 +1,20 @@
+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 .system_database import GetWorkflowsInput, WorkflowStatusString
+
+__all__ = [
+    "ConfigFile",
+    "DBOS",
+    "DBOSConfiguredInstance",
+    "DBOSContextEnsure",
+    "GetWorkflowsInput",
+    "SetWorkflowID",
+    "WorkflowHandle",
+    "WorkflowStatus",
+    "WorkflowStatusString",
+    "load_config",
+    "get_dbos_database_url",
+    "error",
+]

{dbos-0.5.0a12 → dbos-0.6.0a0}/dbos/context.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import json
 import os
 import uuid
 from contextvars import ContextVar
@@ -39,6 +40,9 @@ class TracedAttributes(TypedDict, total=False):
     applicationID: Optional[str]
     applicationVersion: Optional[str]
     executorID: Optional[str]
+    authenticatedUser: Optional[str]
+    authenticatedUserRoles: Optional[str]
+    authenticatedUserAssumedRole: Optional[str]
 
 
 class DBOSContext:
@@ -159,6 +163,13 @@ class DBOSContext:
         attributes["operationUUID"] = (
             self.workflow_id if len(self.workflow_id) > 0 else None
         )
+        attributes["authenticatedUser"] = self.authenticated_user
+        attributes["authenticatedUserRoles"] = (
+            json.dumps(self.authenticated_roles)
+            if self.authenticated_roles is not None
+            else ""
+        )
+        attributes["authenticatedUserAssumedRole"] = self.assumed_role
         span = dbos_tracer.start_span(
             attributes, parent=self.spans[-1] if len(self.spans) > 0 else None
         )
@@ -178,6 +189,11 @@ class DBOSContext:
     ) -> None:
         self.authenticated_user = user
         self.authenticated_roles = roles
+        if user is not None and len(self.spans) > 0:
+            self.spans[-1].set_attribute("authenticatedUser", user)
+            self.spans[-1].set_attribute(
+                "authenticatedUserRoles", json.dumps(roles) if roles is not None else ""
+            )
 
 
 ##############################################################
@@ -217,13 +233,13 @@ class DBOSContextEnsure:
     def __init__(self) -> None:
         self.created_ctx = False
 
-    def __enter__(self) -> DBOSContextEnsure:
+    def __enter__(self) -> DBOSContext:
         # Code to create a basic context
         ctx = get_local_dbos_context()
         if ctx is None:
             self.created_ctx = True
             set_local_dbos_context(DBOSContext())
-        return self
+        return assert_current_dbos_context()
 
     def __exit__(
         self,

{dbos-0.5.0a12 → dbos-0.6.0a0}/dbos/core.py

@@ -1,9 +1,20 @@
+import json
 import sys
 import time
 import traceback
 from concurrent.futures import Future
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, Tuple, TypeVar, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Generic,
+    List,
+    Optional,
+    Tuple,
+    TypeVar,
+    cast,
+)
 
 from dbos.application_database import ApplicationDatabase, TransactionResultInternal
 
@@ -134,6 +145,11 @@ def _init_workflow(
         "executor_id": ctx.executor_id,
         "request": (utils.serialize(ctx.request) if ctx.request is not None else None),
         "recovery_attempts": None,
+        "authenticated_user": ctx.authenticated_user,
+        "authenticated_roles": (
+            json.dumps(ctx.authenticated_roles) if ctx.authenticated_roles else None
+        ),
+        "assumed_role": ctx.assumed_role,
     }
 
     # If we have a class name, the first arg is the instance and do not serialize

{dbos-0.5.0a12 → dbos-0.6.0a0}/dbos/dbos.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import atexit
+import json
 import os
 import sys
 import threading
@@ -19,6 +20,7 @@ from typing import (
     Tuple,
     Type,
     TypeVar,
+    cast,
 )
 
 from opentelemetry.trace import Span
@@ -53,6 +55,7 @@ from .tracer import dbos_tracer
 if TYPE_CHECKING:
     from fastapi import FastAPI
     from .request import Request
+    from flask import Flask
 
 from sqlalchemy.orm import Session
 
@@ -205,6 +208,7 @@ class DBOS:
         *,
         config: Optional[ConfigFile] = None,
         fastapi: Optional["FastAPI"] = None,
+        flask: Optional["Flask"] = None,
     ) -> DBOS:
         global _dbos_global_instance
         global _dbos_global_registry
@@ -219,7 +223,7 @@ class DBOS:
                     )
                 config = _dbos_global_registry.config
             _dbos_global_instance = super().__new__(cls)
-            _dbos_global_instance.__init__(fastapi=fastapi, config=config)  # type: ignore
+            _dbos_global_instance.__init__(fastapi=fastapi, config=config, flask=flask)  # type: ignore
         else:
             if (config is not None and _dbos_global_instance.config is not config) or (
                 _dbos_global_instance.fastapi is not fastapi
@@ -243,6 +247,7 @@ class DBOS:
         *,
         config: Optional[ConfigFile] = None,
         fastapi: Optional["FastAPI"] = None,
+        flask: Optional["Flask"] = None,
     ) -> None:
         if hasattr(self, "_initialized") and self._initialized:
             return
@@ -264,14 +269,44 @@ class DBOS:
         self._admin_server: Optional[AdminServer] = None
         self.stop_events: List[threading.Event] = []
         self.fastapi: Optional["FastAPI"] = fastapi
+        self.flask: Optional["Flask"] = flask
         self._executor: Optional[ThreadPoolExecutor] = None
+
+        # If using FastAPI, set up middleware and lifecycle events
         if self.fastapi is not None:
+            from fastapi.requests import Request as FARequest
+            from fastapi.responses import JSONResponse
+
+            async def dbos_error_handler(
+                request: FARequest, gexc: Exception
+            ) -> JSONResponse:
+                exc: DBOSException = cast(DBOSException, gexc)
+                status_code = 500
+                if exc.status_code is not None:
+                    status_code = exc.status_code
+                return JSONResponse(
+                    status_code=status_code,
+                    content={
+                        "message": str(exc.message),
+                        "dbos_error_code": str(exc.dbos_error_code),
+                        "dbos_error": str(exc.__class__.__name__),
+                    },
+                )
+
+            self.fastapi.add_exception_handler(DBOSException, dbos_error_handler)
+
             from dbos.fastapi import setup_fastapi_middleware
 
             setup_fastapi_middleware(self.fastapi)
             self.fastapi.on_event("startup")(self._launch)
             self.fastapi.on_event("shutdown")(self._destroy)
 
+        # If using Flask, set up middleware
+        if self.flask is not None:
+            from dbos.flask import setup_flask_middleware
+
+            setup_flask_middleware(self.flask)
+
         # Register send_stub as a workflow
         def send_temp_workflow(
             destination_id: str, message: Any, topic: Optional[str]
@@ -502,9 +537,13 @@ class DBOS:
             recovery_attempts=stat["recovery_attempts"],
             class_name=stat["class_name"],
             config_name=stat["config_name"],
-            authenticated_user=None,
-            assumed_role=None,
-            authenticatedRoles=None,
+            authenticated_user=stat["authenticated_user"],
+            assumed_role=stat["assumed_role"],
+            authenticated_roles=(
+                json.loads(stat["authenticated_roles"])
+                if stat["authenticated_roles"] is not None
+                else None
+            ),
         )
 
     @classmethod
@@ -663,6 +702,27 @@ class DBOS:
         ctx = assert_current_dbos_context()
         return ctx.request
 
+    @classproperty
+    def authenticated_user(cls) -> Optional[str]:
+        """Return the current authenticated user, if any, associated with the current context."""
+        ctx = assert_current_dbos_context()
+        return ctx.authenticated_user
+
+    @classproperty
+    def authenticated_roles(cls) -> Optional[List[str]]:
+        """Return the roles granted to the current authenticated user, if any, associated with the current context."""
+        ctx = assert_current_dbos_context()
+        return ctx.authenticated_roles
+
+    @classmethod
+    def set_authentication(
+        cls, authenticated_user: Optional[str], authenticated_roles: Optional[List[str]]
+    ) -> None:
+        """Set the current authenticated user and granted roles into the current context."""
+        ctx = assert_current_dbos_context()
+        ctx.authenticated_user = authenticated_user
+        ctx.authenticated_roles = authenticated_roles
+
 
 @dataclass
 class WorkflowStatus:
@@ -679,7 +739,7 @@ class WorkflowStatus:
         config_name(str): For instance member functions, the name of the class instance for the execution
         authenticated_user(str): The user who invoked the workflow
         assumed_role(str): The access role used by the user to allow access to the workflow function
-        authenticatedRoles(List[str]): List of all access roles available to the authenticated user
+        authenticated_roles(List[str]): List of all access roles available to the authenticated user
         recovery_attempts(int): Number of times the workflow has been restarted (usually by recovery)
 
     """
@@ -691,7 +751,7 @@ class WorkflowStatus:
     config_name: Optional[str]
     authenticated_user: Optional[str]
     assumed_role: Optional[str]
-    authenticatedRoles: Optional[List[str]]
+    authenticated_roles: Optional[List[str]]
     recovery_attempts: Optional[int]
 
 

{dbos-0.5.0a12 → dbos-0.6.0a0}/dbos/dbos_config.py

@@ -2,7 +2,7 @@ import json
 import os
 import re
 from importlib import resources
-from typing import Any, Dict, List, Optional, TypedDict
+from typing import Dict, List, Optional, TypedDict
 
 import yaml
 from jsonschema import ValidationError, validate

{dbos-0.5.0a12 → dbos-0.6.0a0}/dbos/error.py

@@ -17,6 +17,7 @@ class DBOSException(Exception):
     def __init__(self, message: str, dbos_error_code: Optional[int] = None):
         self.message = message
         self.dbos_error_code = dbos_error_code
+        self.status_code: Optional[int] = None
         super().__init__(self.message)
 
     def __str__(self) -> str:
@@ -104,6 +105,7 @@ class DBOSNotAuthorizedError(DBOSException):
             msg,
             dbos_error_code=DBOSErrorCode.NotAuthorized.value,
         )
+        self.status_code = 403
 
 
 class DBOSMaxStepRetriesExceeded(DBOSException):

{dbos-0.5.0a12 → dbos-0.6.0a0}/dbos/fastapi.py

@@ -11,9 +11,7 @@ from .context import (
     TracedAttributes,
     assert_current_dbos_context,
 )
-from .request import Address, Request
-
-request_id_header = "x-request-id"
+from .request import Address, Request, request_id_header
 
 
 def get_or_generate_request_id(request: FastAPIRequest) -> str:

dbos-0.6.0a0/dbos/flask.py

@@ -0,0 +1,79 @@
+import uuid
+from typing import Any
+from urllib.parse import urlparse
+
+from flask import Flask, request
+from werkzeug.wrappers import Request as WRequest
+
+from dbos.context import (
+    EnterDBOSHandler,
+    OperationType,
+    SetWorkflowID,
+    TracedAttributes,
+    assert_current_dbos_context,
+)
+
+from .request import Address, Request, request_id_header
+
+
+class FlaskMiddleware:
+    def __init__(self, app: Any) -> None:
+        self.app = app
+
+    def __call__(self, environ: Any, start_response: Any) -> Any:
+        request = WRequest(environ)
+        attributes: TracedAttributes = {
+            "name": urlparse(request.url).path,
+            "requestID": get_or_generate_request_id(request),
+            "requestIP": (
+                request.remote_addr if request.remote_addr is not None else None
+            ),
+            "requestURL": request.url,
+            "requestMethod": request.method,
+            "operationType": OperationType.HANDLER.value,
+        }
+        with EnterDBOSHandler(attributes):
+            ctx = assert_current_dbos_context()
+            ctx.request = make_request(request)
+            workflow_id = request.headers.get("dbos-idempotency-key", "")
+            with SetWorkflowID(workflow_id):
+                response = self.app(environ, start_response)
+        return response
+
+
+def get_or_generate_request_id(request: WRequest) -> str:
+    request_id = request.headers.get(request_id_header, None)
+    if request_id is not None:
+        return request_id
+    else:
+        return str(uuid.uuid4())
+
+
+def make_request(request: WRequest) -> Request:
+    parsed_url = urlparse(request.url)
+    base_url = f"{parsed_url.scheme}://{parsed_url.netloc}"
+
+    client = None
+    if request.remote_addr:
+        hostname = request.remote_addr
+        port = request.environ.get("REMOTE_PORT")
+        if port:
+            client = Address(hostname=hostname, port=int(port))
+        else:
+            # If port is not available, use 0 as a placeholder
+            client = Address(hostname=hostname, port=0)
+
+    return Request(
+        headers=dict(request.headers),
+        path_params={},
+        query_params=dict(request.args),
+        url=request.url,
+        base_url=base_url,
+        client=client,
+        cookies=dict(request.cookies),
+        method=request.method,
+    )
+
+
+def setup_flask_middleware(app: Flask) -> None:
+    app.wsgi_app = FlaskMiddleware(app.wsgi_app)  # type: ignore

{dbos-0.5.0a12 → dbos-0.6.0a0}/dbos/request.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass
 from typing import Any, Mapping, NamedTuple, Optional
 
+request_id_header = "x-request-id"
+
 
 class Address(NamedTuple):
     hostname: str

{dbos-0.5.0a12 → dbos-0.6.0a0}/dbos/roles.py

@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Any, Callable, List, Optional, Type, TypeVar,
 from dbos.error import DBOSNotAuthorizedError
 
 if TYPE_CHECKING:
-    from dbos.dbos import DBOS, _DBOSRegistry
+    from dbos.dbos import _DBOSRegistry
 
 from dbos.context import DBOSAssumeRole, get_local_dbos_context
 from dbos.registrations import (

{dbos-0.5.0a12 → dbos-0.6.0a0}/dbos/system_database.py

@@ -4,18 +4,7 @@ import select
 import threading
 import time
 from enum import Enum
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Dict,
-    List,
-    Literal,
-    Optional,
-    Sequence,
-    Set,
-    TypedDict,
-    cast,
-)
+from typing import Any, Dict, List, Literal, Optional, Sequence, Set, TypedDict, cast
 
 import psycopg2
 import sqlalchemy as sa
@@ -69,6 +58,9 @@ class WorkflowStatusInternal(TypedDict):
     app_id: Optional[str]
     request: Optional[str]  # JSON (jsonpickle)
     recovery_attempts: Optional[int]
+    authenticated_user: Optional[str]
+    assumed_role: Optional[str]
+    authenticated_roles: Optional[str]  # JSON list of roles.
 
 
 class RecordedResult(TypedDict):
@@ -136,7 +128,7 @@ class WorkflowInformation(TypedDict, total=False):
     authenticated_user: str  # The user who ran the workflow. Empty string if not set.
     assumed_role: str
     # The role used to run this workflow.  Empty string if authorization is not required.
-    authenticatedRoles: List[str]
+    authenticated_roles: List[str]
     # All roles the authenticated user has, if any.
     input: Optional[WorkflowInputs]
     output: Optional[str]
@@ -252,6 +244,9 @@ class SystemDatabase:
             application_version=status["app_version"],
             application_id=status["app_id"],
             request=status["request"],
+            authenticated_user=status["authenticated_user"],
+            authenticated_roles=status["authenticated_roles"],
+            assumed_role=status["assumed_role"],
         )
         if replace:
             cmd = cmd.on_conflict_do_update(
@@ -322,6 +317,9 @@ class SystemDatabase:
                     SystemSchema.workflow_status.c.recovery_attempts,
                     SystemSchema.workflow_status.c.config_name,
                     SystemSchema.workflow_status.c.class_name,
+                    SystemSchema.workflow_status.c.authenticated_user,
+                    SystemSchema.workflow_status.c.authenticated_roles,
+                    SystemSchema.workflow_status.c.assumed_role,
                 ).where(SystemSchema.workflow_status.c.workflow_uuid == workflow_uuid)
             ).fetchone()
             if row is None:
@@ -339,6 +337,9 @@ class SystemDatabase:
                 "executor_id": None,
                 "request": row[2],
                 "recovery_attempts": row[3],
+                "authenticated_user": row[6],
+                "authenticated_roles": row[7],
+                "assumed_role": row[8],
             }
             return status
 
@@ -375,6 +376,9 @@ class SystemDatabase:
                     SystemSchema.workflow_status.c.error,
                     SystemSchema.workflow_status.c.config_name,
                     SystemSchema.workflow_status.c.class_name,
+                    SystemSchema.workflow_status.c.authenticated_user,
+                    SystemSchema.workflow_status.c.authenticated_roles,
+                    SystemSchema.workflow_status.c.assumed_role,
                 ).where(SystemSchema.workflow_status.c.workflow_uuid == workflow_uuid)
             ).fetchone()
             if row is None:
@@ -392,6 +396,9 @@ class SystemDatabase:
                 "executor_id": None,
                 "request": row[2],
                 "recovery_attempts": None,
+                "authenticated_user": row[7],
+                "authenticated_roles": row[8],
+                "assumed_role": row[9],
             }
             return status
 

{dbos-0.5.0a12 → dbos-0.6.0a0}/dbos/tracer.py

@@ -19,6 +19,7 @@ class DBOSTracer:
         self.app_id = os.environ.get("DBOS__APPID", None)
         self.app_version = os.environ.get("DBOS__APPVERSION", None)
         self.executor_id = os.environ.get("DBOS__VMID", "local")
+        self.provider: Optional[TracerProvider] = None
 
     def config(self, config: ConfigFile) -> None:
         if not isinstance(trace.get_tracer_provider(), TracerProvider):
@@ -36,10 +37,17 @@ class DBOSTracer:
                 provider.add_span_processor(processor)
             trace.set_tracer_provider(provider)
 
+    def set_provider(self, provider: Optional[TracerProvider]) -> None:
+        self.provider = provider
+
     def start_span(
         self, attributes: "TracedAttributes", parent: Optional[Span] = None
     ) -> Span:
-        tracer = trace.get_tracer("dbos-tracer")
+        tracer = (
+            self.provider.get_tracer("dbos-tracer")
+            if self.provider is not None
+            else trace.get_tracer("dbos-tracer")
+        )
         context = trace.set_span_in_context(parent) if parent else None
         span: Span = tracer.start_span(name=attributes["name"], context=context)
         attributes["applicationID"] = self.app_id

{dbos-0.5.0a12 → dbos-0.6.0a0}/pyproject.toml

@@ -23,7 +23,7 @@ dependencies = [
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "0.5.0a12"
+version = "0.6.0a0"
 
 [project.license]
 text = "MIT"
@@ -63,6 +63,9 @@ dev = [
     "pdm>=2.18.0",
     "pdm-backend>=2.3.3",
     "GitPython>=3.1.43",
+    "flask>=3.0.3",
+    "pytest-order>=1.3.0",
+    "pyjwt>=2.9.0",
 ]
 
 [tool.black]

{dbos-0.5.0a12 → dbos-0.6.0a0}/tests/conftest.py

@@ -7,6 +7,7 @@ from typing import Any, Generator, Tuple
 import pytest
 import sqlalchemy as sa
 from fastapi import FastAPI
+from flask import Flask
 
 from dbos import DBOS, ConfigFile
 
@@ -135,6 +136,23 @@ def dbos_fastapi(
     DBOS.destroy()
 
 
+@pytest.fixture()
+def dbos_flask(
+    config: ConfigFile, cleanup_test_databases: None
+) -> Generator[Tuple[DBOS, Flask], Any, None]:
+    DBOS.destroy()
+    app = Flask(__name__)
+
+    dbos = DBOS(flask=app, config=config)
+
+    # This is for test convenience.
+    #    Usually fastapi itself does launch, but we are not completing the fastapi lifecycle
+    DBOS.launch()
+
+    yield dbos, app
+    DBOS.destroy()
+
+
 # Pretty-print test names
 def pytest_collection_modifyitems(session: Any, config: Any, items: Any) -> None:
     for item in items:

{dbos-0.5.0a12 → dbos-0.6.0a0}/tests/scheduler/test_scheduler.py

@@ -84,6 +84,9 @@ def test_scheduler_oaoo(dbos: DBOS) -> None:
             "app_version": None,
             "request": None,
             "recovery_attempts": None,
+            "authenticated_user": None,
+            "authenticated_roles": None,
+            "assumed_role": None,
         }
     )
 

{dbos-0.5.0a12 → dbos-0.6.0a0}/tests/test_admin_server.py

@@ -92,6 +92,9 @@ def test_admin_recovery(dbos: DBOS) -> None:
             "app_version": None,
             "request": None,
             "recovery_attempts": None,
+            "authenticated_user": None,
+            "authenticated_roles": None,
+            "assumed_role": None,
         }
     )
     status = dbos.get_workflow_status(wfuuid)