dbos 0.11.0a3__py3-none-any.whl → 0.12.0__py3-none-any.whl

dbos/__init__.py

@@ -1,10 +1,10 @@
-from . import error as error
-from .context import DBOSContextEnsure, DBOSContextSetAuth, SetWorkflowID
-from .dbos import DBOS, DBOSConfiguredInstance, WorkflowHandle, WorkflowStatus
-from .dbos_config import ConfigFile, get_dbos_database_url, load_config
-from .kafka_message import KafkaMessage
-from .queue import Queue
-from .system_database import GetWorkflowsInput, WorkflowStatusString
+from . import _error as error
+from ._context import DBOSContextEnsure, DBOSContextSetAuth, SetWorkflowID
+from ._dbos import DBOS, DBOSConfiguredInstance, WorkflowHandle, WorkflowStatus
+from ._dbos_config import ConfigFile, get_dbos_database_url, load_config
+from ._kafka_message import KafkaMessage
+from ._queue import Queue
+from ._sys_db import GetWorkflowsInput, WorkflowStatusString
 
 __all__ = [
     "ConfigFile",

dbos/{admin_sever.py → _admin_sever.py}

@@ -8,16 +8,16 @@ from typing import TYPE_CHECKING, Any, List, TypedDict
 
 import psutil
 
-from dbos.recovery import _recover_pending_workflows
-
-from .logger import dbos_logger
+from ._logger import dbos_logger
+from ._recovery import recover_pending_workflows
 
 if TYPE_CHECKING:
-    from .dbos import DBOS
+    from ._dbos import DBOS
 
-health_check_path = "/dbos-healthz"
-workflow_recovery_path = "/dbos-workflow-recovery"
-perf_path = "/dbos-perf"
+_health_check_path = "/dbos-healthz"
+_workflow_recovery_path = "/dbos-workflow-recovery"
+_perf_path = "/dbos-perf"
+_deactivate_path = "/deactivate"
 
 
 class AdminServer:
@@ -51,11 +51,11 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
         self._end_headers()
 
     def do_GET(self) -> None:
-        if self.path == health_check_path:
+        if self.path == _health_check_path:
             self.send_response(200)
             self._end_headers()
             self.wfile.write("healthy".encode("utf-8"))
-        elif self.path == perf_path:
+        elif self.path == _perf_path:
             # Compares system CPU times elapsed since last call or module import, returning immediately (non blocking).
             cpu_percent = psutil.cpu_percent(interval=None) / 100.0
             perf_util: PerfUtilization = {
@@ -66,6 +66,15 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
             self.send_response(200)
             self._end_headers()
             self.wfile.write(json.dumps(perf_util).encode("utf-8"))
+        elif self.path == _deactivate_path:
+            dbos_logger.info("Deactivating DBOS")
+            # Stop all scheduled workflows, queues, and kafka loops
+            for event in self.dbos.stop_events:
+                event.set()
+
+            self.send_response(200)
+            self._end_headers()
+            self.wfile.write("deactivated".encode("utf-8"))
         else:
             self.send_response(404)
             self._end_headers()
@@ -76,10 +85,10 @@ class AdminRequestHandler(BaseHTTPRequestHandler):
         )  # <--- Gets the size of data
         post_data = self.rfile.read(content_length)  # <--- Gets the data itself
 
-        if self.path == workflow_recovery_path:
+        if self.path == _workflow_recovery_path:
             executor_ids: List[str] = json.loads(post_data.decode("utf-8"))
             dbos_logger.info("Recovering workflows for executors: %s", executor_ids)
-            workflow_handles = _recover_pending_workflows(self.dbos, executor_ids)
+            workflow_handles = recover_pending_workflows(self.dbos, executor_ids)
             workflow_ids = [handle.workflow_id for handle in workflow_handles]
             self.send_response(200)
             self._end_headers()

dbos/{application_database.py → _app_db.py}

@@ -1,14 +1,13 @@
-from typing import Optional, TypedDict, cast
+from typing import Optional, TypedDict
 
 import sqlalchemy as sa
 import sqlalchemy.dialects.postgresql as pg
 from sqlalchemy.exc import DBAPIError
 from sqlalchemy.orm import Session, sessionmaker
 
-from dbos.error import DBOSWorkflowConflictIDError
-from dbos.schemas.application_database import ApplicationSchema
-
-from .dbos_config import ConfigFile
+from ._dbos_config import ConfigFile
+from ._error import DBOSWorkflowConflictIDError
+from ._schemas.application_database import ApplicationSchema
 
 
 class TransactionResultInternal(TypedDict):

dbos/{decorators.py → _classproperty.py}

@@ -3,7 +3,7 @@ from typing import Any, Callable, Generic, Optional, TypeVar
 G = TypeVar("G")  # A generic type for ClassPropertyDescriptor getters
 
 
-class _ClassPropertyDescriptor(Generic[G]):
+class ClassPropertyDescriptor(Generic[G]):
     def __init__(self, fget: Callable[..., G]) -> None:
         self.fget = fget
 
@@ -15,5 +15,5 @@ class _ClassPropertyDescriptor(Generic[G]):
         return self.fget(objtype)
 
 
-def classproperty(func: Callable[..., G]) -> _ClassPropertyDescriptor[G]:
-    return _ClassPropertyDescriptor(func)
+def classproperty(func: Callable[..., G]) -> ClassPropertyDescriptor[G]:
+    return ClassPropertyDescriptor(func)

dbos/{context.py → _context.py}

@@ -7,14 +7,14 @@ from contextlib import AbstractContextManager
 from contextvars import ContextVar
 from enum import Enum
 from types import TracebackType
-from typing import TYPE_CHECKING, List, Literal, Optional, Type, TypedDict
+from typing import List, Literal, Optional, Type, TypedDict
 
 from opentelemetry.trace import Span, Status, StatusCode
 from sqlalchemy.orm import Session
 
-from .logger import dbos_logger
-from .request import Request
-from .tracer import dbos_tracer
+from ._logger import dbos_logger
+from ._request import Request
+from ._tracer import dbos_tracer
 
 
 # These are used to tag OTel traces
@@ -202,21 +202,21 @@ class DBOSContext:
 ##############################################################
 
 
-dbos_context_var: ContextVar[Optional[DBOSContext]] = ContextVar(
+_dbos_context_var: ContextVar[Optional[DBOSContext]] = ContextVar(
     "dbos_context", default=None
 )
 
 
-def set_local_dbos_context(ctx: Optional[DBOSContext]) -> None:
-    dbos_context_var.set(ctx)
+def _set_local_dbos_context(ctx: Optional[DBOSContext]) -> None:
+    _dbos_context_var.set(ctx)
 
 
-def clear_local_dbos_context() -> None:
-    dbos_context_var.set(None)
+def _clear_local_dbos_context() -> None:
+    _dbos_context_var.set(None)
 
 
 def get_local_dbos_context() -> Optional[DBOSContext]:
-    return dbos_context_var.get()
+    return _dbos_context_var.get()
 
 
 def assert_current_dbos_context() -> DBOSContext:
@@ -239,7 +239,7 @@ class DBOSContextEnsure:
         ctx = get_local_dbos_context()
         if ctx is None:
             self.created_ctx = True
-            set_local_dbos_context(DBOSContext())
+            _set_local_dbos_context(DBOSContext())
         return assert_current_dbos_context()
 
     def __exit__(
@@ -250,7 +250,7 @@ class DBOSContextEnsure:
     ) -> Literal[False]:
         # Code to clean up the basic context if we created it
         if self.created_ctx:
-            clear_local_dbos_context()
+            _clear_local_dbos_context()
         return False  # Did not handle
 
 
@@ -261,7 +261,7 @@ class DBOSContextSwap:
 
     def __enter__(self) -> DBOSContextSwap:
         self.prev_ctx = get_local_dbos_context()
-        set_local_dbos_context(self.next_ctx)
+        _set_local_dbos_context(self.next_ctx)
         return self
 
     def __exit__(
@@ -271,7 +271,7 @@ class DBOSContextSwap:
         traceback: Optional[TracebackType],
     ) -> Literal[False]:
         assert get_local_dbos_context() == self.next_ctx
-        set_local_dbos_context(self.prev_ctx)
+        _set_local_dbos_context(self.prev_ctx)
         return False  # Did not handle
 
 
@@ -301,7 +301,7 @@ class SetWorkflowID:
         ctx = get_local_dbos_context()
         if ctx is None:
             self.created_ctx = True
-            set_local_dbos_context(DBOSContext())
+            _set_local_dbos_context(DBOSContext())
         assert_current_dbos_context().id_assigned_for_next_workflow = self.wfid
         return self
 
@@ -313,7 +313,7 @@ class SetWorkflowID:
     ) -> Literal[False]:
         # Code to clean up the basic context if we created it
         if self.created_ctx:
-            clear_local_dbos_context()
+            _clear_local_dbos_context()
         return False  # Did not handle
 
 
@@ -326,7 +326,7 @@ class SetWorkflowRecovery:
         ctx = get_local_dbos_context()
         if ctx is None:
             self.created_ctx = True
-            set_local_dbos_context(DBOSContext())
+            _set_local_dbos_context(DBOSContext())
         assert_current_dbos_context().in_recovery = True
 
         return self
@@ -341,7 +341,7 @@ class SetWorkflowRecovery:
         assert_current_dbos_context().in_recovery = False
         # Code to clean up the basic context if we created it
         if self.created_ctx:
-            clear_local_dbos_context()
+            _clear_local_dbos_context()
         return False  # Did not handle
 
 
@@ -356,7 +356,7 @@ class EnterDBOSWorkflow(AbstractContextManager[DBOSContext, Literal[False]]):
         if ctx is None:
             self.created_ctx = True
             ctx = DBOSContext()
-            set_local_dbos_context(ctx)
+            _set_local_dbos_context(ctx)
         assert not ctx.is_within_workflow()
         ctx.start_workflow(
             None, self.attributes
@@ -374,7 +374,7 @@ class EnterDBOSWorkflow(AbstractContextManager[DBOSContext, Literal[False]]):
         ctx.end_workflow(exc_value)
         # Code to clean up the basic context if we created it
         if self.created_ctx:
-            clear_local_dbos_context()
+            _clear_local_dbos_context()
         return False  # Did not handle
 
 
@@ -394,7 +394,7 @@ class EnterDBOSChildWorkflow(AbstractContextManager[DBOSContext, Literal[False]]
                 ctx.workflow_id + "-" + str(ctx.function_id)
             )
         self.child_ctx = ctx.create_child()
-        set_local_dbos_context(self.child_ctx)
+        _set_local_dbos_context(self.child_ctx)
         self.child_ctx.start_workflow(None, attributes=self.attributes)
         return self.child_ctx
 
@@ -409,7 +409,7 @@ class EnterDBOSChildWorkflow(AbstractContextManager[DBOSContext, Literal[False]]
         ctx.end_workflow(exc_value)
         # Return to parent ctx
         assert self.parent_ctx
-        set_local_dbos_context(self.parent_ctx)
+        _set_local_dbos_context(self.parent_ctx)
         return False  # Did not handle
 
 
@@ -473,7 +473,7 @@ class EnterDBOSHandler:
         ctx = get_local_dbos_context()
         if ctx is None:
             self.created_ctx = True
-            set_local_dbos_context(DBOSContext())
+            _set_local_dbos_context(DBOSContext())
         ctx = assert_current_dbos_context()
         ctx.start_handler(self.attributes)
         return self
@@ -488,7 +488,7 @@ class EnterDBOSHandler:
         ctx.end_handler(exc_value)
         # Code to clean up the basic context if we created it
         if self.created_ctx:
-            clear_local_dbos_context()
+            _clear_local_dbos_context()
         return False  # Did not handle
 
 
@@ -504,7 +504,7 @@ class DBOSContextSetAuth(DBOSContextEnsure):
         ctx = get_local_dbos_context()
         if ctx is None:
             self.created_ctx = True
-            set_local_dbos_context(DBOSContext())
+            _set_local_dbos_context(DBOSContext())
         ctx = assert_current_dbos_context()
         self.prev_user = ctx.authenticated_user
         self.prev_roles = ctx.authenticated_roles
@@ -521,7 +521,7 @@ class DBOSContextSetAuth(DBOSContextEnsure):
         ctx.set_authentication(self.prev_user, self.prev_roles)
         # Clean up the basic context if we created it
         if self.created_ctx:
-            clear_local_dbos_context()
+            _clear_local_dbos_context()
         return False  # Did not handle