dbos 0.5.0a4__tar.gz → 0.5.0a5__tar.gz

{dbos-0.5.0a4 → dbos-0.5.0a5}/PKG-INFO

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

{dbos-0.5.0a4 → dbos-0.5.0a5}/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-0.5.0a4 → dbos-0.5.0a5}/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-0.5.0a4 → dbos-0.5.0a5}/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 = {
@@ -812,8 +841,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 +928,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 +936,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 +966,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 +1012,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 → dbos-0.5.0a5}/pyproject.toml

@@ -23,7 +23,7 @@ dependencies = [
 ]
 requires-python = ">=3.9"
 readme = "README.md"
-version = "0.5.0a4"
+version = "0.5.0a5"
 
 [project.license]
 text = "MIT"

{dbos-0.5.0a4 → dbos-0.5.0a5}/tests/test_concurrency.py

@@ -34,18 +34,24 @@ def test_concurrent_conflict_uuid(dbos: DBOS) -> None:
     condition = threading.Condition()
     step_count = 0
     txn_count = 0
+    notified = False
 
     @DBOS.step()
     def test_step() -> str:
-        nonlocal step_count
+        nonlocal step_count, notified
         step_count += 1
-        condition.acquire()
         if step_count == 1:
             # Wait for the other one to notify
+            condition.acquire()
             condition.wait()
+            notified = True
+            condition.release()
         else:
-            condition.notify()
-        condition.release()
+            while not notified:
+                condition.acquire()
+                condition.notify()
+                condition.release()
+                time.sleep(0.1)
         return DBOS.workflow_id
 
     @DBOS.workflow()
@@ -61,15 +67,20 @@ def test_concurrent_conflict_uuid(dbos: DBOS) -> None:
     @DBOS.transaction(isolation_level="REPEATABLE READ")
     def test_transaction() -> str:
         DBOS.sql_session.execute(text("SELECT 1")).fetchall()
-        nonlocal txn_count
+        nonlocal txn_count, notified
         txn_count += 1
-        condition.acquire()
         if txn_count == 1:
             # Wait for the other one to notify
+            condition.acquire()
             condition.wait()
+            notified = True
+            condition.release()
         else:
-            condition.notify()
-        condition.release()
+            while not notified:
+                condition.acquire()
+                condition.notify()
+                condition.release()
+                time.sleep(0.1)
         return DBOS.workflow_id
 
     def test_txn_thread(id: str) -> str:
@@ -89,19 +100,21 @@ def test_concurrent_conflict_uuid(dbos: DBOS) -> None:
 
     # Make sure temp workflows can handle conflicts as well.
     step_count = 0
+    notified = False
     wfuuid = str(uuid.uuid4())
     with ThreadPoolExecutor(max_workers=2) as executor:
         future1 = executor.submit(test_comm_thread, wfuuid)
         future2 = executor.submit(test_comm_thread, wfuuid)
 
-        assert future1.result() == wfuuid
-        assert future2.result() == wfuuid
+    assert future1.result() == wfuuid
+    assert future2.result() == wfuuid
 
     # Make sure temp transactions can handle conflicts as well.
     wfuuid = str(uuid.uuid4())
+    notified = False
     with ThreadPoolExecutor(max_workers=2) as executor:
         future1 = executor.submit(test_txn_thread, wfuuid)
         future2 = executor.submit(test_txn_thread, wfuuid)
 
-        assert future1.result() == wfuuid
-        assert future2.result() == wfuuid
+    assert future1.result() == wfuuid
+    assert future2.result() == wfuuid

{dbos-0.5.0a4 → dbos-0.5.0a5}/tests/test_dbos.py

@@ -251,6 +251,10 @@ def test_temp_workflow(dbos: DBOS) -> None:
     res = test_step("var")
     assert res == "var"
 
+    # Flush workflow inputs buffer shouldn't fail due to foreign key violation.
+    # It should properly skip the transaction inputs.
+    dbos.sys_db._flush_workflow_inputs_buffer()
+
     # Wait for buffers to flush
     dbos.sys_db.wait_for_buffer_flush()
     wfs = dbos.sys_db.get_workflows(gwi)