dbos 0.8.0a0__py3-none-any.whl → 0.8.0a3__py3-none-any.whl

dbos/__init__.py

@@ -1,5 +1,5 @@
 from . import error as error
-from .context import DBOSContextEnsure, SetWorkflowID
+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
@@ -11,6 +11,7 @@ __all__ = [
     "DBOS",
     "DBOSConfiguredInstance",
     "DBOSContextEnsure",
+    "DBOSContextSetAuth",
     "GetWorkflowsInput",
     "KafkaMessage",
     "SetWorkflowID",

dbos/context.py

@@ -492,6 +492,39 @@ class EnterDBOSHandler:
         return False  # Did not handle
 
 
+class DBOSContextSetAuth(DBOSContextEnsure):
+    def __init__(self, user: Optional[str], roles: Optional[List[str]]) -> None:
+        self.created_ctx = False
+        self.user = user
+        self.roles = roles
+        self.prev_user: Optional[str] = None
+        self.prev_roles: Optional[List[str]] = None
+
+    def __enter__(self) -> DBOSContext:
+        ctx = get_local_dbos_context()
+        if ctx is None:
+            self.created_ctx = True
+            set_local_dbos_context(DBOSContext())
+        ctx = assert_current_dbos_context()
+        self.prev_user = ctx.authenticated_user
+        self.prev_roles = ctx.authenticated_roles
+        ctx.set_authentication(self.user, self.roles)
+        return ctx
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> Literal[False]:
+        ctx = assert_current_dbos_context()
+        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()
+        return False  # Did not handle
+
+
 class DBOSAssumeRole:
     def __init__(self, assume_role: Optional[str]) -> None:
         self.prior_role: Optional[str] = None

dbos/core.py

@@ -4,17 +4,7 @@ import time
 import traceback
 from concurrent.futures import Future
 from functools import wraps
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Generic,
-    List,
-    Optional,
-    Tuple,
-    TypeVar,
-    cast,
-)
+from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, Tuple, TypeVar, cast
 
 from dbos.application_database import ApplicationDatabase, TransactionResultInternal
 
@@ -61,10 +51,10 @@ from dbos.roles import check_required_roles
 from dbos.system_database import (
     GetEventWorkflowContext,
     OperationResultInternal,
-    WorkflowInputs,
     WorkflowStatusInternal,
     WorkflowStatusString,
 )
+from dbos.utils import WorkflowInputs
 
 if TYPE_CHECKING:
     from dbos.dbos import DBOS, Workflow, WorkflowHandle, WorkflowStatus, _DBOSRegistry
@@ -168,10 +158,10 @@ def _init_workflow(
         # We also have to do this for single-step workflows because of the foreign key constraint on the operation outputs table
         # TODO: Make this transactional (and with the queue step below)
         dbos._sys_db.update_workflow_status(status, False, ctx.in_recovery)
-        dbos._sys_db.update_workflow_inputs(wfid, utils.serialize(inputs))
+        dbos._sys_db.update_workflow_inputs(wfid, utils.serialize_args(inputs))
     else:
         # Buffer the inputs for single-transaction workflows, but don't buffer the status
-        dbos._sys_db.buffer_workflow_inputs(wfid, utils.serialize(inputs))
+        dbos._sys_db.buffer_workflow_inputs(wfid, utils.serialize_args(inputs))
 
     if queue is not None:
         dbos._sys_db.enqueue(wfid, queue)
@@ -203,7 +193,7 @@ def _execute_workflow(
         return output
     except Exception as error:
         status["status"] = "ERROR"
-        status["error"] = utils.serialize(error)
+        status["error"] = utils.serialize_exception(error)
         if status["queue_name"] is not None:
             dbos._sys_db.remove_from_queue(status["workflow_uuid"])
         dbos._sys_db.update_workflow_status(status)
@@ -488,8 +478,10 @@ def _transaction(
                                 )
                                 if recorded_output:
                                     if recorded_output["error"]:
-                                        deserialized_error = utils.deserialize(
-                                            recorded_output["error"]
+                                        deserialized_error = (
+                                            utils.deserialize_exception(
+                                                recorded_output["error"]
+                                            )
                                         )
                                         has_recorded_error = True
                                         raise deserialized_error
@@ -527,7 +519,7 @@ def _transaction(
                         except Exception as error:
                             # Don't record the error if it was already recorded
                             if not has_recorded_error:
-                                txn_output["error"] = utils.serialize(error)
+                                txn_output["error"] = utils.serialize_exception(error)
                                 dbos._app_db.record_transaction_error(txn_output)
                             raise
             return output
@@ -599,7 +591,9 @@ def _step(
                 )
                 if recorded_output:
                     if recorded_output["error"] is not None:
-                        deserialized_error = utils.deserialize(recorded_output["error"])
+                        deserialized_error = utils.deserialize_exception(
+                            recorded_output["error"]
+                        )
                         raise deserialized_error
                     elif recorded_output["output"] is not None:
                         return utils.deserialize(recorded_output["output"])
@@ -639,7 +633,7 @@ def _step(
                                 )
 
                 step_output["error"] = (
-                    utils.serialize(error) if error is not None else None
+                    utils.serialize_exception(error) if error is not None else None
                 )
                 dbos._sys_db.record_operation_result(step_output)
 

dbos/migrations/versions/50f3227f0b4b_fix_job_queue.py

@@ -0,0 +1,34 @@
+"""fix_job_queue
+
+Revision ID: 50f3227f0b4b
+Revises: eab0cc1d9a14
+Create Date: 2024-09-25 14:03:53.308068
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "50f3227f0b4b"
+down_revision: Union[str, None] = "eab0cc1d9a14"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.drop_constraint("job_queue_pkey", "job_queue", schema="dbos", type_="primary")
+
+    op.create_primary_key(
+        "job_queue_pkey", "job_queue", ["workflow_uuid"], schema="dbos"
+    )
+
+
+def downgrade() -> None:
+    # Reverting the changes
+    op.drop_constraint("job_queue_pkey", "job_queue", schema="dbos", type_="primary")
+
+    op.create_primary_key(
+        "job_queue_pkey", "job_queue", ["created_at_epoch_ms"], schema="dbos"
+    )

dbos/queue.py

@@ -1,9 +1,9 @@
 import threading
 import time
+import traceback
 from typing import TYPE_CHECKING, Optional
 
 from dbos.core import P, R, _execute_workflow_id, _start_workflow
-from dbos.error import DBOSInitializationError
 
 if TYPE_CHECKING:
     from dbos.dbos import DBOS, Workflow, WorkflowHandle
@@ -31,6 +31,13 @@ def queue_thread(stop_event: threading.Event, dbos: "DBOS") -> None:
     while not stop_event.is_set():
         time.sleep(1)
         for queue_name, queue in dbos._registry.queue_info_map.items():
-            wf_ids = dbos._sys_db.start_queued_workflows(queue_name, queue.concurrency)
-            for id in wf_ids:
-                _execute_workflow_id(dbos, id)
+            try:
+                wf_ids = dbos._sys_db.start_queued_workflows(
+                    queue_name, queue.concurrency
+                )
+                for id in wf_ids:
+                    _execute_workflow_id(dbos, id)
+            except Exception:
+                dbos.logger.warning(
+                    f"Exception encountered in queue thread: {traceback.format_exc()}"
+                )

dbos/system_database.py

@@ -3,7 +3,18 @@ import os
 import threading
 import time
 from enum import Enum
-from typing import Any, Dict, List, Literal, Optional, Sequence, Set, TypedDict, cast
+from typing import (
+    Any,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Sequence,
+    Set,
+    Tuple,
+    TypedDict,
+    cast,
+)
 
 import psycopg
 import sqlalchemy as sa
@@ -36,11 +47,6 @@ WorkflowStatuses = Literal[
 ]
 
 
-class WorkflowInputs(TypedDict):
-    args: Any
-    kwargs: Any
-
-
 class WorkflowStatusInternal(TypedDict):
     workflow_uuid: str
     status: WorkflowStatuses
@@ -127,7 +133,7 @@ class WorkflowInformation(TypedDict, total=False):
     # The role used to run this workflow.  Empty string if authorization is not required.
     authenticated_roles: List[str]
     # All roles the authenticated user has, if any.
-    input: Optional[WorkflowInputs]
+    input: Optional[utils.WorkflowInputs]
     output: Optional[str]
     error: Optional[str]
     request: Optional[str]
@@ -447,7 +453,7 @@ class SystemDatabase:
         if status == str(WorkflowStatusString.SUCCESS.value):
             return utils.deserialize(stat["output"])
         elif status == str(WorkflowStatusString.ERROR.value):
-            raise utils.deserialize(stat["error"])
+            raise utils.deserialize_exception(stat["error"])
         return None
 
     def get_workflow_info(
@@ -487,7 +493,7 @@ class SystemDatabase:
             self._exported_temp_txn_wf_status.discard(workflow_uuid)
             self._temp_txn_wf_ids.discard(workflow_uuid)
 
-    def get_workflow_inputs(self, workflow_uuid: str) -> Optional[WorkflowInputs]:
+    def get_workflow_inputs(self, workflow_uuid: str) -> Optional[utils.WorkflowInputs]:
         with self.engine.begin() as c:
             row = c.execute(
                 sa.select(SystemSchema.workflow_inputs.c.inputs).where(
@@ -496,7 +502,7 @@ class SystemDatabase:
             ).fetchone()
             if row is None:
                 return None
-            inputs: WorkflowInputs = utils.deserialize(row[0])
+            inputs: utils.WorkflowInputs = utils.deserialize_args(row[0])
             return inputs
 
     def get_workflows(self, input: GetWorkflowsInput) -> GetWorkflowsOutput:

dbos/utils.py

@@ -1,14 +1,55 @@
-from typing import Any
+import types
+from typing import Any, Dict, Tuple, TypedDict
 
 import jsonpickle  # type: ignore
 
 
+class WorkflowInputs(TypedDict):
+    args: Tuple[Any, ...]
+    kwargs: Dict[str, Any]
+
+
+def validate_item(data: Any) -> None:
+    if isinstance(data, (types.FunctionType, types.MethodType)):
+        raise TypeError("Serialized data item should not be a function")
+
+
 def serialize(data: Any) -> str:
     """Serialize an object to a JSON string using jsonpickle."""
-    encoded_data: str = jsonpickle.encode(data)
+    validate_item(data)
+    encoded_data: str = jsonpickle.encode(data, unpicklable=True)
+    return encoded_data
+
+
+def serialize_args(data: WorkflowInputs) -> str:
+    """Serialize args to a JSON string using jsonpickle."""
+    arg: Any
+    for arg in data["args"]:
+        validate_item(arg)
+    for arg in data["kwargs"].values():
+        validate_item(arg)
+    encoded_data: str = jsonpickle.encode(data, unpicklable=True)
+    return encoded_data
+
+
+def serialize_exception(data: Exception) -> str:
+    """Serialize an Exception object to a JSON string using jsonpickle."""
+    encoded_data: str = jsonpickle.encode(data, unpicklable=True)
     return encoded_data
 
 
 def deserialize(serialized_data: str) -> Any:
     """Deserialize a JSON string back to a Python object using jsonpickle."""
     return jsonpickle.decode(serialized_data)
+
+
+def deserialize_args(serialized_data: str) -> WorkflowInputs:
+    """Deserialize a JSON string back to a Python object list using jsonpickle."""
+    args: WorkflowInputs = jsonpickle.decode(serialized_data)
+    return args
+
+
+def deserialize_exception(serialized_data: str) -> Exception:
+    """Deserialize JSON string back to a Python Exception using jsonpickle."""
+    upo: Exception = jsonpickle.decode(serialized_data)
+    return upo

{dbos-0.8.0a0.dist-info → dbos-0.8.0a3.dist-info}/METADATA

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

{dbos-0.8.0a0.dist-info → dbos-0.8.0a3.dist-info}/RECORD

@@ -1,13 +1,13 @@
-dbos-0.8.0a0.dist-info/METADATA,sha256=25gzYG9xJsemmuA1r8k_jnJyp3z3hHnRoc2UzAR2CQI,5010
-dbos-0.8.0a0.dist-info/WHEEL,sha256=Vza3XR51HW1KmFP0iIMUVYIvz0uQuKJpIXKYOBGQyFQ,90
-dbos-0.8.0a0.dist-info/entry_points.txt,sha256=z6GcVANQV7Uw_82H9Ob2axJX6V3imftyZsljdh-M1HU,54
-dbos-0.8.0a0.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
-dbos/__init__.py,sha256=jjlBFzSAzO2e-LD5IKJw7bRqZjrxpF5Sn-_JUJJptHU,680
+dbos-0.8.0a3.dist-info/METADATA,sha256=u3PDfNC0j5qeu05Uq1834FHAScD6ECq64z9_XOQqdiw,5010
+dbos-0.8.0a3.dist-info/WHEEL,sha256=Vza3XR51HW1KmFP0iIMUVYIvz0uQuKJpIXKYOBGQyFQ,90
+dbos-0.8.0a3.dist-info/entry_points.txt,sha256=z6GcVANQV7Uw_82H9Ob2axJX6V3imftyZsljdh-M1HU,54
+dbos-0.8.0a3.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
+dbos/__init__.py,sha256=-h1QgWNL11CiLlHEKa2ycAJVJw5SXYZ4BGNNWBAiE9k,726
 dbos/admin_sever.py,sha256=Qg5T3YRrbPW05PR_99yAaxgo1ugQrAp_uTeTqSfjm_k,3397
 dbos/application_database.py,sha256=knFK8We8y6WrIpnFCKvFq5hvSuFQqUuJqOqDpSVMCPI,5521
 dbos/cli.py,sha256=z5dXbbnGWzSC3E1rfS8Lp1_OIImzcDKM7jP-iu_Q4aI,8602
-dbos/context.py,sha256=1Xp0i-ZP72Vj6eMdHuyfXi9RHnoT2w4MZ2Kx1CrKnQ8,16559
-dbos/core.py,sha256=LnQRHLjfpdlzlRkElduay8nObNMyRNC8SJSZEzqei7o,29161
+dbos/context.py,sha256=4MsxZdoh1WIsgoUsaxo0B6caGN6xq2WC60MzbBppzGk,17738
+dbos/core.py,sha256=nV0w0wCkKm0VRVbUO0DgRnVqX5ue1Bn37UyAsdSXl48,29342
 dbos/dbos-config.schema.json,sha256=azpfmoDZg7WfSy3kvIsk9iEiKB_-VZt03VEOoXJAkqE,5331
 dbos/dbos.py,sha256=RtDcvKe4sm1TlnCGU4cyex-UI7hMMlhgzmOl1NuRLo4,29294
 dbos/dbos_config.py,sha256=NJVze2GkKgYUmcPP31Unb-QpsA0TzImEeQGJgVq6W6k,5352
@@ -20,11 +20,12 @@ dbos/kafka_message.py,sha256=NYvOXNG3Qn7bghn1pv3fg4Pbs86ILZGcK4IB-MLUNu0,409
 dbos/logger.py,sha256=D-aFSZUCHBP34J1IZ5YNkTrJW-rDiH3py_v9jLU4Yrk,3565
 dbos/migrations/env.py,sha256=38SIGVbmn_VV2x2u1aHLcPOoWgZ84eCymf3g_NljmbU,1626
 dbos/migrations/script.py.mako,sha256=MEqL-2qATlST9TAOeYgscMn1uy6HUS9NFvDgl93dMj8,635
+dbos/migrations/versions/50f3227f0b4b_fix_job_queue.py,sha256=ZtnsZFMuon-D0n8V5BR10jQEqJPUsYsOwt29FAoKG8g,868
 dbos/migrations/versions/5c361fc04708_added_system_tables.py,sha256=QMgFMb0aLgC25YicsvPSr6AHRCA6Zd66hyaRUhwKzrQ,6404
 dbos/migrations/versions/a3b18ad34abe_added_triggers.py,sha256=Rv0ZsZYZ_WdgGEULYsPfnp4YzaO5L198gDTgYY39AVA,2022
 dbos/migrations/versions/eab0cc1d9a14_job_queue.py,sha256=_9-FCW-zOpCQfblTS_yRLtFiUaWlC1tM4BoKBTDeH9k,1395
 dbos/py.typed,sha256=QfzXT1Ktfk3Rj84akygc7_42z0lRpCq0Ilh8OXI6Zas,44
-dbos/queue.py,sha256=5unMPX1ThoVWfXOR2IUSeIjgATgcX_YmZS3Fz1Uoc7o,1228
+dbos/queue.py,sha256=ngY1MN3xD7jAvEXlrl_D16FUPpP_vpRgbyERLfPyU9Y,1437
 dbos/recovery.py,sha256=zqtO_ExGoIErLMVnbneU3VeHLVWvhV4jnfqssAVlQQk,2016
 dbos/registrations.py,sha256=gMI-u05tv5bpvyddQGtoUgCsqARx51aOY7p0JXPafQo,6539
 dbos/request.py,sha256=-FIwtknayvRl6OjvqO4V2GySVzSdP1Ft3cc9ZBS-PLY,928
@@ -34,7 +35,7 @@ dbos/scheduler/scheduler.py,sha256=Sz4EIpAtur7so2YajTic64GrTpa4qPw8QxXn0M34v80,1
 dbos/schemas/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/schemas/application_database.py,sha256=q_Wr2XbiZNBYFkOtu7uKavo1T_cSOBblxKGHThYGGsY,962
 dbos/schemas/system_database.py,sha256=ed4c1UntsD-cqXD0ekM4jvcYYEViavDh_G6c0pVDe7k,4938
-dbos/system_database.py,sha256=VSGFSPubbMokGYsZfRb6cQPltLfjoWM-Re_2Gj9qkRc,41844
+dbos/system_database.py,sha256=Os1-qqnrirKm_K37zuMgk14mm4ziXNYDjZnM1o8A450,41864
 dbos/templates/hello/README.md,sha256=GhxhBj42wjTt1fWEtwNriHbJuKb66Vzu89G4pxNHw2g,930
 dbos/templates/hello/__package/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dbos/templates/hello/__package/main.py,sha256=eI0SS9Nwj-fldtiuSzIlIG6dC91GXXwdRsoHxv6S_WI,2719
@@ -46,6 +47,6 @@ dbos/templates/hello/migrations/script.py.mako,sha256=MEqL-2qATlST9TAOeYgscMn1uy
 dbos/templates/hello/migrations/versions/2024_07_31_180642_init.py,sha256=U5thFWGqNN4QLrNXT7wUUqftIFDNE5eSdqD8JNW1mec,942
 dbos/templates/hello/start_postgres_docker.py,sha256=lQVLlYO5YkhGPEgPqwGc7Y8uDKse9HsWv5fynJEFJHM,1681
 dbos/tracer.py,sha256=GaXDhdKKF_IQp5SAMipGXiDVwteRKjNbrXyYCH1mor0,2520
-dbos/utils.py,sha256=hWj9iWDrby2cVEhb0pG-IdnrxLqP64NhkaWUXiLc8bA,402
+dbos/utils.py,sha256=lwRymY-y7GprAS8pKmbICQvOJd5eGxKGTxCMFn0OwaQ,1739
 version/__init__.py,sha256=L4sNxecRuqdtSFdpUGX3TtBi9KL3k7YsZVIvv-fv9-A,1678
-dbos-0.8.0a0.dist-info/RECORD,,
+dbos-0.8.0a3.dist-info/RECORD,,