dbos 0.20.0a9__tar.gz → 0.21.0a3__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 (90) hide show
  1. {dbos-0.20.0a9 → dbos-0.21.0a3}/PKG-INFO +1 -1
  2. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_context.py +0 -30
  3. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_core.py +2 -2
  4. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_dbos.py +1 -1
  5. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_recovery.py +2 -5
  6. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_sys_db.py +81 -61
  7. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_workflow_commands.py +1 -1
  8. {dbos-0.20.0a9 → dbos-0.21.0a3}/pyproject.toml +1 -1
  9. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_dbos.py +50 -88
  10. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_failures.py +0 -4
  11. dbos-0.21.0a3/tests/test_sqlalchemy.py +113 -0
  12. {dbos-0.20.0a9 → dbos-0.21.0a3}/LICENSE +0 -0
  13. {dbos-0.20.0a9 → dbos-0.21.0a3}/README.md +0 -0
  14. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/__init__.py +0 -0
  15. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_admin_server.py +0 -0
  16. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_app_db.py +0 -0
  17. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_classproperty.py +0 -0
  18. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_cloudutils/authentication.py +0 -0
  19. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_cloudutils/cloudutils.py +0 -0
  20. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_cloudutils/databases.py +0 -0
  21. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_croniter.py +0 -0
  22. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_db_wizard.py +0 -0
  23. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_dbos_config.py +0 -0
  24. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_error.py +0 -0
  25. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_fastapi.py +0 -0
  26. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_flask.py +0 -0
  27. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_kafka.py +0 -0
  28. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_kafka_message.py +0 -0
  29. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_logger.py +0 -0
  30. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_migrations/env.py +0 -0
  31. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_migrations/script.py.mako +0 -0
  32. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_migrations/versions/04ca4f231047_workflow_queues_executor_id.py +0 -0
  33. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_migrations/versions/50f3227f0b4b_fix_job_queue.py +0 -0
  34. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_migrations/versions/5c361fc04708_added_system_tables.py +0 -0
  35. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_migrations/versions/a3b18ad34abe_added_triggers.py +0 -0
  36. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_migrations/versions/d76646551a6b_job_queue_limiter.py +0 -0
  37. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_migrations/versions/d76646551a6c_workflow_queue.py +0 -0
  38. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_migrations/versions/eab0cc1d9a14_job_queue.py +0 -0
  39. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_outcome.py +0 -0
  40. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_queue.py +0 -0
  41. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_registrations.py +0 -0
  42. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_request.py +0 -0
  43. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_roles.py +0 -0
  44. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_scheduler.py +0 -0
  45. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_schemas/__init__.py +0 -0
  46. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_schemas/application_database.py +0 -0
  47. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_schemas/system_database.py +0 -0
  48. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_serialization.py +0 -0
  49. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_templates/dbos-db-starter/README.md +0 -0
  50. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_templates/dbos-db-starter/__package/__init__.py +0 -0
  51. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_templates/dbos-db-starter/__package/main.py +0 -0
  52. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_templates/dbos-db-starter/__package/schema.py +0 -0
  53. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_templates/dbos-db-starter/alembic.ini +0 -0
  54. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_templates/dbos-db-starter/dbos-config.yaml.dbos +0 -0
  55. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_templates/dbos-db-starter/migrations/env.py.dbos +0 -0
  56. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_templates/dbos-db-starter/migrations/script.py.mako +0 -0
  57. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_templates/dbos-db-starter/migrations/versions/2024_07_31_180642_init.py +0 -0
  58. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_templates/dbos-db-starter/start_postgres_docker.py +0 -0
  59. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/_tracer.py +0 -0
  60. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/cli/_github_init.py +0 -0
  61. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/cli/_template_init.py +0 -0
  62. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/cli/cli.py +0 -0
  63. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/dbos-config.schema.json +0 -0
  64. {dbos-0.20.0a9 → dbos-0.21.0a3}/dbos/py.typed +0 -0
  65. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/__init__.py +0 -0
  66. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/atexit_no_ctor.py +0 -0
  67. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/atexit_no_launch.py +0 -0
  68. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/classdefs.py +0 -0
  69. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/conftest.py +0 -0
  70. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/more_classdefs.py +0 -0
  71. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/queuedworkflow.py +0 -0
  72. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_admin_server.py +0 -0
  73. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_async.py +0 -0
  74. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_classdecorators.py +0 -0
  75. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_concurrency.py +0 -0
  76. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_config.py +0 -0
  77. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_croniter.py +0 -0
  78. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_fastapi.py +0 -0
  79. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_fastapi_roles.py +0 -0
  80. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_flask.py +0 -0
  81. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_kafka.py +0 -0
  82. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_outcome.py +0 -0
  83. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_package.py +0 -0
  84. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_queue.py +0 -0
  85. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_scheduler.py +0 -0
  86. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_schema_migration.py +0 -0
  87. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_singleton.py +0 -0
  88. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_spans.py +0 -0
  89. {dbos-0.20.0a9 → dbos-0.21.0a3}/tests/test_workflow_cmds.py +0 -0
  90. {dbos-0.20.0a9 → dbos-0.21.0a3}/version/__init__.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: dbos
3
- Version: 0.20.0a9
3
+ Version: 0.21.0a3
4
4
  Summary: Ultra-lightweight durable execution in Python
5
5
  Author-Email: "DBOS, Inc." <contact@dbos.dev>
6
6
  License: MIT
@@ -63,7 +63,6 @@ class DBOSContext:
63
63
  self.parent_workflow_fid: int = -1
64
64
  self.workflow_id: str = ""
65
65
  self.function_id: int = -1
66
- self.in_recovery: bool = False
67
66
 
68
67
  self.curr_step_function_id: int = -1
69
68
  self.curr_tx_function_id: int = -1
@@ -82,7 +81,6 @@ class DBOSContext:
82
81
  rv.is_within_set_workflow_id_block = self.is_within_set_workflow_id_block
83
82
  rv.parent_workflow_id = self.workflow_id
84
83
  rv.parent_workflow_fid = self.function_id
85
- rv.in_recovery = self.in_recovery
86
84
  rv.authenticated_user = self.authenticated_user
87
85
  rv.authenticated_roles = (
88
86
  self.authenticated_roles[:]
@@ -335,34 +333,6 @@ class SetWorkflowID:
335
333
  return False # Did not handle
336
334
 
337
335
 
338
- class SetWorkflowRecovery:
339
- def __init__(self) -> None:
340
- self.created_ctx = False
341
-
342
- def __enter__(self) -> SetWorkflowRecovery:
343
- # Code to create a basic context
344
- ctx = get_local_dbos_context()
345
- if ctx is None:
346
- self.created_ctx = True
347
- _set_local_dbos_context(DBOSContext())
348
- assert_current_dbos_context().in_recovery = True
349
-
350
- return self
351
-
352
- def __exit__(
353
- self,
354
- exc_type: Optional[Type[BaseException]],
355
- exc_value: Optional[BaseException],
356
- traceback: Optional[TracebackType],
357
- ) -> Literal[False]:
358
- assert assert_current_dbos_context().in_recovery == True
359
- assert_current_dbos_context().in_recovery = False
360
- # Code to clean up the basic context if we created it
361
- if self.created_ctx:
362
- _clear_local_dbos_context()
363
- return False # Did not handle
364
-
365
-
366
336
  class EnterDBOSWorkflow(AbstractContextManager[DBOSContext, Literal[False]]):
367
337
  def __init__(self, attributes: TracedAttributes) -> None:
368
338
  self.created_ctx = False
@@ -185,8 +185,8 @@ def _init_workflow(
185
185
  # Synchronously record the status and inputs for workflows and single-step workflows
186
186
  # We also have to do this for single-step workflows because of the foreign key constraint on the operation outputs table
187
187
  # TODO: Make this transactional (and with the queue step below)
188
- wf_status = dbos._sys_db.update_workflow_status(
189
- status, False, ctx.in_recovery, max_recovery_attempts=max_recovery_attempts
188
+ wf_status = dbos._sys_db.insert_workflow_status(
189
+ status, max_recovery_attempts=max_recovery_attempts
190
190
  )
191
191
  # TODO: Modify the inputs if they were changed by `update_workflow_inputs`
192
192
  dbos._sys_db.update_workflow_inputs(wfid, _serialization.serialize_args(inputs))
@@ -801,7 +801,7 @@ class DBOS:
801
801
  def cancel_workflow(cls, workflow_id: str) -> None:
802
802
  """Cancel a workflow by ID."""
803
803
  _get_dbos_instance()._sys_db.set_workflow_status(
804
- workflow_id, WorkflowStatusString.CANCELLED, False
804
+ workflow_id, WorkflowStatusString.CANCELLED
805
805
  )
806
806
 
807
807
  @classmethod
@@ -4,7 +4,6 @@ import time
4
4
  import traceback
5
5
  from typing import TYPE_CHECKING, Any, List
6
6
 
7
- from ._context import SetWorkflowRecovery
8
7
  from ._core import execute_workflow_by_id
9
8
  from ._error import DBOSWorkflowFunctionNotFoundError
10
9
 
@@ -19,8 +18,7 @@ def startup_recovery_thread(dbos: "DBOS", workflow_ids: List[str]) -> None:
19
18
  while not stop_event.is_set() and len(workflow_ids) > 0:
20
19
  try:
21
20
  for workflowID in list(workflow_ids):
22
- with SetWorkflowRecovery():
23
- execute_workflow_by_id(dbos, workflowID)
21
+ execute_workflow_by_id(dbos, workflowID)
24
22
  workflow_ids.remove(workflowID)
25
23
  except DBOSWorkflowFunctionNotFoundError:
26
24
  time.sleep(1)
@@ -45,8 +43,7 @@ def recover_pending_workflows(
45
43
  dbos.logger.debug(f"Pending workflows: {workflow_ids}")
46
44
 
47
45
  for workflowID in workflow_ids:
48
- with SetWorkflowRecovery():
49
- handle = execute_workflow_by_id(dbos, workflowID)
46
+ handle = execute_workflow_by_id(dbos, workflowID)
50
47
  workflow_handles.append(handle)
51
48
 
52
49
  dbos.logger.info("Recovered pending workflows")
@@ -243,66 +243,50 @@ class SystemDatabase:
243
243
  dbos_logger.debug("Waiting for system buffers to be exported")
244
244
  time.sleep(1)
245
245
 
246
- def update_workflow_status(
246
+ def insert_workflow_status(
247
247
  self,
248
248
  status: WorkflowStatusInternal,
249
- replace: bool = True,
250
- in_recovery: bool = False,
251
249
  *,
252
- conn: Optional[sa.Connection] = None,
253
250
  max_recovery_attempts: int = DEFAULT_MAX_RECOVERY_ATTEMPTS,
254
251
  ) -> WorkflowStatuses:
255
252
  wf_status: WorkflowStatuses = status["status"]
256
253
 
257
- cmd = pg.insert(SystemSchema.workflow_status).values(
258
- workflow_uuid=status["workflow_uuid"],
259
- status=status["status"],
260
- name=status["name"],
261
- class_name=status["class_name"],
262
- config_name=status["config_name"],
263
- output=status["output"],
264
- error=status["error"],
265
- executor_id=status["executor_id"],
266
- application_version=status["app_version"],
267
- application_id=status["app_id"],
268
- request=status["request"],
269
- authenticated_user=status["authenticated_user"],
270
- authenticated_roles=status["authenticated_roles"],
271
- assumed_role=status["assumed_role"],
272
- queue_name=status["queue_name"],
273
- )
274
- if replace:
275
- cmd = cmd.on_conflict_do_update(
276
- index_elements=["workflow_uuid"],
277
- set_=dict(
278
- status=status["status"],
279
- output=status["output"],
280
- error=status["error"],
281
- ),
282
- )
283
- elif in_recovery:
284
- cmd = cmd.on_conflict_do_update(
285
- index_elements=["workflow_uuid"],
286
- set_=dict(
287
- recovery_attempts=SystemSchema.workflow_status.c.recovery_attempts
288
- + 1,
254
+ cmd = (
255
+ pg.insert(SystemSchema.workflow_status)
256
+ .values(
257
+ workflow_uuid=status["workflow_uuid"],
258
+ status=status["status"],
259
+ name=status["name"],
260
+ class_name=status["class_name"],
261
+ config_name=status["config_name"],
262
+ output=status["output"],
263
+ error=status["error"],
264
+ executor_id=status["executor_id"],
265
+ application_version=status["app_version"],
266
+ application_id=status["app_id"],
267
+ request=status["request"],
268
+ authenticated_user=status["authenticated_user"],
269
+ authenticated_roles=status["authenticated_roles"],
270
+ assumed_role=status["assumed_role"],
271
+ queue_name=status["queue_name"],
272
+ recovery_attempts=(
273
+ 1 if wf_status != WorkflowStatusString.ENQUEUED.value else 0
289
274
  ),
290
275
  )
291
- else:
292
- # A blank update so that we can return the existing status
293
- cmd = cmd.on_conflict_do_update(
276
+ .on_conflict_do_update(
294
277
  index_elements=["workflow_uuid"],
295
278
  set_=dict(
296
- recovery_attempts=SystemSchema.workflow_status.c.recovery_attempts
279
+ recovery_attempts=(
280
+ SystemSchema.workflow_status.c.recovery_attempts + 1
281
+ ),
297
282
  ),
298
283
  )
284
+ )
285
+
299
286
  cmd = cmd.returning(SystemSchema.workflow_status.c.recovery_attempts, SystemSchema.workflow_status.c.status, SystemSchema.workflow_status.c.name, SystemSchema.workflow_status.c.class_name, SystemSchema.workflow_status.c.config_name, SystemSchema.workflow_status.c.queue_name) # type: ignore
300
287
 
301
- if conn is not None:
302
- results = conn.execute(cmd)
303
- else:
304
- with self.engine.begin() as c:
305
- results = c.execute(cmd)
288
+ with self.engine.begin() as c:
289
+ results = c.execute(cmd)
306
290
 
307
291
  row = results.fetchone()
308
292
  if row is not None:
@@ -325,7 +309,9 @@ class SystemDatabase:
325
309
  if err_msg is not None:
326
310
  raise DBOSConflictingWorkflowError(status["workflow_uuid"], err_msg)
327
311
 
328
- if in_recovery and recovery_attempts > max_recovery_attempts:
312
+ # Every time we start executing a workflow (and thus attempt to insert its status), we increment `recovery_attempts` by 1.
313
+ # When this number becomes equal to `maxRetries + 1`, we mark the workflow as `RETRIES_EXCEEDED`.
314
+ if recovery_attempts > max_recovery_attempts + 1:
329
315
  with self.engine.begin() as c:
330
316
  c.execute(
331
317
  sa.delete(SystemSchema.workflow_queue).where(
@@ -352,17 +338,62 @@ class SystemDatabase:
352
338
  status["workflow_uuid"], max_recovery_attempts
353
339
  )
354
340
 
355
- # Record we have exported status for this single-transaction workflow
341
+ return wf_status
342
+
343
+ def update_workflow_status(
344
+ self,
345
+ status: WorkflowStatusInternal,
346
+ *,
347
+ conn: Optional[sa.Connection] = None,
348
+ ) -> None:
349
+ wf_status: WorkflowStatuses = status["status"]
350
+
351
+ cmd = (
352
+ pg.insert(SystemSchema.workflow_status)
353
+ .values(
354
+ workflow_uuid=status["workflow_uuid"],
355
+ status=status["status"],
356
+ name=status["name"],
357
+ class_name=status["class_name"],
358
+ config_name=status["config_name"],
359
+ output=status["output"],
360
+ error=status["error"],
361
+ executor_id=status["executor_id"],
362
+ application_version=status["app_version"],
363
+ application_id=status["app_id"],
364
+ request=status["request"],
365
+ authenticated_user=status["authenticated_user"],
366
+ authenticated_roles=status["authenticated_roles"],
367
+ assumed_role=status["assumed_role"],
368
+ queue_name=status["queue_name"],
369
+ recovery_attempts=(
370
+ 1 if wf_status != WorkflowStatusString.ENQUEUED.value else 0
371
+ ),
372
+ )
373
+ .on_conflict_do_update(
374
+ index_elements=["workflow_uuid"],
375
+ set_=dict(
376
+ status=status["status"],
377
+ output=status["output"],
378
+ error=status["error"],
379
+ ),
380
+ )
381
+ )
382
+
383
+ if conn is not None:
384
+ conn.execute(cmd)
385
+ else:
386
+ with self.engine.begin() as c:
387
+ c.execute(cmd)
388
+
389
+ # If this is a single-transaction workflow, record that its status has been exported
356
390
  if status["workflow_uuid"] in self._temp_txn_wf_ids:
357
391
  self._exported_temp_txn_wf_status.add(status["workflow_uuid"])
358
392
 
359
- return wf_status
360
-
361
393
  def set_workflow_status(
362
394
  self,
363
395
  workflow_uuid: str,
364
396
  status: WorkflowStatusString,
365
- reset_recovery_attempts: bool,
366
397
  ) -> None:
367
398
  with self.engine.begin() as c:
368
399
  stmt = (
@@ -374,17 +405,6 @@ class SystemDatabase:
374
405
  )
375
406
  c.execute(stmt)
376
407
 
377
- if reset_recovery_attempts:
378
- with self.engine.begin() as c:
379
- stmt = (
380
- sa.update(SystemSchema.workflow_status)
381
- .where(
382
- SystemSchema.workflow_status.c.workflow_uuid == workflow_uuid
383
- )
384
- .values(recovery_attempts=reset_recovery_attempts)
385
- )
386
- c.execute(stmt)
387
-
388
408
  def get_workflow_status(
389
409
  self, workflow_uuid: str
390
410
  ) -> Optional[WorkflowStatusInternal]:
@@ -116,7 +116,7 @@ def _cancel_workflow(config: ConfigFile, uuid: str) -> None:
116
116
 
117
117
  try:
118
118
  sys_db = SystemDatabase(config)
119
- sys_db.set_workflow_status(uuid, WorkflowStatusString.CANCELLED, False)
119
+ sys_db.set_workflow_status(uuid, WorkflowStatusString.CANCELLED)
120
120
  return
121
121
 
122
122
  except Exception as e:
@@ -27,7 +27,7 @@ dependencies = [
27
27
  ]
28
28
  requires-python = ">=3.9"
29
29
  readme = "README.md"
30
- version = "0.20.0a9"
30
+ version = "0.21.0a3"
31
31
 
32
32
  [project.license]
33
33
  text = "MIT"
@@ -14,6 +14,7 @@ from dbos import DBOS, ConfigFile, SetWorkflowID, WorkflowHandle, WorkflowStatus
14
14
  # Private API because this is a test
15
15
  from dbos._context import assert_current_dbos_context, get_local_dbos_context
16
16
  from dbos._error import DBOSMaxStepRetriesExceeded
17
+ from dbos._schemas.system_database import SystemSchema
17
18
  from dbos._sys_db import GetWorkflowsInput
18
19
 
19
20
 
@@ -63,6 +64,23 @@ def test_simple_workflow(dbos: DBOS) -> None:
63
64
  assert wf_counter == 4
64
65
 
65
66
 
67
+ def test_simple_workflow_attempts_counter(dbos: DBOS) -> None:
68
+ @DBOS.workflow()
69
+ def noop() -> None:
70
+ pass
71
+
72
+ wfuuid = str(uuid.uuid4())
73
+ with dbos._sys_db.engine.connect() as c:
74
+ stmt = sa.select(SystemSchema.workflow_status.c.recovery_attempts).where(
75
+ SystemSchema.workflow_status.c.workflow_uuid == wfuuid
76
+ )
77
+ for i in range(10):
78
+ with SetWorkflowID(wfuuid):
79
+ noop()
80
+ result = c.execute(stmt).scalar()
81
+ assert result == i + 1
82
+
83
+
66
84
  def test_child_workflow(dbos: DBOS) -> None:
67
85
  txn_counter: int = 0
68
86
  wf_counter: int = 0
@@ -365,26 +383,12 @@ def test_recovery_workflow(dbos: DBOS) -> None:
365
383
 
366
384
  dbos._sys_db.wait_for_buffer_flush()
367
385
  # Change the workflow status to pending
368
- dbos._sys_db.update_workflow_status(
369
- {
370
- "workflow_uuid": wfuuid,
371
- "status": "PENDING",
372
- "name": test_workflow.__qualname__,
373
- "class_name": None,
374
- "config_name": None,
375
- "output": None,
376
- "error": None,
377
- "executor_id": None,
378
- "app_id": None,
379
- "app_version": None,
380
- "request": None,
381
- "recovery_attempts": None,
382
- "authenticated_user": None,
383
- "authenticated_roles": None,
384
- "assumed_role": None,
385
- "queue_name": None,
386
- }
387
- )
386
+ with dbos._sys_db.engine.begin() as c:
387
+ c.execute(
388
+ sa.update(SystemSchema.workflow_status)
389
+ .values({"status": "PENDING", "name": test_workflow.__qualname__})
390
+ .where(SystemSchema.workflow_status.c.workflow_uuid == wfuuid)
391
+ )
388
392
 
389
393
  # Recovery should execute the workflow again but skip the transaction
390
394
  workflow_handles = DBOS.recover_pending_workflows()
@@ -397,7 +401,7 @@ def test_recovery_workflow(dbos: DBOS) -> None:
397
401
  # Test that there was a recovery attempt of this
398
402
  stat = workflow_handles[0].get_status()
399
403
  assert stat
400
- assert stat.recovery_attempts == 1
404
+ assert stat.recovery_attempts == 2 # original attempt + recovery attempt
401
405
 
402
406
 
403
407
  def test_recovery_workflow_step(dbos: DBOS) -> None:
@@ -425,26 +429,12 @@ def test_recovery_workflow_step(dbos: DBOS) -> None:
425
429
 
426
430
  dbos._sys_db.wait_for_buffer_flush()
427
431
  # Change the workflow status to pending
428
- dbos._sys_db.update_workflow_status(
429
- {
430
- "workflow_uuid": wfuuid,
431
- "status": "PENDING",
432
- "name": test_workflow.__qualname__,
433
- "class_name": None,
434
- "config_name": None,
435
- "output": None,
436
- "error": None,
437
- "executor_id": None,
438
- "app_id": None,
439
- "app_version": None,
440
- "request": None,
441
- "recovery_attempts": None,
442
- "authenticated_user": None,
443
- "authenticated_roles": None,
444
- "assumed_role": None,
445
- "queue_name": None,
446
- }
447
- )
432
+ with dbos._sys_db.engine.begin() as c:
433
+ c.execute(
434
+ sa.update(SystemSchema.workflow_status)
435
+ .values({"status": "PENDING", "name": test_workflow.__qualname__})
436
+ .where(SystemSchema.workflow_status.c.workflow_uuid == wfuuid)
437
+ )
448
438
 
449
439
  # Recovery should execute the workflow again but skip the transaction
450
440
  workflow_handles = DBOS.recover_pending_workflows()
@@ -456,7 +446,7 @@ def test_recovery_workflow_step(dbos: DBOS) -> None:
456
446
  # Test that there was a recovery attempt of this
457
447
  stat = workflow_handles[0].get_status()
458
448
  assert stat
459
- assert stat.recovery_attempts == 1
449
+ assert stat.recovery_attempts == 2
460
450
 
461
451
 
462
452
  def test_workflow_returns_none(dbos: DBOS) -> None:
@@ -484,26 +474,12 @@ def test_workflow_returns_none(dbos: DBOS) -> None:
484
474
  assert wf_counter == 2
485
475
 
486
476
  # Change the workflow status to pending
487
- dbos._sys_db.update_workflow_status(
488
- {
489
- "workflow_uuid": wfuuid,
490
- "status": "PENDING",
491
- "name": test_workflow.__qualname__,
492
- "class_name": None,
493
- "config_name": None,
494
- "output": None,
495
- "error": None,
496
- "executor_id": None,
497
- "app_id": None,
498
- "app_version": None,
499
- "request": None,
500
- "recovery_attempts": None,
501
- "authenticated_user": None,
502
- "authenticated_roles": None,
503
- "assumed_role": None,
504
- "queue_name": None,
505
- }
506
- )
477
+ with dbos._sys_db.engine.begin() as c:
478
+ c.execute(
479
+ sa.update(SystemSchema.workflow_status)
480
+ .values({"status": "PENDING", "name": test_workflow.__qualname__})
481
+ .where(SystemSchema.workflow_status.c.workflow_uuid == wfuuid)
482
+ )
507
483
 
508
484
  workflow_handles = DBOS.recover_pending_workflows()
509
485
  assert len(workflow_handles) == 1
@@ -513,7 +489,7 @@ def test_workflow_returns_none(dbos: DBOS) -> None:
513
489
  # Test that there was a recovery attempt of this
514
490
  stat = workflow_handles[0].get_status()
515
491
  assert stat
516
- assert stat.recovery_attempts == 1
492
+ assert stat.recovery_attempts == 3 # 2 calls to test_workflow + 1 recovery attempt
517
493
 
518
494
 
519
495
  def test_recovery_temp_workflow(dbos: DBOS) -> None:
@@ -545,26 +521,12 @@ def test_recovery_temp_workflow(dbos: DBOS) -> None:
545
521
  assert wfi["name"].startswith("<temp>")
546
522
 
547
523
  # Change the workflow status to pending
548
- dbos._sys_db.update_workflow_status(
549
- {
550
- "workflow_uuid": wfuuid,
551
- "status": "PENDING",
552
- "name": wfi["name"],
553
- "class_name": None,
554
- "config_name": None,
555
- "output": None,
556
- "error": None,
557
- "executor_id": None,
558
- "app_id": None,
559
- "app_version": None,
560
- "request": None,
561
- "recovery_attempts": None,
562
- "authenticated_user": None,
563
- "authenticated_roles": None,
564
- "assumed_role": None,
565
- "queue_name": None,
566
- }
567
- )
524
+ with dbos._sys_db.engine.begin() as c:
525
+ c.execute(
526
+ sa.update(SystemSchema.workflow_status)
527
+ .values({"status": "PENDING", "name": wfi["name"]})
528
+ .where(SystemSchema.workflow_status.c.workflow_uuid == wfuuid)
529
+ )
568
530
 
569
531
  # Recovery should execute the workflow again but skip the transaction
570
532
  workflow_handles = DBOS.recover_pending_workflows()
@@ -787,19 +749,19 @@ def test_retrieve_workflow_in_workflow(dbos: DBOS) -> None:
787
749
  with SetWorkflowID("parent_b"):
788
750
  assert test_workflow_status_b() == "PENDINGrun_this_once_bSUCCESS"
789
751
 
790
- # Test that there were no recovery attempts of this
752
+ # Test that the number of attempts matches the number of calls
791
753
  stat = dbos.get_workflow_status("parent_a")
792
754
  assert stat
793
- assert stat.recovery_attempts == 0
755
+ assert stat.recovery_attempts == 2
794
756
  stat = dbos.get_workflow_status("parent_b")
795
757
  assert stat
796
- assert stat.recovery_attempts == 0
758
+ assert stat.recovery_attempts == 2
797
759
  stat = dbos.get_workflow_status("run_this_once_a")
798
760
  assert stat
799
- assert stat.recovery_attempts == 0
761
+ assert stat.recovery_attempts == 2
800
762
  stat = dbos.get_workflow_status("run_this_once_b")
801
763
  assert stat
802
- assert stat.recovery_attempts == 0
764
+ assert stat.recovery_attempts == 2
803
765
 
804
766
 
805
767
  def test_sleep(dbos: DBOS) -> None:
@@ -200,10 +200,6 @@ def test_dead_letter_queue(dbos: DBOS) -> None:
200
200
  assert exc_info.errisinstance(DBOSDeadLetterQueueError)
201
201
  assert handle.get_status().status == WorkflowStatusString.RETRIES_EXCEEDED.value
202
202
 
203
- with SetWorkflowID(handle.get_workflow_id()):
204
- DBOS.start_workflow(dead_letter_workflow)
205
- assert recovery_count == max_recovery_attempts + 2
206
-
207
203
  event.set()
208
204
  assert handle.get_result() == None
209
205
  dbos._sys_db.wait_for_buffer_flush()
@@ -0,0 +1,113 @@
1
+ import uuid
2
+
3
+ import pytest
4
+ import sqlalchemy as sa
5
+ from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
6
+
7
+ # Public API
8
+ from dbos import DBOS, SetWorkflowID
9
+
10
+
11
+ # Declare a SQLAlchemy ORM base class
12
+ class Base(DeclarativeBase):
13
+ pass
14
+
15
+
16
+ # Declare a SQLAlchemy ORM class for accessing the database table.
17
+ class Hello(Base):
18
+ __tablename__ = "dbos_hello"
19
+ greet_count: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
20
+ name: Mapped[str] = mapped_column(nullable=False)
21
+
22
+ def __repr__(self) -> str:
23
+ return f"Hello(greet_count={self.greet_count!r}, name={self.name!r})"
24
+
25
+
26
+ def test_simple_transaction(dbos: DBOS, postgres_db_engine: sa.Engine) -> None:
27
+ txn_counter: int = 0
28
+ assert dbos._app_db_field is not None
29
+ Base.metadata.drop_all(dbos._app_db_field.engine)
30
+ Base.metadata.create_all(dbos._app_db_field.engine)
31
+
32
+ @DBOS.transaction()
33
+ def test_transaction(name: str) -> str:
34
+ new_greeting = Hello(name=name)
35
+ DBOS.sql_session.add(new_greeting)
36
+ stmt = (
37
+ sa.select(Hello)
38
+ .where(Hello.name == name)
39
+ .order_by(Hello.greet_count.desc())
40
+ .limit(1)
41
+ )
42
+ row = DBOS.sql_session.scalar(stmt)
43
+ assert row is not None
44
+ greet_count = row.greet_count
45
+ nonlocal txn_counter
46
+ txn_counter += 1
47
+ return name + str(greet_count)
48
+
49
+ assert test_transaction("alice") == "alice1"
50
+ assert test_transaction("alice") == "alice2"
51
+ assert txn_counter == 2
52
+
53
+ # Test OAOO
54
+ wfuuid = str(uuid.uuid4())
55
+ with SetWorkflowID(wfuuid):
56
+ assert test_transaction("alice") == "alice3"
57
+ with SetWorkflowID(wfuuid):
58
+ assert test_transaction("alice") == "alice3"
59
+ assert txn_counter == 3 # Only increment once
60
+
61
+ Base.metadata.drop_all(dbos._app_db_field.engine)
62
+
63
+ # Make sure no transactions are left open
64
+ with postgres_db_engine.begin() as conn:
65
+ result = conn.execute(
66
+ sa.text(
67
+ "select * from pg_stat_activity where state = 'idle in transaction'"
68
+ )
69
+ ).fetchall()
70
+ assert len(result) == 0
71
+
72
+
73
+ def test_error_transaction(dbos: DBOS, postgres_db_engine: sa.Engine) -> None:
74
+ txn_counter: int = 0
75
+ assert dbos._app_db_field is not None
76
+ # Drop the database but don't re-create. Should fail.
77
+ Base.metadata.drop_all(dbos._app_db_field.engine)
78
+
79
+ @DBOS.transaction()
80
+ def test_transaction(name: str) -> str:
81
+ nonlocal txn_counter
82
+ txn_counter += 1
83
+ new_greeting = Hello(name=name)
84
+ DBOS.sql_session.add(new_greeting)
85
+ return name
86
+
87
+ with pytest.raises(Exception) as exc_info:
88
+ test_transaction("alice")
89
+ assert 'relation "dbos_hello" does not exist' in str(exc_info.value)
90
+ assert txn_counter == 1
91
+
92
+ # Test OAOO
93
+ wfuuid = str(uuid.uuid4())
94
+ with SetWorkflowID(wfuuid):
95
+ with pytest.raises(Exception) as exc_info:
96
+ test_transaction("alice")
97
+ assert 'relation "dbos_hello" does not exist' in str(exc_info.value)
98
+ assert txn_counter == 2
99
+
100
+ with SetWorkflowID(wfuuid):
101
+ with pytest.raises(Exception) as exc_info:
102
+ test_transaction("alice")
103
+ assert 'relation "dbos_hello" does not exist' in str(exc_info.value)
104
+ assert txn_counter == 2
105
+
106
+ # Make sure no transactions are left open
107
+ with postgres_db_engine.begin() as conn:
108
+ result = conn.execute(
109
+ sa.text(
110
+ "select * from pg_stat_activity where state = 'idle in transaction'"
111
+ )
112
+ ).fetchall()
113
+ assert len(result) == 0
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes