dbos 0.5.0a11__tar.gz → 0.6.0a0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of dbos might be problematic. Click here for more details.

Files changed (69) hide show
  1. dbos-0.6.0a0/PKG-INFO +130 -0
  2. dbos-0.6.0a0/README.md +107 -0
  3. dbos-0.6.0a0/dbos/__init__.py +20 -0
  4. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/admin_sever.py +2 -2
  5. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/cli.py +19 -6
  6. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/context.py +18 -2
  7. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/core.py +17 -1
  8. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/dbos.py +66 -6
  9. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/dbos_config.py +1 -1
  10. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/error.py +2 -0
  11. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/fastapi.py +1 -3
  12. dbos-0.6.0a0/dbos/flask.py +79 -0
  13. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/request.py +2 -0
  14. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/roles.py +1 -1
  15. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/system_database.py +21 -14
  16. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/tracer.py +9 -1
  17. {dbos-0.5.0a11 → dbos-0.6.0a0}/pyproject.toml +5 -2
  18. {dbos-0.5.0a11 → dbos-0.6.0a0}/tests/conftest.py +18 -0
  19. {dbos-0.5.0a11 → dbos-0.6.0a0}/tests/scheduler/test_scheduler.py +3 -0
  20. {dbos-0.5.0a11 → dbos-0.6.0a0}/tests/test_admin_server.py +3 -0
  21. {dbos-0.5.0a11 → dbos-0.6.0a0}/tests/test_dbos.py +9 -51
  22. {dbos-0.5.0a11 → dbos-0.6.0a0}/tests/test_fastapi.py +3 -0
  23. dbos-0.6.0a0/tests/test_fastapi_roles.py +342 -0
  24. dbos-0.6.0a0/tests/test_flask.py +104 -0
  25. dbos-0.5.0a11/PKG-INFO +0 -78
  26. dbos-0.5.0a11/README.md +0 -55
  27. dbos-0.5.0a11/dbos/__init__.py +0 -24
  28. {dbos-0.5.0a11 → dbos-0.6.0a0}/LICENSE +0 -0
  29. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/application_database.py +0 -0
  30. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/dbos-config.schema.json +0 -0
  31. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/decorators.py +0 -0
  32. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/logger.py +0 -0
  33. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/migrations/env.py +0 -0
  34. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/migrations/script.py.mako +0 -0
  35. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/migrations/versions/5c361fc04708_added_system_tables.py +0 -0
  36. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/migrations/versions/a3b18ad34abe_added_triggers.py +0 -0
  37. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/py.typed +0 -0
  38. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/recovery.py +0 -0
  39. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/registrations.py +0 -0
  40. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/scheduler/croniter.py +0 -0
  41. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/scheduler/scheduler.py +0 -0
  42. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/schemas/__init__.py +0 -0
  43. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/schemas/application_database.py +0 -0
  44. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/schemas/system_database.py +0 -0
  45. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/templates/hello/README.md +0 -0
  46. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/templates/hello/__package/__init__.py +0 -0
  47. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/templates/hello/__package/main.py +0 -0
  48. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/templates/hello/__package/schema.py +0 -0
  49. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/templates/hello/alembic.ini +0 -0
  50. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/templates/hello/dbos-config.yaml.dbos +0 -0
  51. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/templates/hello/migrations/env.py.dbos +0 -0
  52. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/templates/hello/migrations/script.py.mako +0 -0
  53. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/templates/hello/migrations/versions/2024_07_31_180642_init.py +0 -0
  54. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/templates/hello/start_postgres_docker.py +0 -0
  55. {dbos-0.5.0a11 → dbos-0.6.0a0}/dbos/utils.py +0 -0
  56. {dbos-0.5.0a11 → dbos-0.6.0a0}/tests/__init__.py +0 -0
  57. {dbos-0.5.0a11 → dbos-0.6.0a0}/tests/atexit_no_ctor.py +0 -0
  58. {dbos-0.5.0a11 → dbos-0.6.0a0}/tests/atexit_no_launch.py +0 -0
  59. {dbos-0.5.0a11 → dbos-0.6.0a0}/tests/classdefs.py +0 -0
  60. {dbos-0.5.0a11 → dbos-0.6.0a0}/tests/more_classdefs.py +0 -0
  61. {dbos-0.5.0a11 → dbos-0.6.0a0}/tests/scheduler/test_croniter.py +0 -0
  62. {dbos-0.5.0a11 → dbos-0.6.0a0}/tests/test_classdecorators.py +0 -0
  63. {dbos-0.5.0a11 → dbos-0.6.0a0}/tests/test_concurrency.py +0 -0
  64. {dbos-0.5.0a11 → dbos-0.6.0a0}/tests/test_config.py +0 -0
  65. {dbos-0.5.0a11 → dbos-0.6.0a0}/tests/test_failures.py +0 -0
  66. {dbos-0.5.0a11 → dbos-0.6.0a0}/tests/test_package.py +0 -0
  67. {dbos-0.5.0a11 → dbos-0.6.0a0}/tests/test_schema_migration.py +0 -0
  68. {dbos-0.5.0a11 → dbos-0.6.0a0}/tests/test_singleton.py +0 -0
  69. {dbos-0.5.0a11 → dbos-0.6.0a0}/version/__init__.py +0 -0
dbos-0.6.0a0/PKG-INFO ADDED
@@ -0,0 +1,130 @@
1
+ Metadata-Version: 2.1
2
+ Name: dbos
3
+ Version: 0.6.0a0
4
+ Summary: Ultra-lightweight durable execution in Python
5
+ Author-Email: "DBOS, Inc." <contact@dbos.dev>
6
+ License: MIT
7
+ Requires-Python: >=3.9
8
+ Requires-Dist: pyyaml>=6.0.2
9
+ Requires-Dist: jsonschema>=4.23.0
10
+ Requires-Dist: alembic>=1.13.2
11
+ Requires-Dist: psycopg2-binary>=2.9.9
12
+ Requires-Dist: typing-extensions>=4.12.2; python_version < "3.10"
13
+ Requires-Dist: typer>=0.12.3
14
+ Requires-Dist: jsonpickle>=3.2.2
15
+ Requires-Dist: opentelemetry-api>=1.26.0
16
+ Requires-Dist: opentelemetry-sdk>=1.26.0
17
+ Requires-Dist: opentelemetry-exporter-otlp-proto-http>=1.26.0
18
+ Requires-Dist: python-dateutil>=2.9.0.post0
19
+ Requires-Dist: fastapi[standard]>=0.112.1
20
+ Requires-Dist: psutil>=6.0.0
21
+ Requires-Dist: tomlkit>=0.13.2
22
+ Description-Content-Type: text/markdown
23
+
24
+ ## 🚀 DBOS Transact - Ultra-Lightweight Durable Execution in Python 🚀
25
+
26
+ ---
27
+
28
+ 📚 **Documentation**: Under Construction 🚧
29
+
30
+ 💬 **Join the Discussion**: [Discord Community](https://discord.gg/fMwQjeW5zg)
31
+
32
+ ---
33
+
34
+
35
+ **DBOS Python is under construction! 🚧🚧🚧 Check back regularly for updates, release coming in mid-September!**
36
+
37
+ DBOS Transact is a **Python library** providing ultra-lightweight durable execution.
38
+ For example:
39
+
40
+ ```python
41
+ @DBOS.step()
42
+ def step_one():
43
+ ...
44
+
45
+ @DBOS.step()
46
+ def step_two():
47
+ ...
48
+
49
+ @DBOS.workflow()
50
+ def workflow()
51
+ step_one()
52
+ step_two()
53
+ ```
54
+
55
+ Durable execution means your program is **resilient to any failure**.
56
+ If it is ever interrupted or crashes, all your workflows will automatically resume from the last completed step.
57
+ If you want to see durable execution in action, check out [this demo app](https://demo-widget-store.cloud.dbos.dev/) (source code [here](https://github.com/dbos-inc/dbos-demo-apps/tree/main/python/widget-store)).
58
+ No matter how many times you try to crash it, it always resumes from exactly where it left off!
59
+
60
+ Under the hood, DBOS Transact works by storing your program's execution state (which workflows are currently executing and which steps they've completed) in a Postgres database.
61
+ So all you need to use it is a Postgres database to connect to&mdash;there's no need for a "workflow server."
62
+ This approach is also incredibly fast, for example [25x faster than AWS Step Functions](https://www.dbos.dev/blog/dbos-vs-aws-step-functions-benchmark).
63
+
64
+ Some more cool features include:
65
+
66
+ - Scheduled jobs&mdash;run your workflows exactly-once per time interval.
67
+ - Exactly-once event processing&mdash;use workflows to process incoming events (for example, from a Kafka topic) exactly-once.
68
+ - Observability&mdash;all workflows automatically emit [OpenTelemetry](https://opentelemetry.io/) traces.
69
+
70
+ ## Getting Started
71
+
72
+ To try out the latest pre-release version, install and configure with:
73
+
74
+ ```shell
75
+ pip install --pre dbos
76
+ dbos init --config
77
+ ```
78
+
79
+ Try it out with this simple program (requires Postgres):
80
+
81
+ ```python
82
+ from fastapi import FastAPI
83
+ from dbos import DBOS
84
+
85
+ app = FastAPI()
86
+ DBOS(fastapi=app)
87
+
88
+ @DBOS.step()
89
+ def step_one():
90
+ print("Step one completed!")
91
+
92
+ @DBOS.step()
93
+ def step_two():
94
+ print("Step two completed!")
95
+
96
+ @DBOS.workflow()
97
+ def workflow():
98
+ step_one()
99
+ for _ in range(5):
100
+ print("Press Control + \ to stop the app...")
101
+ DBOS.sleep(1)
102
+ step_two()
103
+
104
+ @app.get("/")
105
+ def endpoint():
106
+ workflow()
107
+ ```
108
+
109
+ Save the program into `main.py`, tell it your local Postgres password via `export PGPASSWORD=<your password>` and start it with `fastapi run`.
110
+ Visit `localhost:8000` in your browser (or curl it) to start the workflow.
111
+ When prompted, press `Control + \` to force quit your application.
112
+ It should crash midway through the workflow, having completed step one but not step two.
113
+ Then, restart your app with `fastapi run`.
114
+ It should resume the workflow from where it left off, completing step two without re-executing step one.
115
+
116
+ To learn how to build more complex examples, see our programming guide (coming soon).
117
+
118
+ ## Documentation
119
+
120
+ Coming soon! 🚧
121
+
122
+ ## Examples
123
+
124
+ Check out some cool demo apps here: [https://github.com/dbos-inc/dbos-demo-apps/tree/main/python](https://github.com/dbos-inc/dbos-demo-apps/tree/main/python)
125
+
126
+ ## Community
127
+
128
+ If you're interested in building with us, please star our repository and join our community on [Discord](https://discord.gg/fMwQjeW5zg)!
129
+ If you see a bug or have a feature request, don't hesitate to open an issue here on GitHub.
130
+ If you're interested in contributing, check out our [contributions guide](./CONTRIBUTING.md).
dbos-0.6.0a0/README.md ADDED
@@ -0,0 +1,107 @@
1
+ ## 🚀 DBOS Transact - Ultra-Lightweight Durable Execution in Python 🚀
2
+
3
+ ---
4
+
5
+ 📚 **Documentation**: Under Construction 🚧
6
+
7
+ 💬 **Join the Discussion**: [Discord Community](https://discord.gg/fMwQjeW5zg)
8
+
9
+ ---
10
+
11
+
12
+ **DBOS Python is under construction! 🚧🚧🚧 Check back regularly for updates, release coming in mid-September!**
13
+
14
+ DBOS Transact is a **Python library** providing ultra-lightweight durable execution.
15
+ For example:
16
+
17
+ ```python
18
+ @DBOS.step()
19
+ def step_one():
20
+ ...
21
+
22
+ @DBOS.step()
23
+ def step_two():
24
+ ...
25
+
26
+ @DBOS.workflow()
27
+ def workflow()
28
+ step_one()
29
+ step_two()
30
+ ```
31
+
32
+ Durable execution means your program is **resilient to any failure**.
33
+ If it is ever interrupted or crashes, all your workflows will automatically resume from the last completed step.
34
+ If you want to see durable execution in action, check out [this demo app](https://demo-widget-store.cloud.dbos.dev/) (source code [here](https://github.com/dbos-inc/dbos-demo-apps/tree/main/python/widget-store)).
35
+ No matter how many times you try to crash it, it always resumes from exactly where it left off!
36
+
37
+ Under the hood, DBOS Transact works by storing your program's execution state (which workflows are currently executing and which steps they've completed) in a Postgres database.
38
+ So all you need to use it is a Postgres database to connect to&mdash;there's no need for a "workflow server."
39
+ This approach is also incredibly fast, for example [25x faster than AWS Step Functions](https://www.dbos.dev/blog/dbos-vs-aws-step-functions-benchmark).
40
+
41
+ Some more cool features include:
42
+
43
+ - Scheduled jobs&mdash;run your workflows exactly-once per time interval.
44
+ - Exactly-once event processing&mdash;use workflows to process incoming events (for example, from a Kafka topic) exactly-once.
45
+ - Observability&mdash;all workflows automatically emit [OpenTelemetry](https://opentelemetry.io/) traces.
46
+
47
+ ## Getting Started
48
+
49
+ To try out the latest pre-release version, install and configure with:
50
+
51
+ ```shell
52
+ pip install --pre dbos
53
+ dbos init --config
54
+ ```
55
+
56
+ Try it out with this simple program (requires Postgres):
57
+
58
+ ```python
59
+ from fastapi import FastAPI
60
+ from dbos import DBOS
61
+
62
+ app = FastAPI()
63
+ DBOS(fastapi=app)
64
+
65
+ @DBOS.step()
66
+ def step_one():
67
+ print("Step one completed!")
68
+
69
+ @DBOS.step()
70
+ def step_two():
71
+ print("Step two completed!")
72
+
73
+ @DBOS.workflow()
74
+ def workflow():
75
+ step_one()
76
+ for _ in range(5):
77
+ print("Press Control + \ to stop the app...")
78
+ DBOS.sleep(1)
79
+ step_two()
80
+
81
+ @app.get("/")
82
+ def endpoint():
83
+ workflow()
84
+ ```
85
+
86
+ Save the program into `main.py`, tell it your local Postgres password via `export PGPASSWORD=<your password>` and start it with `fastapi run`.
87
+ Visit `localhost:8000` in your browser (or curl it) to start the workflow.
88
+ When prompted, press `Control + \` to force quit your application.
89
+ It should crash midway through the workflow, having completed step one but not step two.
90
+ Then, restart your app with `fastapi run`.
91
+ It should resume the workflow from where it left off, completing step two without re-executing step one.
92
+
93
+ To learn how to build more complex examples, see our programming guide (coming soon).
94
+
95
+ ## Documentation
96
+
97
+ Coming soon! 🚧
98
+
99
+ ## Examples
100
+
101
+ Check out some cool demo apps here: [https://github.com/dbos-inc/dbos-demo-apps/tree/main/python](https://github.com/dbos-inc/dbos-demo-apps/tree/main/python)
102
+
103
+ ## Community
104
+
105
+ If you're interested in building with us, please star our repository and join our community on [Discord](https://discord.gg/fMwQjeW5zg)!
106
+ If you see a bug or have a feature request, don't hesitate to open an issue here on GitHub.
107
+ If you're interested in contributing, check out our [contributions guide](./CONTRIBUTING.md).
@@ -0,0 +1,20 @@
1
+ from . import error as error
2
+ from .context import DBOSContextEnsure, SetWorkflowID
3
+ from .dbos import DBOS, DBOSConfiguredInstance, WorkflowHandle, WorkflowStatus
4
+ from .dbos_config import ConfigFile, get_dbos_database_url, load_config
5
+ from .system_database import GetWorkflowsInput, WorkflowStatusString
6
+
7
+ __all__ = [
8
+ "ConfigFile",
9
+ "DBOS",
10
+ "DBOSConfiguredInstance",
11
+ "DBOSContextEnsure",
12
+ "GetWorkflowsInput",
13
+ "SetWorkflowID",
14
+ "WorkflowHandle",
15
+ "WorkflowStatus",
16
+ "WorkflowStatusString",
17
+ "load_config",
18
+ "get_dbos_database_url",
19
+ "error",
20
+ ]
@@ -28,11 +28,11 @@ class AdminServer:
28
28
  self.server_thread = threading.Thread(target=self.server.serve_forever)
29
29
  self.server_thread.daemon = True
30
30
 
31
- dbos_logger.info("Starting DBOS admin server on port %d", self.port)
31
+ dbos_logger.debug("Starting DBOS admin server on port %d", self.port)
32
32
  self.server_thread.start()
33
33
 
34
34
  def stop(self) -> None:
35
- dbos_logger.info("Stopping DBOS admin server")
35
+ dbos_logger.debug("Stopping DBOS admin server")
36
36
  self.server.shutdown()
37
37
  self.server.server_close()
38
38
  self.server_thread.join()
@@ -119,7 +119,7 @@ def copy_template_dir(src_dir: str, dst_dir: str, ctx: dict[str, str]) -> None:
119
119
  shutil.copy(src, dst)
120
120
 
121
121
 
122
- def copy_template(src_dir: str, project_name: str) -> None:
122
+ def copy_template(src_dir: str, project_name: str, config_mode: bool) -> None:
123
123
 
124
124
  dst_dir = path.abspath(".")
125
125
 
@@ -131,10 +131,17 @@ def copy_template(src_dir: str, project_name: str) -> None:
131
131
  "db_name": db_name,
132
132
  }
133
133
 
134
- copy_template_dir(src_dir, dst_dir, ctx)
135
- copy_template_dir(
136
- path.join(src_dir, "__package"), path.join(dst_dir, package_name), ctx
137
- )
134
+ if config_mode:
135
+ copy_dbos_template(
136
+ os.path.join(src_dir, "dbos-config.yaml.dbos"),
137
+ os.path.join(dst_dir, "dbos-config.yaml"),
138
+ ctx,
139
+ )
140
+ else:
141
+ copy_template_dir(src_dir, dst_dir, ctx)
142
+ copy_template_dir(
143
+ path.join(src_dir, "__package"), path.join(dst_dir, package_name), ctx
144
+ )
138
145
 
139
146
 
140
147
  def get_project_name() -> typing.Union[str, None]:
@@ -173,6 +180,10 @@ def init(
173
180
  typing.Optional[str],
174
181
  typer.Option("--template", "-t", help="Specify template to use"),
175
182
  ] = None,
183
+ config: Annotated[
184
+ bool,
185
+ typer.Option("--config", "-c", help="Only add dbos-config.yaml"),
186
+ ] = False,
176
187
  ) -> None:
177
188
  try:
178
189
  if project_name is None:
@@ -199,7 +210,9 @@ def init(
199
210
  if template not in templates:
200
211
  raise Exception(f"template {template} not found in {templates_dir}")
201
212
 
202
- copy_template(path.join(templates_dir, template), project_name)
213
+ copy_template(
214
+ path.join(templates_dir, template), project_name, config_mode=config
215
+ )
203
216
  except Exception as e:
204
217
  print(f"[red]{e}[/red]")
205
218
 
@@ -1,5 +1,6 @@
1
1
  from __future__ import annotations
2
2
 
3
+ import json
3
4
  import os
4
5
  import uuid
5
6
  from contextvars import ContextVar
@@ -39,6 +40,9 @@ class TracedAttributes(TypedDict, total=False):
39
40
  applicationID: Optional[str]
40
41
  applicationVersion: Optional[str]
41
42
  executorID: Optional[str]
43
+ authenticatedUser: Optional[str]
44
+ authenticatedUserRoles: Optional[str]
45
+ authenticatedUserAssumedRole: Optional[str]
42
46
 
43
47
 
44
48
  class DBOSContext:
@@ -159,6 +163,13 @@ class DBOSContext:
159
163
  attributes["operationUUID"] = (
160
164
  self.workflow_id if len(self.workflow_id) > 0 else None
161
165
  )
166
+ attributes["authenticatedUser"] = self.authenticated_user
167
+ attributes["authenticatedUserRoles"] = (
168
+ json.dumps(self.authenticated_roles)
169
+ if self.authenticated_roles is not None
170
+ else ""
171
+ )
172
+ attributes["authenticatedUserAssumedRole"] = self.assumed_role
162
173
  span = dbos_tracer.start_span(
163
174
  attributes, parent=self.spans[-1] if len(self.spans) > 0 else None
164
175
  )
@@ -178,6 +189,11 @@ class DBOSContext:
178
189
  ) -> None:
179
190
  self.authenticated_user = user
180
191
  self.authenticated_roles = roles
192
+ if user is not None and len(self.spans) > 0:
193
+ self.spans[-1].set_attribute("authenticatedUser", user)
194
+ self.spans[-1].set_attribute(
195
+ "authenticatedUserRoles", json.dumps(roles) if roles is not None else ""
196
+ )
181
197
 
182
198
 
183
199
  ##############################################################
@@ -217,13 +233,13 @@ class DBOSContextEnsure:
217
233
  def __init__(self) -> None:
218
234
  self.created_ctx = False
219
235
 
220
- def __enter__(self) -> DBOSContextEnsure:
236
+ def __enter__(self) -> DBOSContext:
221
237
  # Code to create a basic context
222
238
  ctx = get_local_dbos_context()
223
239
  if ctx is None:
224
240
  self.created_ctx = True
225
241
  set_local_dbos_context(DBOSContext())
226
- return self
242
+ return assert_current_dbos_context()
227
243
 
228
244
  def __exit__(
229
245
  self,
@@ -1,9 +1,20 @@
1
+ import json
1
2
  import sys
2
3
  import time
3
4
  import traceback
4
5
  from concurrent.futures import Future
5
6
  from functools import wraps
6
- from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, Tuple, TypeVar, cast
7
+ from typing import (
8
+ TYPE_CHECKING,
9
+ Any,
10
+ Callable,
11
+ Generic,
12
+ List,
13
+ Optional,
14
+ Tuple,
15
+ TypeVar,
16
+ cast,
17
+ )
7
18
 
8
19
  from dbos.application_database import ApplicationDatabase, TransactionResultInternal
9
20
 
@@ -134,6 +145,11 @@ def _init_workflow(
134
145
  "executor_id": ctx.executor_id,
135
146
  "request": (utils.serialize(ctx.request) if ctx.request is not None else None),
136
147
  "recovery_attempts": None,
148
+ "authenticated_user": ctx.authenticated_user,
149
+ "authenticated_roles": (
150
+ json.dumps(ctx.authenticated_roles) if ctx.authenticated_roles else None
151
+ ),
152
+ "assumed_role": ctx.assumed_role,
137
153
  }
138
154
 
139
155
  # If we have a class name, the first arg is the instance and do not serialize
@@ -1,6 +1,7 @@
1
1
  from __future__ import annotations
2
2
 
3
3
  import atexit
4
+ import json
4
5
  import os
5
6
  import sys
6
7
  import threading
@@ -19,6 +20,7 @@ from typing import (
19
20
  Tuple,
20
21
  Type,
21
22
  TypeVar,
23
+ cast,
22
24
  )
23
25
 
24
26
  from opentelemetry.trace import Span
@@ -53,6 +55,7 @@ from .tracer import dbos_tracer
53
55
  if TYPE_CHECKING:
54
56
  from fastapi import FastAPI
55
57
  from .request import Request
58
+ from flask import Flask
56
59
 
57
60
  from sqlalchemy.orm import Session
58
61
 
@@ -205,6 +208,7 @@ class DBOS:
205
208
  *,
206
209
  config: Optional[ConfigFile] = None,
207
210
  fastapi: Optional["FastAPI"] = None,
211
+ flask: Optional["Flask"] = None,
208
212
  ) -> DBOS:
209
213
  global _dbos_global_instance
210
214
  global _dbos_global_registry
@@ -219,7 +223,7 @@ class DBOS:
219
223
  )
220
224
  config = _dbos_global_registry.config
221
225
  _dbos_global_instance = super().__new__(cls)
222
- _dbos_global_instance.__init__(fastapi=fastapi, config=config) # type: ignore
226
+ _dbos_global_instance.__init__(fastapi=fastapi, config=config, flask=flask) # type: ignore
223
227
  else:
224
228
  if (config is not None and _dbos_global_instance.config is not config) or (
225
229
  _dbos_global_instance.fastapi is not fastapi
@@ -243,6 +247,7 @@ class DBOS:
243
247
  *,
244
248
  config: Optional[ConfigFile] = None,
245
249
  fastapi: Optional["FastAPI"] = None,
250
+ flask: Optional["Flask"] = None,
246
251
  ) -> None:
247
252
  if hasattr(self, "_initialized") and self._initialized:
248
253
  return
@@ -264,14 +269,44 @@ class DBOS:
264
269
  self._admin_server: Optional[AdminServer] = None
265
270
  self.stop_events: List[threading.Event] = []
266
271
  self.fastapi: Optional["FastAPI"] = fastapi
272
+ self.flask: Optional["Flask"] = flask
267
273
  self._executor: Optional[ThreadPoolExecutor] = None
274
+
275
+ # If using FastAPI, set up middleware and lifecycle events
268
276
  if self.fastapi is not None:
277
+ from fastapi.requests import Request as FARequest
278
+ from fastapi.responses import JSONResponse
279
+
280
+ async def dbos_error_handler(
281
+ request: FARequest, gexc: Exception
282
+ ) -> JSONResponse:
283
+ exc: DBOSException = cast(DBOSException, gexc)
284
+ status_code = 500
285
+ if exc.status_code is not None:
286
+ status_code = exc.status_code
287
+ return JSONResponse(
288
+ status_code=status_code,
289
+ content={
290
+ "message": str(exc.message),
291
+ "dbos_error_code": str(exc.dbos_error_code),
292
+ "dbos_error": str(exc.__class__.__name__),
293
+ },
294
+ )
295
+
296
+ self.fastapi.add_exception_handler(DBOSException, dbos_error_handler)
297
+
269
298
  from dbos.fastapi import setup_fastapi_middleware
270
299
 
271
300
  setup_fastapi_middleware(self.fastapi)
272
301
  self.fastapi.on_event("startup")(self._launch)
273
302
  self.fastapi.on_event("shutdown")(self._destroy)
274
303
 
304
+ # If using Flask, set up middleware
305
+ if self.flask is not None:
306
+ from dbos.flask import setup_flask_middleware
307
+
308
+ setup_flask_middleware(self.flask)
309
+
275
310
  # Register send_stub as a workflow
276
311
  def send_temp_workflow(
277
312
  destination_id: str, message: Any, topic: Optional[str]
@@ -502,9 +537,13 @@ class DBOS:
502
537
  recovery_attempts=stat["recovery_attempts"],
503
538
  class_name=stat["class_name"],
504
539
  config_name=stat["config_name"],
505
- authenticated_user=None,
506
- assumed_role=None,
507
- authenticatedRoles=None,
540
+ authenticated_user=stat["authenticated_user"],
541
+ assumed_role=stat["assumed_role"],
542
+ authenticated_roles=(
543
+ json.loads(stat["authenticated_roles"])
544
+ if stat["authenticated_roles"] is not None
545
+ else None
546
+ ),
508
547
  )
509
548
 
510
549
  @classmethod
@@ -663,6 +702,27 @@ class DBOS:
663
702
  ctx = assert_current_dbos_context()
664
703
  return ctx.request
665
704
 
705
+ @classproperty
706
+ def authenticated_user(cls) -> Optional[str]:
707
+ """Return the current authenticated user, if any, associated with the current context."""
708
+ ctx = assert_current_dbos_context()
709
+ return ctx.authenticated_user
710
+
711
+ @classproperty
712
+ def authenticated_roles(cls) -> Optional[List[str]]:
713
+ """Return the roles granted to the current authenticated user, if any, associated with the current context."""
714
+ ctx = assert_current_dbos_context()
715
+ return ctx.authenticated_roles
716
+
717
+ @classmethod
718
+ def set_authentication(
719
+ cls, authenticated_user: Optional[str], authenticated_roles: Optional[List[str]]
720
+ ) -> None:
721
+ """Set the current authenticated user and granted roles into the current context."""
722
+ ctx = assert_current_dbos_context()
723
+ ctx.authenticated_user = authenticated_user
724
+ ctx.authenticated_roles = authenticated_roles
725
+
666
726
 
667
727
  @dataclass
668
728
  class WorkflowStatus:
@@ -679,7 +739,7 @@ class WorkflowStatus:
679
739
  config_name(str): For instance member functions, the name of the class instance for the execution
680
740
  authenticated_user(str): The user who invoked the workflow
681
741
  assumed_role(str): The access role used by the user to allow access to the workflow function
682
- authenticatedRoles(List[str]): List of all access roles available to the authenticated user
742
+ authenticated_roles(List[str]): List of all access roles available to the authenticated user
683
743
  recovery_attempts(int): Number of times the workflow has been restarted (usually by recovery)
684
744
 
685
745
  """
@@ -691,7 +751,7 @@ class WorkflowStatus:
691
751
  config_name: Optional[str]
692
752
  authenticated_user: Optional[str]
693
753
  assumed_role: Optional[str]
694
- authenticatedRoles: Optional[List[str]]
754
+ authenticated_roles: Optional[List[str]]
695
755
  recovery_attempts: Optional[int]
696
756
 
697
757
 
@@ -2,7 +2,7 @@ import json
2
2
  import os
3
3
  import re
4
4
  from importlib import resources
5
- from typing import Any, Dict, List, Optional, TypedDict
5
+ from typing import Dict, List, Optional, TypedDict
6
6
 
7
7
  import yaml
8
8
  from jsonschema import ValidationError, validate
@@ -17,6 +17,7 @@ class DBOSException(Exception):
17
17
  def __init__(self, message: str, dbos_error_code: Optional[int] = None):
18
18
  self.message = message
19
19
  self.dbos_error_code = dbos_error_code
20
+ self.status_code: Optional[int] = None
20
21
  super().__init__(self.message)
21
22
 
22
23
  def __str__(self) -> str:
@@ -104,6 +105,7 @@ class DBOSNotAuthorizedError(DBOSException):
104
105
  msg,
105
106
  dbos_error_code=DBOSErrorCode.NotAuthorized.value,
106
107
  )
108
+ self.status_code = 403
107
109
 
108
110
 
109
111
  class DBOSMaxStepRetriesExceeded(DBOSException):
@@ -11,9 +11,7 @@ from .context import (
11
11
  TracedAttributes,
12
12
  assert_current_dbos_context,
13
13
  )
14
- from .request import Address, Request
15
-
16
- request_id_header = "x-request-id"
14
+ from .request import Address, Request, request_id_header
17
15
 
18
16
 
19
17
  def get_or_generate_request_id(request: FastAPIRequest) -> str:
@@ -0,0 +1,79 @@
1
+ import uuid
2
+ from typing import Any
3
+ from urllib.parse import urlparse
4
+
5
+ from flask import Flask, request
6
+ from werkzeug.wrappers import Request as WRequest
7
+
8
+ from dbos.context import (
9
+ EnterDBOSHandler,
10
+ OperationType,
11
+ SetWorkflowID,
12
+ TracedAttributes,
13
+ assert_current_dbos_context,
14
+ )
15
+
16
+ from .request import Address, Request, request_id_header
17
+
18
+
19
+ class FlaskMiddleware:
20
+ def __init__(self, app: Any) -> None:
21
+ self.app = app
22
+
23
+ def __call__(self, environ: Any, start_response: Any) -> Any:
24
+ request = WRequest(environ)
25
+ attributes: TracedAttributes = {
26
+ "name": urlparse(request.url).path,
27
+ "requestID": get_or_generate_request_id(request),
28
+ "requestIP": (
29
+ request.remote_addr if request.remote_addr is not None else None
30
+ ),
31
+ "requestURL": request.url,
32
+ "requestMethod": request.method,
33
+ "operationType": OperationType.HANDLER.value,
34
+ }
35
+ with EnterDBOSHandler(attributes):
36
+ ctx = assert_current_dbos_context()
37
+ ctx.request = make_request(request)
38
+ workflow_id = request.headers.get("dbos-idempotency-key", "")
39
+ with SetWorkflowID(workflow_id):
40
+ response = self.app(environ, start_response)
41
+ return response
42
+
43
+
44
+ def get_or_generate_request_id(request: WRequest) -> str:
45
+ request_id = request.headers.get(request_id_header, None)
46
+ if request_id is not None:
47
+ return request_id
48
+ else:
49
+ return str(uuid.uuid4())
50
+
51
+
52
+ def make_request(request: WRequest) -> Request:
53
+ parsed_url = urlparse(request.url)
54
+ base_url = f"{parsed_url.scheme}://{parsed_url.netloc}"
55
+
56
+ client = None
57
+ if request.remote_addr:
58
+ hostname = request.remote_addr
59
+ port = request.environ.get("REMOTE_PORT")
60
+ if port:
61
+ client = Address(hostname=hostname, port=int(port))
62
+ else:
63
+ # If port is not available, use 0 as a placeholder
64
+ client = Address(hostname=hostname, port=0)
65
+
66
+ return Request(
67
+ headers=dict(request.headers),
68
+ path_params={},
69
+ query_params=dict(request.args),
70
+ url=request.url,
71
+ base_url=base_url,
72
+ client=client,
73
+ cookies=dict(request.cookies),
74
+ method=request.method,
75
+ )
76
+
77
+
78
+ def setup_flask_middleware(app: Flask) -> None:
79
+ app.wsgi_app = FlaskMiddleware(app.wsgi_app) # type: ignore