dbos 1.4.0a1__tar.gz → 1.5.0__tar.gz

{dbos-1.4.0a1 → dbos-1.5.0}/PKG-INFO

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

{dbos-1.4.0a1 → dbos-1.5.0}/dbos/_admin_server.py

@@ -5,7 +5,9 @@ import re
 import threading
 from functools import partial
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
-from typing import TYPE_CHECKING, Any, List, Optional, TypedDict
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, TypedDict
+
+from dbos._workflow_commands import garbage_collect, global_timeout
 
 from ._context import SetWorkflowID
 from ._error import DBOSException
@@ -20,6 +22,10 @@ _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"
+_queued_workflows_path = "/queues"
+_workflows_path = "/workflows"
 # /workflows/:workflow_id/cancel
 # /workflows/:workflow_id/resume
 # /workflows/:workflow_id/restart
@@ -100,10 +106,24 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
             steps_match = re.match(
                 r"^/workflows/(?P<workflow_id>[^/]+)/steps$", self.path
             )
+            workflow_match = re.match(r"^/workflows/(?P<workflow_id>[^/]+)$", self.path)
 
             if steps_match:
                 workflow_id = steps_match.group("workflow_id")
                 self._handle_steps(workflow_id)
+            elif workflow_match:
+                workflow_id = workflow_match.group("workflow_id")
+                workflows = self.dbos.list_workflows(workflow_ids=[workflow_id])
+                if not workflows:
+                    self.send_response(404)
+                    self._end_headers()
+                    return
+                response_body = json.dumps(workflows[0].__dict__).encode("utf-8")
+                self.send_response(200)
+                self.send_header("Content-Type", "application/json")
+                self.send_header("Content-Length", str(len(response_body)))
+                self._end_headers()
+                self.wfile.write(response_body)
             else:
                 self.send_response(404)
                 self._end_headers()
@@ -122,8 +142,50 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
             self.send_response(200)
             self._end_headers()
             self.wfile.write(json.dumps(workflow_ids).encode("utf-8"))
+        elif self.path == _workflows_path:
+            try:
+                filters = json.loads(post_data.decode("utf-8")) if post_data else {}
+                self._handle_workflows(filters)
+            except (json.JSONDecodeError, AttributeError) as e:
+                self.send_response(400)
+                self.send_header("Content-Type", "application/json")
+                self.end_headers()
+                self.wfile.write(
+                    json.dumps({"error": f"Invalid JSON input: {str(e)}"}).encode(
+                        "utf-8"
+                    )
+                )
+        elif self.path == _queued_workflows_path:
+            try:
+                filters = json.loads(post_data.decode("utf-8")) if post_data else {}
+                self._handle_queued_workflows(filters)
+            except (json.JSONDecodeError, AttributeError) as e:
+                self.send_response(400)
+                self.send_header("Content-Type", "application/json")
+                self.end_headers()
+                self.wfile.write(
+                    json.dumps({"error": f"Invalid JSON input: {str(e)}"}).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"))
+            cutoff_epoch_timestamp_ms = inputs.get("cutoff_epoch_timestamp_ms", None)
+            global_timeout(self.dbos, cutoff_epoch_timestamp_ms)
+            self.send_response(204)
+            self._end_headers()
         else:
-
             restart_match = re.match(
                 r"^/workflows/(?P<workflow_id>[^/]+)/restart$", self.path
             )
@@ -262,6 +324,50 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
         self._end_headers()
         self.wfile.write(json_steps)
 
+    def _handle_workflows(self, filters: Dict[str, Any]) -> None:
+        workflows = self.dbos.list_workflows(
+            workflow_ids=filters.get("workflow_ids"),
+            name=filters.get("name"),
+            start_time=filters.get("start_time"),
+            end_time=filters.get("end_time"),
+            status=filters.get("status"),
+            app_version=filters.get("application_version"),
+            limit=filters.get("limit"),
+            offset=filters.get("offset"),
+            sort_desc=filters.get("sort_desc", False),
+            workflow_id_prefix=filters.get("workflow_id_prefix"),
+        )
+
+        response_body = json.dumps(
+            [workflow.__dict__ for workflow in workflows]
+        ).encode("utf-8")
+        self.send_response(200)
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Length", str(len(response_body)))
+        self._end_headers()
+        self.wfile.write(response_body)
+
+    def _handle_queued_workflows(self, filters: Dict[str, Any]) -> None:
+        workflows = self.dbos.list_queued_workflows(
+            queue_name=filters.get("queue_name"),
+            name=filters.get("name"),
+            start_time=filters.get("start_time"),
+            end_time=filters.get("end_time"),
+            status=filters.get("status"),
+            limit=filters.get("limit"),
+            offset=filters.get("offset"),
+            sort_desc=filters.get("sort_desc", False),
+        )
+
+        response_body = json.dumps(
+            [workflow.__dict__ for workflow in workflows]
+        ).encode("utf-8")
+        self.send_response(200)
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Length", str(len(response_body)))
+        self._end_headers()
+        self.wfile.write(response_body)
+
 
 # Be consistent with DBOS-TS response.
 class PerfUtilization(TypedDict):

{dbos-1.4.0a1 → dbos-1.5.0}/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.0a1 → dbos-1.5.0}/dbos/_conductor/conductor.py

@@ -13,7 +13,9 @@ from websockets.sync.connection import Connection
 from dbos._context import SetWorkflowID
 from dbos._utils import GlobalParams
 from dbos._workflow_commands import (
+    garbage_collect,
     get_workflow,
+    global_timeout,
     list_queued_workflows,
     list_workflow_steps,
     list_workflows,
@@ -356,6 +358,41 @@ class ConductorWebsocket(threading.Thread):
                                 error_message=error_message,
                             )
                             websocket.send(list_steps_response.to_json())
+                        elif msg_type == p.MessageType.RETENTION:
+                            retention_message = p.RetentionRequest.from_json(message)
+                            success = True
+                            try:
+                                garbage_collect(
+                                    self.dbos,
+                                    cutoff_epoch_timestamp_ms=retention_message.body[
+                                        "gc_cutoff_epoch_ms"
+                                    ],
+                                    rows_threshold=retention_message.body[
+                                        "gc_rows_threshold"
+                                    ],
+                                )
+                                if (
+                                    retention_message.body["timeout_cutoff_epoch_ms"]
+                                    is not None
+                                ):
+                                    global_timeout(
+                                        self.dbos,
+                                        retention_message.body[
+                                            "timeout_cutoff_epoch_ms"
+                                        ],
+                                    )
+                            except Exception as e:
+                                error_message = f"Exception encountered during enforcing retention policy: {traceback.format_exc()}"
+                                self.dbos.logger.error(error_message)
+                                success = False
+
+                            retention_response = p.RetentionResponse(
+                                type=p.MessageType.RETENTION,
+                                request_id=base_message.request_id,
+                                success=success,
+                                error_message=error_message,
+                            )
+                            websocket.send(retention_response.to_json())
                         else:
                             self.dbos.logger.warning(
                                 f"Unexpected message type: {msg_type}"

{dbos-1.4.0a1 → dbos-1.5.0}/dbos/_conductor/protocol.py

@@ -18,6 +18,7 @@ class MessageType(str, Enum):
     EXIST_PENDING_WORKFLOWS = "exist_pending_workflows"
     LIST_STEPS = "list_steps"
     FORK_WORKFLOW = "fork_workflow"
+    RETENTION = "retention"
 
 
 T = TypeVar("T", bound="BaseMessage")
@@ -280,3 +281,20 @@ class ForkWorkflowRequest(BaseMessage):
 class ForkWorkflowResponse(BaseMessage):
     new_workflow_id: Optional[str]
     error_message: Optional[str] = None
+
+
+class RetentionBody(TypedDict):
+    gc_cutoff_epoch_ms: Optional[int]
+    gc_rows_threshold: Optional[int]
+    timeout_cutoff_epoch_ms: Optional[int]
+
+
+@dataclass
+class RetentionRequest(BaseMessage):
+    body: RetentionBody
+
+
+@dataclass
+class RetentionResponse(BaseMessage):
+    success: bool
+    error_message: Optional[str] = None

{dbos-1.4.0a1 → dbos-1.5.0}/dbos/_core.py

@@ -404,9 +404,9 @@ def _execute_workflow_wthread(
                     return dbos._background_event_loop.submit_coroutine(
                         cast(Pending[R], result)()
                     )
-            except Exception:
+            except Exception as e:
                 dbos.logger.error(
-                    f"Exception encountered in asynchronous workflow: {traceback.format_exc()}"
+                    f"Exception encountered in asynchronous workflow:", exc_info=e
                 )
                 raise
 
@@ -430,9 +430,9 @@ async def _execute_workflow_async(
                     _get_wf_invoke_func(dbos, status)
                 )
                 return await result()
-            except Exception:
+            except Exception as e:
                 dbos.logger.error(
-                    f"Exception encountered in asynchronous workflow: {traceback.format_exc()}"
+                    f"Exception encountered in asynchronous workflow:", exc_info=e
                 )
                 raise
 
@@ -1123,7 +1123,7 @@ def decorate_step(
                 stepOutcome = stepOutcome.retry(
                     max_attempts,
                     on_exception,
-                    lambda i: DBOSMaxStepRetriesExceeded(func.__name__, i),
+                    lambda i, e: DBOSMaxStepRetriesExceeded(func.__name__, i, e),
                 )
 
             outcome = (

{dbos-1.4.0a1 → dbos-1.5.0}/dbos/_dbos.py

@@ -521,8 +521,8 @@ class DBOS:
                 handler.flush()
             add_otlp_to_all_loggers()
             add_transformer_to_all_loggers()
-        except Exception:
-            dbos_logger.error(f"DBOS failed to launch: {traceback.format_exc()}")
+        except Exception as e:
+            dbos_logger.error(f"DBOS failed to launch:", exc_info=e)
             raise
 
     @classmethod

{dbos-1.4.0a1 → dbos-1.5.0}/dbos/_dbos_config.py

@@ -31,6 +31,7 @@ class DBOSConfig(TypedDict, total=False):
         otlp_logs_endpoints: List[str]: OTLP logs endpoints
         admin_port (int): Admin port
         run_admin_server (bool): Whether to run the DBOS admin server
+        otlp_attributes (dict[str, str]): A set of custom attributes to apply OTLP-exported logs and traces
     """
 
     name: str
@@ -43,6 +44,7 @@ class DBOSConfig(TypedDict, total=False):
     otlp_logs_endpoints: Optional[List[str]]
     admin_port: Optional[int]
     run_admin_server: Optional[bool]
+    otlp_attributes: Optional[dict[str, str]]
 
 
 class RuntimeConfig(TypedDict, total=False):
@@ -84,6 +86,7 @@ class LoggerConfig(TypedDict, total=False):
 class TelemetryConfig(TypedDict, total=False):
     logs: Optional[LoggerConfig]
     OTLPExporter: Optional[OTLPExporterConfig]
+    otlp_attributes: Optional[dict[str, str]]
 
 
 class ConfigFile(TypedDict, total=False):
@@ -145,7 +148,8 @@ def translate_dbos_config_to_config_file(config: DBOSConfig) -> ConfigFile:
 
     # Telemetry config
     telemetry: TelemetryConfig = {
-        "OTLPExporter": {"tracesEndpoint": [], "logsEndpoint": []}
+        "OTLPExporter": {"tracesEndpoint": [], "logsEndpoint": []},
+        "otlp_attributes": config.get("otlp_attributes", {}),
     }
     # For mypy
     assert telemetry["OTLPExporter"] is not None
@@ -431,7 +435,6 @@ def is_valid_database_url(database_url: str) -> bool:
     url = make_url(database_url)
     required_fields = [
         ("username", "Username must be specified in the connection URL"),
-        ("password", "Password must be specified in the connection URL"),
         ("host", "Host must be specified in the connection URL"),
         ("database", "Database name must be specified in the connection URL"),
     ]

{dbos-1.4.0a1 → dbos-1.5.0}/dbos/_debug.py

@@ -15,11 +15,11 @@ class PythonModule:
 
 
 def debug_workflow(workflow_id: str, entrypoint: Union[str, PythonModule]) -> None:
-    # include the current directory (represented by empty string) in the search path
-    # if it not already included
-    if "" not in sys.path:
-        sys.path.insert(0, "")
     if isinstance(entrypoint, str):
+        # ensure the entrypoint parent directory is in sys.path
+        parent = str(Path(entrypoint).parent)
+        if parent not in sys.path:
+            sys.path.insert(0, parent)
         runpy.run_path(entrypoint)
     elif isinstance(entrypoint, PythonModule):
         runpy.run_module(entrypoint.module_name)

{dbos-1.4.0a1 → dbos-1.5.0}/dbos/_error.py

@@ -150,9 +150,12 @@ class DBOSNotAuthorizedError(DBOSException):
 class DBOSMaxStepRetriesExceeded(DBOSException):
     """Exception raised when a step was retried the maximimum number of times without success."""
 
-    def __init__(self, step_name: str, max_retries: int) -> None:
+    def __init__(
+        self, step_name: str, max_retries: int, errors: list[Exception]
+    ) -> None:
         self.step_name = step_name
         self.max_retries = max_retries
+        self.errors = errors
         super().__init__(
             f"Step {step_name} has exceeded its maximum of {max_retries} retries",
             dbos_error_code=DBOSErrorCode.MaxStepRetriesExceeded.value,
@@ -160,7 +163,7 @@ class DBOSMaxStepRetriesExceeded(DBOSException):
 
     def __reduce__(self) -> Any:
         # Tell jsonpickle how to reconstruct this object
-        return (self.__class__, (self.step_name, self.max_retries))
+        return (self.__class__, (self.step_name, self.max_retries, self.errors))
 
 
 class DBOSConflictingRegistrationError(DBOSException):

{dbos-1.4.0a1 → dbos-1.5.0}/dbos/_logger.py

@@ -20,14 +20,17 @@ _otlp_handler, _dbos_log_transformer = None, None
 
 
 class DBOSLogTransformer(logging.Filter):
-    def __init__(self) -> None:
+    def __init__(self, config: "ConfigFile") -> None:
         super().__init__()
         self.app_id = os.environ.get("DBOS__APPID", "")
+        self.otlp_attributes: dict[str, str] = config.get("telemetry", {}).get("otlp_attributes", {})  # type: ignore
 
     def filter(self, record: Any) -> bool:
         record.applicationID = self.app_id
         record.applicationVersion = GlobalParams.app_version
         record.executorID = GlobalParams.executor_id
+        for k, v in self.otlp_attributes.items():
+            setattr(record, k, v)
 
         # If available, decorate the log entry with Workflow ID and Trace ID
         from dbos._context import get_local_dbos_context
@@ -98,7 +101,7 @@ def config_logger(config: "ConfigFile") -> None:
 
     # Attach DBOS-specific attributes to all log entries.
     global _dbos_log_transformer
-    _dbos_log_transformer = DBOSLogTransformer()
+    _dbos_log_transformer = DBOSLogTransformer(config)
     dbos_logger.addFilter(_dbos_log_transformer)
 
 

{dbos-1.4.0a1 → dbos-1.5.0}/dbos/_outcome.py

@@ -37,7 +37,7 @@ class Outcome(Protocol[T]):
         self,
         attempts: int,
         on_exception: Callable[[int, BaseException], float],
-        exceeded_retries: Callable[[int], BaseException],
+        exceeded_retries: Callable[[int, list[Exception]], Exception],
     ) -> "Outcome[T]": ...
 
     def intercept(
@@ -96,23 +96,25 @@ class Immediate(Outcome[T]):
         func: Callable[[], T],
         attempts: int,
         on_exception: Callable[[int, BaseException], float],
-        exceeded_retries: Callable[[int], BaseException],
+        exceeded_retries: Callable[[int, list[Exception]], Exception],
     ) -> T:
+        errors: list[Exception] = []
         for i in range(attempts):
             try:
                 with EnterDBOSStepRetry(i, attempts):
                     return func()
             except Exception as exp:
+                errors.append(exp)
                 wait_time = on_exception(i, exp)
                 time.sleep(wait_time)
 
-        raise exceeded_retries(attempts)
+        raise exceeded_retries(attempts, errors)
 
     def retry(
         self,
         attempts: int,
         on_exception: Callable[[int, BaseException], float],
-        exceeded_retries: Callable[[int], BaseException],
+        exceeded_retries: Callable[[int, list[Exception]], Exception],
     ) -> "Immediate[T]":
         assert attempts > 0
         return Immediate[T](
@@ -183,23 +185,25 @@ class Pending(Outcome[T]):
         func: Callable[[], Coroutine[Any, Any, T]],
         attempts: int,
         on_exception: Callable[[int, BaseException], float],
-        exceeded_retries: Callable[[int], BaseException],
+        exceeded_retries: Callable[[int, list[Exception]], Exception],
     ) -> T:
+        errors: list[Exception] = []
         for i in range(attempts):
             try:
                 with EnterDBOSStepRetry(i, attempts):
                     return await func()
             except Exception as exp:
+                errors.append(exp)
                 wait_time = on_exception(i, exp)
                 await asyncio.sleep(wait_time)
 
-        raise exceeded_retries(attempts)
+        raise exceeded_retries(attempts, errors)
 
     def retry(
         self,
         attempts: int,
         on_exception: Callable[[int, BaseException], float],
-        exceeded_retries: Callable[[int], BaseException],
+        exceeded_retries: Callable[[int, list[Exception]], Exception],
     ) -> "Pending[T]":
         assert attempts > 0
         return Pending[T](

{dbos-1.4.0a1 → dbos-1.5.0}/dbos/_queue.py

@@ -1,5 +1,4 @@
 import threading
-import traceback
 from typing import TYPE_CHECKING, Any, Callable, Coroutine, Optional, TypedDict
 
 from psycopg import errors

{dbos-1.4.0a1 → dbos-1.5.0}/dbos/_recovery.py

@@ -1,7 +1,5 @@
-import os
 import threading
 import time
-import traceback
 from typing import TYPE_CHECKING, Any, List
 
 from dbos._utils import GlobalParams
@@ -39,9 +37,9 @@ def startup_recovery_thread(
             time.sleep(1)
         except Exception as e:
             dbos.logger.error(
-                f"Exception encountered when recovering workflows: {traceback.format_exc()}"
+                f"Exception encountered when recovering workflows:", exc_info=e
             )
-            raise e
+            raise
 
 
 def recover_pending_workflows(
@@ -59,9 +57,9 @@ def recover_pending_workflows(
                 workflow_handles.append(handle)
             except Exception as e:
                 dbos.logger.error(
-                    f"Exception encountered when recovering workflows: {traceback.format_exc()}"
+                    f"Exception encountered when recovering workflows:", exc_info=e
                 )
-                raise e
+                raise
         dbos.logger.info(
             f"Recovering {len(pending_workflows)} workflows for executor {executor_id} from version {GlobalParams.app_version}"
         )

{dbos-1.4.0a1 → dbos-1.5.0}/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.0a1 → dbos-1.5.0}/dbos/_tracer.py

@@ -19,11 +19,14 @@ if TYPE_CHECKING:
 
 class DBOSTracer:
 
+    otlp_attributes: dict[str, str] = {}
+
     def __init__(self) -> None:
         self.app_id = os.environ.get("DBOS__APPID", None)
         self.provider: Optional[TracerProvider] = None
 
     def config(self, config: ConfigFile) -> None:
+        self.otlp_attributes = config.get("telemetry", {}).get("otlp_attributes", {})  # type: ignore
         if not isinstance(trace.get_tracer_provider(), TracerProvider):
             resource = Resource(
                 attributes={
@@ -63,6 +66,8 @@ class DBOSTracer:
         for k, v in attributes.items():
             if k != "name" and v is not None and isinstance(v, (str, bool, int, float)):
                 span.set_attribute(k, v)
+        for k, v in self.otlp_attributes.items():
+            span.set_attribute(k, v)
         return span
 
     def end_span(self, span: Span) -> None:

{dbos-1.4.0a1 → dbos-1.5.0}/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,31 @@ 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", cutoff_epoch_timestamp_ms: int) -> None:
+    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)