dbos 0.5.0a4__py3-none-any.whl → 0.5.0a7__py3-none-any.whl

dbos/application_database.py

@@ -3,6 +3,7 @@ from typing import Optional, TypedDict, cast
 import sqlalchemy as sa
 import sqlalchemy.dialects.postgresql as pg
 import sqlalchemy.exc as sa_exc
+from sqlalchemy.exc import DBAPIError
 from sqlalchemy.orm import Session, sessionmaker
 
 from dbos.error import DBOSWorkflowConflictIDError
@@ -95,8 +96,10 @@ class ApplicationDatabase:
                     ),
                 )
             )
-        except sa_exc.IntegrityError:
-            raise DBOSWorkflowConflictIDError(output["workflow_uuid"])
+        except DBAPIError as dbapi_error:
+            if dbapi_error.orig.pgcode == "23505":  # type: ignore
+                raise DBOSWorkflowConflictIDError(output["workflow_uuid"])
+            raise dbapi_error
         except Exception as e:
             raise e
 
@@ -118,8 +121,10 @@ class ApplicationDatabase:
                         ),
                     )
                 )
-        except sa_exc.IntegrityError:
-            raise DBOSWorkflowConflictIDError(output["workflow_uuid"])
+        except DBAPIError as dbapi_error:
+            if dbapi_error.orig.pgcode == "23505":  # type: ignore
+                raise DBOSWorkflowConflictIDError(output["workflow_uuid"])
+            raise dbapi_error
         except Exception as e:
             raise e
 

dbos/core.py

@@ -136,14 +136,18 @@ def _init_workflow(
         "recovery_attempts": None,
     }
 
+    # If we have a class name, the first arg is the instance and do not serialize
+    if class_name is not None:
+        inputs = {"args": inputs["args"][1:], "kwargs": inputs["kwargs"]}
+
     if temp_wf_type != "transaction":
-        # Don't synchronously record the status or inputs for single transaction workflows
+        # Synchronously record the status and inputs for workflows and single-step workflows
+        # We also have to do this for single-step workflows because of the foreign key constraint on the operation outputs table
         dbos.sys_db.update_workflow_status(status, False, ctx.in_recovery)
-
-        # If we have a class name, the first arg is the instance and do not serialize
-        if class_name is not None:
-            inputs = {"args": inputs["args"][1:], "kwargs": inputs["kwargs"]}
         dbos.sys_db.update_workflow_inputs(wfid, utils.serialize(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))
 
     return status
 
@@ -173,20 +177,6 @@ def _execute_workflow(
         status["error"] = utils.serialize(error)
         dbos.sys_db.update_workflow_status(status)
         raise error
-    finally:
-        if get_temp_workflow_type(func) == "transaction":
-            # Buffer the inputs for single transaction workflows
-            inputs: WorkflowInputs = {
-                "args": args,
-                "kwargs": kwargs,
-            }
-            # If we have a class name, the first arg is the instance and do not serialize
-            class_name = get_dbos_class_name(get_func_info(func), func, args)
-            if class_name is not None:
-                inputs = {"args": inputs["args"][1:], "kwargs": inputs["kwargs"]}
-            dbos.sys_db.buffer_workflow_inputs(
-                status["workflow_uuid"], utils.serialize(inputs)
-            )
 
     return output
 
@@ -617,6 +607,7 @@ def _step(
                     utils.serialize(error) if error is not None else None
                 )
                 dbos.sys_db.record_operation_result(step_output)
+
                 if error is not None:
                     raise error
                 return output

dbos/dbos.py

@@ -72,7 +72,7 @@ from dbos.error import DBOSException, DBOSNonExistentWorkflowError
 
 from .application_database import ApplicationDatabase
 from .dbos_config import ConfigFile, load_config, set_env_vars
-from .logger import config_logger, dbos_logger, init_logger
+from .logger import add_otlp_to_all_loggers, config_logger, dbos_logger, init_logger
 from .system_database import SystemDatabase
 
 # Most DBOS functions are just any callable F, so decorators / wrappers work on F
@@ -204,7 +204,6 @@ class DBOS:
         cls: Type[DBOS],
         fastapi: Optional["FastAPI"] = None,
         config: Optional[ConfigFile] = None,
-        launch: bool = True,
     ) -> DBOS:
         global _dbos_global_instance
         global _dbos_global_registry
@@ -219,7 +218,7 @@ class DBOS:
                     )
                 config = _dbos_global_registry.config
             _dbos_global_instance = super().__new__(cls)
-            _dbos_global_instance.__init__(fastapi=fastapi, config=config, launch=launch)  # type: ignore
+            _dbos_global_instance.__init__(fastapi=fastapi, config=config)  # type: ignore
         else:
             if (config is not None and _dbos_global_instance.config is not config) or (
                 _dbos_global_instance.fastapi is not fastapi
@@ -242,7 +241,6 @@ class DBOS:
         self,
         fastapi: Optional["FastAPI"] = None,
         config: Optional[ConfigFile] = None,
-        launch: bool = True,
     ) -> None:
         if hasattr(self, "_initialized") and self._initialized:
             return
@@ -270,7 +268,6 @@ class DBOS:
 
             setup_fastapi_middleware(self.fastapi)
             self.fastapi.on_event("startup")(self.launch)
-            launch = False
 
         # Register send_stub as a workflow
         def send_temp_workflow(
@@ -283,8 +280,8 @@ class DBOS:
         set_temp_workflow_type(send_temp_workflow, "send")
         self._registry.register_wf_function(TEMP_SEND_WF_NAME, temp_send_wf)
 
-        if launch:
-            self.launch()
+        for handler in dbos_logger.handlers:
+            handler.flush()
 
     @property
     def executor(self) -> ThreadPoolExecutor:
@@ -342,9 +339,13 @@ class DBOS:
             self.executor.submit(func, *args, **kwargs)
         self._registry.pollers = []
 
-        dbos_logger.info("DBOS initialized")
+        dbos_logger.info("DBOS launched")
+
+        # Flush handlers and add OTLP to all loggers if enabled
+        # to enable their export in DBOS Cloud
         for handler in dbos_logger.handlers:
             handler.flush()
+        add_otlp_to_all_loggers()
 
     def _destroy(self) -> None:
         self._initialized = False

dbos/logger.py

@@ -12,6 +12,7 @@ if TYPE_CHECKING:
     from dbos.dbos_config import ConfigFile
 
 dbos_logger = logging.getLogger("dbos")
+otlp_handler, otlp_transformer = None, None
 
 
 class DBOSLogTransformer(logging.Filter):
@@ -73,25 +74,28 @@ def config_logger(config: "ConfigFile") -> None:
                 export_timeout_millis=5000,
             )
         )
+        global otlp_handler
         otlp_handler = LoggingHandler(logger_provider=log_provider)
 
         # Attach DBOS-specific attributes to all log entries.
+        global otlp_transformer
         otlp_transformer = DBOSLogTransformer()
 
-        # Direct all logs to OTLP
-        add_otlp_to_all_loggers(otlp_handler, otlp_transformer)
+        # Direct DBOS logs to OTLP
+        dbos_logger.addHandler(otlp_handler)
+        dbos_logger.addFilter(otlp_transformer)
 
 
-def add_otlp_to_all_loggers(
-    otlp_handler: LoggingHandler, otlp_transformer: DBOSLogTransformer
-) -> None:
-    root = logging.root
+def add_otlp_to_all_loggers() -> None:
+    if otlp_handler is not None and otlp_transformer is not None:
+        root = logging.root
 
-    root.addHandler(otlp_handler)
-    root.addFilter(otlp_transformer)
+        root.addHandler(otlp_handler)
+        root.addFilter(otlp_transformer)
 
-    for logger_name in root.manager.loggerDict:
-        logger = logging.getLogger(logger_name)
-        if not logger.propagate:
-            logger.addHandler(otlp_handler)
-        logger.addFilter(otlp_transformer)
+        for logger_name in root.manager.loggerDict:
+            if logger_name != dbos_logger.name:
+                logger = logging.getLogger(logger_name)
+                if not logger.propagate:
+                    logger.addHandler(otlp_handler)
+                logger.addFilter(otlp_transformer)

dbos/system_database.py

@@ -4,13 +4,25 @@ import select
 import threading
 import time
 from enum import Enum
-from typing import Any, Dict, List, Literal, Optional, Sequence, TypedDict, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Sequence,
+    Set,
+    TypedDict,
+    cast,
+)
 
 import psycopg2
 import sqlalchemy as sa
 import sqlalchemy.dialects.postgresql as pg
 from alembic import command
 from alembic.config import Config
+from sqlalchemy.exc import DBAPIError
 
 import dbos.utils as utils
 from dbos.error import (
@@ -199,6 +211,9 @@ class SystemDatabase:
         # Initialize the workflow status and inputs buffers
         self._workflow_status_buffer: Dict[str, WorkflowStatusInternal] = {}
         self._workflow_inputs_buffer: Dict[str, str] = {}
+        # Two sets for tracking which single-transaction workflows have been exported to the status table
+        self._exported_temp_txn_wf_status: Set[str] = set()
+        self._temp_txn_wf_ids: Set[str] = set()
         self._is_flushing_status_buffer = False
 
         # Now we can run background processes
@@ -264,6 +279,10 @@ class SystemDatabase:
             with self.engine.begin() as c:
                 c.execute(cmd)
 
+        # Record we have exported status for this single-transaction workflow
+        if status["workflow_uuid"] in self._temp_txn_wf_ids:
+            self._exported_temp_txn_wf_status.add(status["workflow_uuid"])
+
     def set_workflow_status(
         self,
         workflow_uuid: str,
@@ -454,6 +473,11 @@ class SystemDatabase:
             with self.engine.begin() as c:
                 c.execute(cmd)
 
+        if workflow_uuid in self._temp_txn_wf_ids:
+            # Clean up the single-transaction tracking sets
+            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]:
         with self.engine.begin() as c:
             row = c.execute(
@@ -533,8 +557,10 @@ class SystemDatabase:
             else:
                 with self.engine.begin() as c:
                     c.execute(sql)
-        except sa.exc.IntegrityError:
-            raise DBOSWorkflowConflictIDError(result["workflow_uuid"])
+        except DBAPIError as dbapi_error:
+            if dbapi_error.orig.pgcode == "23505":  # type: ignore
+                raise DBOSWorkflowConflictIDError(result["workflow_uuid"])
+            raise dbapi_error
         except Exception as e:
             raise e
 
@@ -588,8 +614,11 @@ class SystemDatabase:
                         message=utils.serialize(message),
                     )
                 )
-            except sa.exc.IntegrityError:
-                raise DBOSNonExistentWorkflowError(destination_uuid)
+            except DBAPIError as dbapi_error:
+                # Foreign key violation
+                if dbapi_error.orig.pgcode == "23503":  # type: ignore
+                    raise DBOSNonExistentWorkflowError(destination_uuid)
+                raise dbapi_error
             except Exception as e:
                 raise e
             output: OperationResultInternal = {
@@ -705,7 +734,6 @@ class SystemDatabase:
                 )
                 notification_cursor = self.notification_conn.cursor()
 
-                dbos_logger.info("Listening to notifications")
                 notification_cursor.execute("LISTEN dbos_notifications_channel")
                 notification_cursor.execute("LISTEN dbos_workflow_events_channel")
                 while self._run_background_processes:
@@ -812,8 +840,10 @@ class SystemDatabase:
                         value=utils.serialize(message),
                     )
                 )
-            except sa.exc.IntegrityError:
-                raise DBOSDuplicateWorkflowEventError(workflow_uuid, key)
+            except DBAPIError as dbapi_error:
+                if dbapi_error.orig.pgcode == "23505":  # type: ignore
+                    raise DBOSDuplicateWorkflowEventError(workflow_uuid, key)
+                raise dbapi_error
             except Exception as e:
                 raise e
             output: OperationResultInternal = {
@@ -897,7 +927,7 @@ class SystemDatabase:
         return value
 
     def _flush_workflow_status_buffer(self) -> None:
-        """Export the workflow status buffer to the database, up to the batch size."""
+        """Export the workflow status buffer to the database, up to the batch size"""
         if len(self._workflow_status_buffer) == 0:
             return
 
@@ -905,13 +935,16 @@ class SystemDatabase:
         exported_status: Dict[str, WorkflowStatusInternal] = {}
         with self.engine.begin() as c:
             exported = 0
-            local_batch_size = min(
-                buffer_flush_batch_size, len(self._workflow_status_buffer)
-            )
-            while exported < local_batch_size:
+            status_iter = iter(list(self._workflow_status_buffer))
+            wf_id: Optional[str] = None
+            while (
+                exported < buffer_flush_batch_size
+                and (wf_id := next(status_iter, None)) is not None
+            ):
                 # Pop the first key in the buffer (FIFO)
-                wf_id = next(iter(self._workflow_status_buffer))
-                status = self._workflow_status_buffer.pop(wf_id)
+                status = self._workflow_status_buffer.pop(wf_id, None)
+                if status is None:
+                    continue
                 exported_status[wf_id] = status
                 try:
                     self.update_workflow_status(status, conn=c)
@@ -932,12 +965,18 @@ class SystemDatabase:
         exported_inputs: Dict[str, str] = {}
         with self.engine.begin() as c:
             exported = 0
-            local_batch_size = min(
-                buffer_flush_batch_size, len(self._workflow_inputs_buffer)
-            )
-            while exported < local_batch_size:
-                wf_id = next(iter(self._workflow_inputs_buffer))
-                inputs = self._workflow_inputs_buffer.pop(wf_id)
+            input_iter = iter(list(self._workflow_inputs_buffer))
+            wf_id: Optional[str] = None
+            while (
+                exported < buffer_flush_batch_size
+                and (wf_id := next(input_iter, None)) is not None
+            ):
+                if wf_id not in self._exported_temp_txn_wf_status:
+                    # Skip exporting inputs if the status has not been exported yet
+                    continue
+                inputs = self._workflow_inputs_buffer.pop(wf_id, None)
+                if inputs is None:
+                    continue
                 exported_inputs[wf_id] = inputs
                 try:
                     self.update_workflow_inputs(wf_id, inputs, conn=c)
@@ -972,6 +1011,7 @@ class SystemDatabase:
     def buffer_workflow_inputs(self, workflow_id: str, inputs: str) -> None:
         # inputs is a serialized WorkflowInputs string
         self._workflow_inputs_buffer[workflow_id] = inputs
+        self._temp_txn_wf_ids.add(workflow_id)
 
     @property
     def _is_buffers_empty(self) -> bool:

{dbos-0.5.0a4.dist-info → dbos-0.5.0a7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dbos
-Version: 0.5.0a4
+Version: 0.5.0a7
 Summary: A Python framework for backends that scale
 Author-Email: "DBOS, Inc." <contact@dbos.dev>
 License: MIT

{dbos-0.5.0a4.dist-info → dbos-0.5.0a7.dist-info}/RECORD

@@ -1,20 +1,20 @@
-dbos-0.5.0a4.dist-info/METADATA,sha256=l_v0uEPtP7uqtxzLfs6FCFNe3jLyhnG1of_HACqUUyo,5420
-dbos-0.5.0a4.dist-info/WHEEL,sha256=rSwsxJWe3vzyR5HCwjWXQruDgschpei4h_giTm0dJVE,90
-dbos-0.5.0a4.dist-info/entry_points.txt,sha256=3PmOPbM4FYxEmggRRdJw0oAsiBzKR8U0yx7bmwUmMOM,39
-dbos-0.5.0a4.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
+dbos-0.5.0a7.dist-info/METADATA,sha256=LdnWhAfO7Lnl85g-bvhXhoFYrCHkz7Jxmj_yqEKpPVc,5420
+dbos-0.5.0a7.dist-info/WHEEL,sha256=rSwsxJWe3vzyR5HCwjWXQruDgschpei4h_giTm0dJVE,90
+dbos-0.5.0a7.dist-info/entry_points.txt,sha256=3PmOPbM4FYxEmggRRdJw0oAsiBzKR8U0yx7bmwUmMOM,39
+dbos-0.5.0a7.dist-info/licenses/LICENSE,sha256=VGZit_a5-kdw9WT6fY5jxAWVwGQzgLFyPWrcVVUhVNU,1067
 dbos/__init__.py,sha256=X1LdP36NomDtvPfFwoMNtgXf81TO05jj7vltsp79UUw,787
 dbos/admin_sever.py,sha256=KtzH6aKyskCm4h3yulpy9jb5PIqRlYI2sjctw5mvaKY,3395
-dbos/application_database.py,sha256=cBSfyB-5KGFBtigLpFPWEAgmrKONAUnG1825vXPifBk,5411
+dbos/application_database.py,sha256=1K3kE96BgGi_QWOd2heXluyNTwFAwlUVuAR6JKKUqf0,5659
 dbos/cli.py,sha256=QnbGtZ8S963q3iyFvXNBcL4DB35r4SFMarlb5DRqN6M,7915
 dbos/context.py,sha256=JZMV2RtSpTK7lnyyWxeBmGPwrZSB00XZEP6R6MT9ygQ,15690
-dbos/core.py,sha256=8KkmHYQtxbtfFI_2sAf4DYmiznwYxLLhE6z-FLc3Gho,28675
+dbos/core.py,sha256=HfKnPpIaQqIBAHzP2hD67aSIchTHp87NgD21CcujKkE,28300
 dbos/dbos-config.schema.json,sha256=azpfmoDZg7WfSy3kvIsk9iEiKB_-VZt03VEOoXJAkqE,5331
-dbos/dbos.py,sha256=HngS2BUWSbWPmloXGr-KE81BQ8dpZtlvOXM4tx4_Qhg,26246
+dbos/dbos.py,sha256=-zrxmo_yN4vPTKQdyDtAcGlZI4-RV4PZBiIFayNSHyI,26342
 dbos/dbos_config.py,sha256=EkO0c0xaIM7_vAAqqnvNNEAKG5fOJbmmalqnZvaKYZA,5312
 dbos/decorators.py,sha256=lbPefsLK6Cya4cb7TrOcLglOpGT3pc6qjZdsQKlfZLg,629
 dbos/error.py,sha256=nBdLC4hxGO_K9V26YbDGOo7xi1CKuN4PsE_cBv7K8Cc,3798
 dbos/fastapi.py,sha256=ZFcMizyv3pizo5zf0sSF6U4GoR3rQH8LxGipkQIGHfU,2282
-dbos/logger.py,sha256=cfybbu6F1zsgYLEPW8D8V6h033u-YedLXnGMnQQM6-4,3341
+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/5c361fc04708_added_system_tables.py,sha256=QMgFMb0aLgC25YicsvPSr6AHRCA6Zd66hyaRUhwKzrQ,6404
@@ -28,7 +28,7 @@ dbos/scheduler/scheduler.py,sha256=uO4_9jmWW2rLv1ODL3lc1cE_37ZaVTgnvmFx_FAlN50,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=5V3vqnEzry0Hn7ZbVS9Gs_dJKia8uX8p7mGC82Ru8rk,4303
-dbos/system_database.py,sha256=h04PngaTdHxr1zfXcH6rSdac_vX0mSsD4SWHyOEWpJQ,38147
+dbos/system_database.py,sha256=84c53iAel113SRb7DcgFJ8XQNWBhD4VrCRCb0s5Oe8Y,39635
 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=hJgp3S14cseT7zWIZsPwjqdzwTCw1aLo8kPKsTvYz0Y,2976
@@ -42,4 +42,4 @@ dbos/templates/hello/start_postgres_docker.py,sha256=lQVLlYO5YkhGPEgPqwGc7Y8uDKs
 dbos/tracer.py,sha256=RPW9oxmX9tSc0Yq7O-FAhpQWBg1QT7Ni1Q06uwhtNDk,2237
 dbos/utils.py,sha256=hWj9iWDrby2cVEhb0pG-IdnrxLqP64NhkaWUXiLc8bA,402
 version/__init__.py,sha256=L4sNxecRuqdtSFdpUGX3TtBi9KL3k7YsZVIvv-fv9-A,1678
-dbos-0.5.0a4.dist-info/RECORD,,
+dbos-0.5.0a7.dist-info/RECORD,,