@dbos-inc/kysely-datasource 4.26.6-preview → 4.26.7-preview

package/index.ts

@@ -1,10 +1,12 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 // using https://kysely.dev/
-import { DBOS, DBOSWorkflowConflictError, FunctionName } from '@dbos-inc/dbos-sdk';
+import { DBOS, FunctionName } from '@dbos-inc/dbos-sdk';
 import {
   type DataSourceTransactionHandler,
   isPGRetriableTransactionError,
-  isPGKeyConflictError,
+  DBOSError,
+  DBOSStepAlreadyRecordedError,
+  replayRecordedStep,
   registerTransaction,
   runTransaction,
   DBOSDataSource,
@@ -136,23 +138,31 @@ class KyselyTransactionHandler implements DataSourceTransactionHandler {
   }
 
   async #recordError(workflowID: string, stepID: number, error: string): Promise<void> {
-    try {
-      await this.#kyselyDB
-        .insertInto('dbos.transaction_completion')
-        .values({
-          workflow_id: workflowID,
-          function_num: stepID,
-          error,
-          output: null,
-        })
-        .execute();
-    } catch (error) {
-      if (isPGKeyConflictError(error)) {
-        throw new DBOSWorkflowConflictError(workflowID);
-      } else {
-        throw error;
-      }
+    const inserted = await this.#kyselyDB
+      .insertInto('dbos.transaction_completion')
+      .values({
+        workflow_id: workflowID,
+        function_num: stepID,
+        error,
+        output: null,
+      })
+      .onConflict((oc) => oc.columns(['workflow_id', 'function_num']).doNothing())
+      .returning('workflow_id')
+      .executeTakeFirst();
+    if (inserted === undefined) {
+      throw new DBOSStepAlreadyRecordedError(workflowID, stepID);
+    }
+  }
+
+  // A duplicate execution won the race, so its recorded outcome is the durable one.
+  async #replayConflictingStep<Return>(workflowID: string, stepID: number): Promise<Return> {
+    const recorded = await this.#checkExecution(workflowID, stepID);
+    if (recorded === undefined) {
+      throw new DBOSError(
+        `Step ${stepID} of workflow ${workflowID} conflicted with a concurrent execution, but no recorded outcome was found`,
+      );
     }
+    return replayRecordedStep<Return>(recorded);
   }
 
   static async #recordOutput(
@@ -161,22 +171,19 @@ class KyselyTransactionHandler implements DataSourceTransactionHandler {
     stepID: number,
     output: string | null,
   ): Promise<void> {
-    try {
-      await client
-        .insertInto('dbos.transaction_completion')
-        .values({
-          workflow_id: workflowID,
-          function_num: stepID,
-          output,
-          error: null,
-        })
-        .execute();
-    } catch (error) {
-      if (isPGKeyConflictError(error)) {
-        throw new DBOSWorkflowConflictError(workflowID);
-      } else {
-        throw error;
-      }
+    const inserted = await client
+      .insertInto('dbos.transaction_completion')
+      .values({
+        workflow_id: workflowID,
+        function_num: stepID,
+        output,
+        error: null,
+      })
+      .onConflict((oc) => oc.columns(['workflow_id', 'function_num']).doNothing())
+      .returning('workflow_id')
+      .executeTakeFirst();
+    if (inserted === undefined) {
+      throw new DBOSStepAlreadyRecordedError(workflowID, stepID);
     }
   }
 
@@ -204,12 +211,7 @@ class KyselyTransactionHandler implements DataSourceTransactionHandler {
       // Check to see if this tx has already been executed
       const previousResult = saveResults ? await this.#checkExecution(workflowID, stepID!) : undefined;
       if (previousResult) {
-        DBOS.span?.setAttribute('cached', true);
-
-        if ('error' in previousResult) {
-          throw SuperJSON.parse(previousResult.error);
-        }
-        return (previousResult.output ? SuperJSON.parse(previousResult.output) : null) as Return;
+        return replayRecordedStep<Return>(previousResult);
       }
 
       try {
@@ -236,6 +238,9 @@ class KyselyTransactionHandler implements DataSourceTransactionHandler {
 
         return result;
       } catch (error) {
+        if (saveResults && error instanceof DBOSStepAlreadyRecordedError) {
+          return await this.#replayConflictingStep<Return>(workflowID, stepID!);
+        }
         if (isPGRetriableTransactionError(error)) {
           DBOS.span?.addEvent('TXN SERIALIZATION FAILURE', { retryWaitMillis: retryWaitMS }, performance.now());
           await new Promise((resolve) => setTimeout(resolve, retryWaitMS));
@@ -244,7 +249,14 @@ class KyselyTransactionHandler implements DataSourceTransactionHandler {
         } else {
           if (saveResults) {
             const message = SuperJSON.stringify(error);
-            await this.#recordError(workflowID, stepID!, message);
+            try {
+              await this.#recordError(workflowID, stepID!, message);
+            } catch (recordError) {
+              if (recordError instanceof DBOSStepAlreadyRecordedError) {
+                return await this.#replayConflictingStep<Return>(workflowID, stepID!);
+              }
+              throw recordError;
+            }
           }
 
           throw error;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@dbos-inc/kysely-datasource",
-  "version": "4.26.6-preview",
+  "version": "4.26.7-preview",
   "description": "DBOS DataSource library for Kysely Query Builder with PostgreSQL support",
   "license": "MIT",
   "main": "dist/index.js",

package/tests/datasource.test.ts

@@ -298,6 +298,57 @@ describe('KyselyDataSource', () => {
     ]);
     expect(txResults.rows.length).toBe(0);
   });
+
+  test('duplicate execution replays the winner result', async () => {
+    await userDB.query('DELETE FROM race_side_effects');
+    raceState.callCount = 0;
+
+    // A loser whose body succeeds: the collision happens on the result-recording insert.
+    raceState.plantWinner = true;
+    raceState.plantWinnerError = false;
+    raceState.fail = false;
+    const wfid1 = randomUUID();
+    await expect(DBOS.withNextWorkflowID(wfid1, () => regRaceWorkflow())).resolves.toBe('winner-result');
+    expect(raceState.callCount).toBe(1); // the loser did run its body
+
+    // A loser whose body fails: the collision moves to the error-recording insert,
+    // and the winner's result still wins over the loser's own error.
+    raceState.plantWinner = true;
+    raceState.plantWinnerError = false;
+    raceState.fail = true;
+    const wfid2 = randomUUID();
+    await expect(DBOS.withNextWorkflowID(wfid2, () => regRaceWorkflow())).resolves.toBe('winner-result');
+    expect(raceState.callCount).toBe(2);
+
+    // A winner that recorded an error: the loser replays that throw instead of its own result.
+    raceState.plantWinner = true;
+    raceState.plantWinnerError = true;
+    raceState.fail = false;
+    const wfid3 = randomUUID();
+    const winnerError = await throws(() => DBOS.withNextWorkflowID(wfid3, () => regRaceWorkflow()));
+    expect((winnerError as Error).message).toBe('winner-error');
+    expect(raceState.callCount).toBe(3);
+
+    // Every loser's writes were discarded and the winners' records still stand.
+    const { rows: tags } = await userDB.query<{ tag: string }>('SELECT tag FROM race_side_effects');
+    expect(tags).toHaveLength(0);
+    const { rows: txOutput } = await userDB.query<transaction_completion>(
+      'SELECT * FROM dbos.transaction_completion WHERE workflow_id = ANY($1)',
+      [[wfid1, wfid2]],
+    );
+    expect(txOutput).toHaveLength(2);
+    for (const row of txOutput) {
+      expect(row.error).toBeNull(); // no loser error was ever recorded
+      expect(SuperJSON.parse(row.output!)).toBe('winner-result');
+    }
+    const { rows: txError } = await userDB.query<transaction_completion>(
+      'SELECT * FROM dbos.transaction_completion WHERE workflow_id = $1',
+      [wfid3],
+    );
+    expect(txError).toHaveLength(1);
+    expect(txError[0].output).toBeNull(); // the loser never overwrote the winner's error with its own result
+    expect(SuperJSON.parse<Error>(txError[0].error!).message).toBe('winner-error');
+  });
 });
 
 export interface greetings {
@@ -325,6 +376,7 @@ async function createDatabases(userDB: Pool, createTxCompletion: boolean) {
       await client.query(
         'CREATE TABLE greetings(name text NOT NULL, greet_count integer DEFAULT 0, PRIMARY KEY(name))',
       );
+      await client.query('CREATE TABLE race_side_effects(tag text NOT NULL)');
     } finally {
       client.release();
     }
@@ -351,6 +403,47 @@ async function errorFunction(user: string) {
   throw new Error(`test error ${Date.now()}`);
 }
 
+const raceState = { callCount: 0, plantWinner: false, plantWinnerError: false, fail: false };
+
+async function raceFunction() {
+  raceState.callCount += 1;
+  await sql`INSERT INTO race_side_effects(tag) VALUES (${`run-${raceState.callCount}`})`.execute(dataSource.client);
+  if (raceState.plantWinner) {
+    raceState.plantWinner = false;
+    await plantWinnerCompletion();
+  }
+  if (raceState.fail) {
+    throw new Error("loser's own failure");
+  }
+  return `result-${raceState.callCount}`;
+}
+
+// A duplicate execution commits its result while this transaction is still open, so this
+// execution's pre-check missed it but its completion insert will collide with it.
+async function plantWinnerCompletion() {
+  const winner = new Client(config.connection);
+  try {
+    await winner.connect();
+    const [column, value] = raceState.plantWinnerError
+      ? ['error', SuperJSON.stringify(new Error('winner-error'))]
+      : ['output', SuperJSON.stringify('winner-result')];
+    await winner.query(
+      `INSERT INTO dbos.transaction_completion (workflow_id, function_num, ${column}) VALUES ($1, $2, $3)`,
+      [DBOS.workflowID, DBOS.stepID, value],
+    );
+  } finally {
+    await winner.end();
+  }
+}
+
+const regRaceFunction = dataSource.registerTransaction(raceFunction);
+
+async function raceWorkflow() {
+  return await regRaceFunction();
+}
+
+const regRaceWorkflow = DBOS.registerWorkflow(raceWorkflow);
+
 async function readFunction(user: string) {
   const row = await dataSource.client
     .selectFrom('greetings')