dbos 1.4.1__tar.gz → 1.5.0a2__tar.gz

{dbos-1.4.1 → dbos-1.5.0a2}/PKG-INFO

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

{dbos-1.4.1 → dbos-1.5.0a2}/dbos/_admin_server.py

@@ -7,6 +7,8 @@ from functools import partial
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from typing import TYPE_CHECKING, Any, List, Optional, TypedDict
 
+from dbos._workflow_commands import garbage_collect, global_timeout
+
 from ._context import SetWorkflowID
 from ._error import DBOSException
 from ._logger import dbos_logger
@@ -20,6 +22,8 @@ _health_check_path = "/dbos-healthz"
 _workflow_recovery_path = "/dbos-workflow-recovery"
 _deactivate_path = "/deactivate"
 _workflow_queues_metadata_path = "/dbos-workflow-queues-metadata"
+_garbage_collect_path = "/dbos-garbage-collect"
+_global_timeout_path = "/dbos-global-timeout"
 # /workflows/:workflow_id/cancel
 # /workflows/:workflow_id/resume
 # /workflows/:workflow_id/restart
@@ -122,6 +126,23 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
             self.send_response(200)
             self._end_headers()
             self.wfile.write(json.dumps(workflow_ids).encode("utf-8"))
+        elif self.path == _garbage_collect_path:
+            inputs = json.loads(post_data.decode("utf-8"))
+            cutoff_epoch_timestamp_ms = inputs.get("cutoff_epoch_timestamp_ms", None)
+            rows_threshold = inputs.get("rows_threshold", None)
+            garbage_collect(
+                self.dbos,
+                cutoff_epoch_timestamp_ms=cutoff_epoch_timestamp_ms,
+                rows_threshold=rows_threshold,
+            )
+            self.send_response(204)
+            self._end_headers()
+        elif self.path == _global_timeout_path:
+            inputs = json.loads(post_data.decode("utf-8"))
+            timeout_ms = inputs.get("timeout_ms", None)
+            global_timeout(self.dbos, timeout_ms)
+            self.send_response(204)
+            self._end_headers()
         else:
 
             restart_match = re.match(

{dbos-1.4.1 → dbos-1.5.0a2}/dbos/_app_db.py

@@ -256,3 +256,21 @@ class ApplicationDatabase:
             )
 
             conn.execute(insert_stmt)
+
+    def garbage_collect(
+        self, cutoff_epoch_timestamp_ms: int, pending_workflow_ids: list[str]
+    ) -> None:
+        with self.engine.begin() as c:
+            delete_query = sa.delete(ApplicationSchema.transaction_outputs).where(
+                ApplicationSchema.transaction_outputs.c.created_at
+                < cutoff_epoch_timestamp_ms
+            )
+
+            if len(pending_workflow_ids) > 0:
+                delete_query = delete_query.where(
+                    ~ApplicationSchema.transaction_outputs.c.workflow_uuid.in_(
+                        pending_workflow_ids
+                    )
+                )
+
+            c.execute(delete_query)

{dbos-1.4.1 → dbos-1.5.0a2}/dbos/_sys_db.py

@@ -1852,6 +1852,62 @@ class SystemDatabase:
             dbos_logger.error(f"Error connecting to the DBOS system database: {e}")
             raise
 
+    def garbage_collect(
+        self, cutoff_epoch_timestamp_ms: Optional[int], rows_threshold: Optional[int]
+    ) -> Optional[tuple[int, list[str]]]:
+        if rows_threshold is not None:
+            with self.engine.begin() as c:
+                # Get the created_at timestamp of the rows_threshold newest row
+                result = c.execute(
+                    sa.select(SystemSchema.workflow_status.c.created_at)
+                    .order_by(SystemSchema.workflow_status.c.created_at.desc())
+                    .limit(1)
+                    .offset(rows_threshold - 1)
+                ).fetchone()
+
+                if result is not None:
+                    rows_based_cutoff = result[0]
+                    # Use the more restrictive cutoff (higher timestamp = more recent = more deletion)
+                    if (
+                        cutoff_epoch_timestamp_ms is None
+                        or rows_based_cutoff > cutoff_epoch_timestamp_ms
+                    ):
+                        cutoff_epoch_timestamp_ms = rows_based_cutoff
+
+        if cutoff_epoch_timestamp_ms is None:
+            return None
+
+        with self.engine.begin() as c:
+            # Delete all workflows older than cutoff that are NOT PENDING or ENQUEUED
+            c.execute(
+                sa.delete(SystemSchema.workflow_status)
+                .where(
+                    SystemSchema.workflow_status.c.created_at
+                    < cutoff_epoch_timestamp_ms
+                )
+                .where(
+                    ~SystemSchema.workflow_status.c.status.in_(
+                        [
+                            WorkflowStatusString.PENDING.value,
+                            WorkflowStatusString.ENQUEUED.value,
+                        ]
+                    )
+                )
+            )
+
+            # Then, get the IDs of all remaining old workflows
+            pending_enqueued_result = c.execute(
+                sa.select(SystemSchema.workflow_status.c.workflow_uuid).where(
+                    SystemSchema.workflow_status.c.created_at
+                    < cutoff_epoch_timestamp_ms
+                )
+            ).fetchall()
+
+            # Return the final cutoff and workflow IDs
+            return cutoff_epoch_timestamp_ms, [
+                row[0] for row in pending_enqueued_result
+            ]
+
 
 def reset_system_database(postgres_db_url: sa.URL, sysdb_name: str) -> None:
     try:

{dbos-1.4.1 → dbos-1.5.0a2}/dbos/_workflow_commands.py

@@ -1,8 +1,9 @@
+import time
 import uuid
-from typing import List, Optional
+from datetime import datetime
+from typing import TYPE_CHECKING, List, Optional
 
 from dbos._context import get_local_dbos_context
-from dbos._error import DBOSException
 
 from ._app_db import ApplicationDatabase
 from ._sys_db import (
@@ -11,8 +12,12 @@ from ._sys_db import (
     StepInfo,
     SystemDatabase,
     WorkflowStatus,
+    WorkflowStatusString,
 )
 
+if TYPE_CHECKING:
+    from ._dbos import DBOS
+
 
 def list_workflows(
     sys_db: SystemDatabase,
@@ -118,3 +123,32 @@ def fork_workflow(
         application_version=application_version,
     )
     return forked_workflow_id
+
+
+def garbage_collect(
+    dbos: "DBOS",
+    cutoff_epoch_timestamp_ms: Optional[int],
+    rows_threshold: Optional[int],
+) -> None:
+    if cutoff_epoch_timestamp_ms is None and rows_threshold is None:
+        return
+    result = dbos._sys_db.garbage_collect(
+        cutoff_epoch_timestamp_ms=cutoff_epoch_timestamp_ms,
+        rows_threshold=rows_threshold,
+    )
+    if result is not None:
+        cutoff_epoch_timestamp_ms, pending_workflow_ids = result
+        dbos._app_db.garbage_collect(cutoff_epoch_timestamp_ms, pending_workflow_ids)
+
+
+def global_timeout(dbos: "DBOS", timeout_ms: int) -> None:
+    cutoff_epoch_timestamp_ms = int(time.time() * 1000) - timeout_ms
+    cutoff_iso = datetime.fromtimestamp(cutoff_epoch_timestamp_ms / 1000).isoformat()
+    for workflow in dbos.list_workflows(
+        status=WorkflowStatusString.PENDING.value, end_time=cutoff_iso
+    ):
+        dbos.cancel_workflow(workflow.workflow_id)
+    for workflow in dbos.list_workflows(
+        status=WorkflowStatusString.ENQUEUED.value, end_time=cutoff_iso
+    ):
+        dbos.cancel_workflow(workflow.workflow_id)

{dbos-1.4.1 → dbos-1.5.0a2}/pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "1.4.1"
+version = "1.5.0a2"
 
 [project.license]
 text = "MIT"

{dbos-1.4.1 → dbos-1.5.0a2}/tests/test_admin_server.py

@@ -453,3 +453,43 @@ def test_admin_workflow_fork(dbos: DBOS, sys_db: SystemDatabase) -> None:
     ), f"Expected application version to be {new_version}, but got {handle.get_status().app_version}"
 
     assert worked, "Workflow did not finish successfully"
+
+
+def test_admin_garbage_collect(dbos: DBOS) -> None:
+
+    @DBOS.workflow()
+    def workflow() -> str:
+        return DBOS.workflow_id
+
+    workflow()
+
+    assert len(DBOS.list_workflows()) == 1
+
+    response = requests.post(
+        f"http://localhost:3001/dbos-garbage-collect",
+        json={"cutoff_epoch_timestamp_ms": int(time.time() * 1000)},
+        timeout=5,
+    )
+    response.raise_for_status()
+
+    assert len(DBOS.list_workflows()) == 0
+
+
+def test_admin_global_timeout(dbos: DBOS) -> None:
+
+    @DBOS.workflow()
+    def workflow() -> None:
+        while True:
+            DBOS.sleep(0.1)
+
+    handle = DBOS.start_workflow(workflow)
+    time.sleep(1)
+
+    response = requests.post(
+        f"http://localhost:3001/dbos-global-timeout",
+        json={"timeout_ms": 1000},
+        timeout=5,
+    )
+    response.raise_for_status()
+    with pytest.raises(DBOSWorkflowCancelledError):
+        handle.get_result()

{dbos-1.4.1 → dbos-1.5.0a2}/tests/test_workflow_management.py

@@ -1,14 +1,18 @@
 import threading
+import time
 import uuid
 from typing import Callable
 
 import pytest
+import sqlalchemy as sa
 
 # Public API
 from dbos import DBOS, Queue, SetWorkflowID
 from dbos._dbos import DBOSConfiguredInstance
-from dbos._error import DBOSException, DBOSWorkflowCancelledError
+from dbos._error import DBOSWorkflowCancelledError
+from dbos._schemas.application_database import ApplicationSchema
 from dbos._utils import INTERNAL_QUEUE_NAME, GlobalParams
+from dbos._workflow_commands import garbage_collect, global_timeout
 from tests.conftest import queue_entries_are_cleaned_up
 
 
@@ -624,3 +628,120 @@ def test_fork_version(
     GlobalParams.app_version = new_version
     assert handle.get_result() == output
     assert queue_entries_are_cleaned_up(dbos)
+
+
+def test_garbage_collection(dbos: DBOS) -> None:
+    event = threading.Event()
+
+    @DBOS.step()
+    def step(x: int) -> int:
+        return x
+
+    @DBOS.transaction()
+    def txn(x: int) -> int:
+        DBOS.sql_session.execute(sa.text("SELECT 1")).fetchall()
+        return x
+
+    @DBOS.workflow()
+    def workflow(x: int) -> int:
+        step(x)
+        txn(x)
+        return x
+
+    @DBOS.workflow()
+    def blocked_workflow() -> str:
+        txn(0)
+        event.wait()
+        return DBOS.workflow_id
+
+    num_workflows = 10
+
+    handle = DBOS.start_workflow(blocked_workflow)
+    for i in range(num_workflows):
+        assert workflow(i) == i
+
+    # Garbage collect all but one workflow
+    garbage_collect(dbos, cutoff_epoch_timestamp_ms=None, rows_threshold=1)
+    # Verify two workflows remain: the newest and the blocked workflow
+    workflows = DBOS.list_workflows()
+    assert len(workflows) == 2
+    assert workflows[0].workflow_id == handle.workflow_id
+    # Verify txn outputs are preserved only for the remaining workflows
+    with dbos._app_db.engine.begin() as c:
+        rows = c.execute(
+            sa.select(
+                ApplicationSchema.transaction_outputs.c.workflow_uuid,
+            )
+        ).all()
+        assert len(rows) == 2
+
+    # Garbage collect all previous workflows
+    garbage_collect(
+        dbos, cutoff_epoch_timestamp_ms=int(time.time() * 1000), rows_threshold=None
+    )
+    # Verify only the blocked workflow remains
+    workflows = DBOS.list_workflows()
+    assert len(workflows) == 1
+    assert workflows[0].workflow_id == handle.workflow_id
+    # Verify txn outputs are preserved only for the remaining workflow
+    with dbos._app_db.engine.begin() as c:
+        rows = c.execute(
+            sa.select(
+                ApplicationSchema.transaction_outputs.c.workflow_uuid,
+            )
+        ).all()
+        assert len(rows) == 1
+
+    # Finish the blocked workflow, garbage collect everything
+    event.set()
+    assert handle.get_result() is not None
+    garbage_collect(
+        dbos, cutoff_epoch_timestamp_ms=int(time.time() * 1000), rows_threshold=None
+    )
+    # Verify only the blocked workflow remains
+    workflows = DBOS.list_workflows()
+    assert len(workflows) == 0
+
+    # Verify GC runs without error on a blank table
+    garbage_collect(dbos, cutoff_epoch_timestamp_ms=None, rows_threshold=1)
+
+    # Run workflows, wait, run them again
+    for i in range(num_workflows):
+        assert workflow(i) == i
+    time.sleep(1)
+    for i in range(num_workflows):
+        assert workflow(i) == i
+
+    # GC the first half, verify only half were GC'ed
+    garbage_collect(
+        dbos,
+        cutoff_epoch_timestamp_ms=int(time.time() * 1000) - 1000,
+        rows_threshold=None,
+    )
+    workflows = DBOS.list_workflows()
+    assert len(workflows) == num_workflows
+
+
+def test_global_timeout(dbos: DBOS) -> None:
+    event = threading.Event()
+
+    @DBOS.workflow()
+    def blocked_workflow() -> str:
+        while not event.wait(0):
+            DBOS.sleep(0.1)
+        return DBOS.workflow_id
+
+    num_workflows = 10
+    handles = [DBOS.start_workflow(blocked_workflow) for _ in range(num_workflows)]
+
+    # Wait one second, start one final workflow, then timeout all workflows started more than one second ago
+    time.sleep(1)
+    final_handle = DBOS.start_workflow(blocked_workflow)
+    global_timeout(dbos, 1000)
+
+    # Verify all workflows started before the global timeout are cancelled
+    for handle in handles:
+        with pytest.raises(DBOSWorkflowCancelledError):
+            handle.get_result()
+    event.set()
+    final_handle.get_result() is not None