Datasource transaction raises IntegrityError instead of replaying the recorded result when a duplicate execution wins the race

[SamuelXing]

When two executions of the same workflow run a datasource transaction step concurrently — which is exactly what recovery produces when it re-runs a workflow whose original executor is still alive — the losing execution doesn't get the recorded result. It gets a raw IntegrityError.

The witness insert in _datasource.py is a plain INSERT with no conflict handling:

dbos-transact-py/dbos/_datasource.py

Lines 566 to 583 in 5d920f7

	def _record_result(
	self,
	conn: Union[sa.Connection, Session],
	workflow_id: str,
	step_id: int,
	output: Optional[str],
	error: Optional[str],
	serialization: Optional[str],
	) -> None:
	conn.execute(
	sa.insert(DatasourceSchema.datasource_outputs).values(
	workflow_id=workflow_id,
	step_id=step_id,
	output=output,
	error=error,
	serialization=serialization,
	)
	)

and datasource_outputs has PRIMARY KEY (workflow_id, step_id), so the race plays out like this:

1. Both executions call _check_execution — no row yet, so both run the function.
2. Execution A commits: app writes + its datasource_outputs row, atomically.
3. Execution B's insert hits the primary-key violation and its transaction rolls back. So far so good — B's app writes are undone and exactly once holds.
4. But 23505 is neither in retriable_postgres_exception (connection/resource errors only) nor a serialization error (40001/40P01 only), so B falls through to the error-recording path — and _record_error does another plain insert into the same table, hits the same violation, and raises that IntegrityError out of the error handler, masking the original.

Durable state ends up correct, but the losing execution surfaces an unrelated database error where it should return A's result — the same thing _check_execution would have handed it a moment later.

Repro

The script below doesn't actually race two processes, instead it runs the losing execution once and plants the winner's row by hand, from inside the transaction body, through a second engine. Separate connection, separate transaction — to Postgres and to run_tx_step this looks exactly like a concurrent executor finishing the same step, and it fails every run:

loser (SDK transaction, conn #1)        "winner" (second engine, conn #2)
--------------------------------        ---------------------------------
_check_execution → no row
transaction opens, body runs
                                        INSERT witness row, COMMIT
body returns 42
_record_result INSERT  ← collides

import os

import sqlalchemy as sa

from dbos import DBOS, DBOSConfig
from dbos._context import get_local_dbos_context
from dbos._datasource import SQLAlchemyDatasource

DB_URL = os.environ["DBOS_DATABASE_URL"]

DBOS(config={"name": "ds-race-repro", "system_database_url": DB_URL, "run_admin_server": False})

ds = SQLAlchemyDatasource.create(DB_URL)
rival = sa.create_engine(sa.make_url(DB_URL).set(drivername="postgresql+psycopg"))


@ds.transaction()
def racy_txn() -> int:
    ctx = get_local_dbos_context()
    assert ctx is not None
    # a concurrent execution of this same step commits first
    with rival.begin() as conn:
        conn.execute(
            sa.text(
                "INSERT INTO dbos.datasource_outputs"
                " (workflow_id, step_id, output, error, serialization)"
                " VALUES (:w, :s, :o, NULL, 'json')"
            ),
            {"w": ctx.workflow_id, "s": ctx.curr_step_function_id, "o": "999"},
        )
    ds.sql_session().execute(sa.text("SELECT 1"))
    return 42


@DBOS.workflow()
def wf() -> int:
    return racy_txn()


DBOS.launch()
try:
    print("workflow returned:", wf())
except Exception as e:
    print(f"workflow raised: {type(e).__name__}")
    cause = e.__context__
    while cause is not None:
        print(f"  ...while handling: {type(cause).__name__}")
        cause = cause.__context__
finally:
    DBOS.destroy()

Run it against a fresh Postgres database. Output on current main:

Possible fix

We can make the witness insert ON CONFLICT DO NOTHING (plus the sqlite equivalent) and treat rowcount == 0 as "someone already recorded this step": roll back, re-run _check_execution, return _replay_recorded(...). And give _record_error the same guard so the error path can't throw on an already-recorded step.

[kraftp]

Good catch, I'll take a look