@dbos-inc/node-pg-datasource 3.0.6-preview

package/index.ts

@@ -0,0 +1,240 @@
+// using https://github.com/brianc/node-postgres
+
+import { DBOS, DBOSWorkflowConflictError } from '@dbos-inc/dbos-sdk';
+import {
+  type DataSourceTransactionHandler,
+  createTransactionCompletionSchemaPG,
+  createTransactionCompletionTablePG,
+  isPGRetriableTransactionError,
+  isPGKeyConflictError,
+  registerTransaction,
+  runTransaction,
+  PGTransactionConfig as NodePostgresTransactionOptions,
+  DBOSDataSource,
+  registerDataSource,
+} from '@dbos-inc/dbos-sdk/datasource';
+import { Client, type ClientBase, type ClientConfig, Pool, type PoolConfig } from 'pg';
+import { AsyncLocalStorage } from 'node:async_hooks';
+import { SuperJSON } from 'superjson';
+
+interface NodePostgresDataSourceContext {
+  client: ClientBase;
+}
+
+export { NodePostgresTransactionOptions };
+
+const asyncLocalCtx = new AsyncLocalStorage<NodePostgresDataSourceContext>();
+
+class NodePGDSTH implements DataSourceTransactionHandler {
+  readonly name: string;
+  readonly dsType = 'NodePostgresDataSource';
+  readonly #pool: Pool;
+
+  constructor(name: string, config: PoolConfig) {
+    this.name = name;
+    this.#pool = new Pool(config);
+  }
+
+  initialize(): Promise<void> {
+    return Promise.resolve();
+  }
+
+  destroy(): Promise<void> {
+    return this.#pool.end();
+  }
+
+  static async #checkExecution(
+    client: ClientBase,
+    workflowID: string,
+    functionNum: number,
+  ): Promise<{ output: string | null } | undefined> {
+    const { rows } = await client.query<{ output: string }>(
+      /*sql*/ `SELECT output FROM dbos.transaction_completion
+       WHERE workflow_id = $1 AND function_num = $2`,
+      [workflowID, functionNum],
+    );
+    return rows.length > 0 ? { output: rows[0].output } : undefined;
+  }
+
+  static async #recordOutput(
+    client: ClientBase,
+    workflowID: string,
+    functionNum: number,
+    output: string | null,
+  ): Promise<void> {
+    try {
+      await client.query(
+        /*sql*/
+        `INSERT INTO dbos.transaction_completion (workflow_id, function_num, output) 
+         VALUES ($1, $2, $3)`,
+        [workflowID, functionNum, output],
+      );
+    } catch (error) {
+      if (isPGKeyConflictError(error)) {
+        throw new DBOSWorkflowConflictError(workflowID);
+      } else {
+        throw error;
+      }
+    }
+  }
+
+  async #transaction<Return>(
+    func: (client: ClientBase) => Promise<Return>,
+    config: NodePostgresTransactionOptions = {},
+  ): Promise<Return> {
+    const isolationLevel = config?.isolationLevel ? `ISOLATION LEVEL ${config.isolationLevel}` : '';
+    const readOnly = config?.readOnly ?? false;
+    const accessMode = config?.readOnly === undefined ? '' : readOnly ? 'READ ONLY' : 'READ WRITE';
+
+    const client = await this.#pool.connect();
+    try {
+      await client.query(/*sql*/ `BEGIN ${isolationLevel} ${accessMode}`);
+      const result = await func(client);
+      await client.query(/*sql*/ `COMMIT`);
+      return result;
+    } catch (error) {
+      await client.query(/*sql*/ `ROLLBACK`);
+      throw error;
+    } finally {
+      client.release();
+    }
+  }
+
+  async invokeTransactionFunction<This, Args extends unknown[], Return>(
+    config: NodePostgresTransactionOptions | undefined,
+    target: This,
+    func: (this: This, ...args: Args) => Promise<Return>,
+    ...args: Args
+  ): Promise<Return> {
+    const workflowID = DBOS.workflowID;
+    if (workflowID === undefined) {
+      throw new Error('Workflow ID is not set.');
+    }
+    const functionNum = DBOS.stepID;
+    if (functionNum === undefined) {
+      throw new Error('Function Number is not set.');
+    }
+
+    const readOnly = config?.readOnly ?? false;
+    let retryWaitMS = 1;
+    const backoffFactor = 1.5;
+    const maxRetryWaitMS = 2000;
+
+    while (true) {
+      try {
+        const result = await this.#transaction<Return>(
+          async (client) => {
+            // Check to see if this tx has already been executed
+            const previousResult =
+              readOnly || !workflowID ? undefined : await NodePGDSTH.#checkExecution(client, workflowID, functionNum);
+            if (previousResult) {
+              return (previousResult.output ? SuperJSON.parse(previousResult.output) : null) as Return;
+            }
+
+            // execute user's transaction function
+            const result = await asyncLocalCtx.run({ client }, async () => {
+              return (await func.call(target, ...args)) as Return;
+            });
+
+            // save the output of read/write transactions
+            if (!readOnly && workflowID) {
+              await NodePGDSTH.#recordOutput(client, workflowID, functionNum, SuperJSON.stringify(result));
+
+              // Note, existing code wraps #recordOutput call in a try/catch block that
+              // converts DB error with code 25P02 to DBOSFailedSqlTransactionError.
+              // However, existing code doesn't make any logic decisions based on that error type.
+              // DBOSFailedSqlTransactionError does stored WF ID and function name, so I assume that info is logged out somewhere
+            }
+
+            return result;
+          },
+          { isolationLevel: config?.isolationLevel, readOnly: config?.readOnly },
+        );
+        // TODO: span.setStatus({ code: SpanStatusCode.OK });
+        // TODO: this.tracer.endSpan(span);
+
+        return result;
+      } catch (error) {
+        if (isPGRetriableTransactionError(error)) {
+          // TODO: span.addEvent('TXN SERIALIZATION FAILURE', { retryWaitMillis: retryWaitMillis }, performance.now());
+          await new Promise((resolve) => setTimeout(resolve, retryWaitMS));
+          retryWaitMS = Math.min(retryWaitMS * backoffFactor, maxRetryWaitMS);
+          continue;
+        } else {
+          // TODO: span.setStatus({ code: SpanStatusCode.ERROR, message: e.message });
+          // TODO: this.tracer.endSpan(span);
+
+          // TODO: currently, we are *not* recording errors in the txOutput table.
+          // For normal execution, this is fine because we also store tx step results (output and error) in the sysdb operation output table.
+          // However, I'm concerned that we have a dueling execution hole where one tx fails while another succeeds.
+          // This implies that we can end up in a situation where the step output records an error but the txOutput table records success.
+
+          throw error;
+        }
+      }
+    }
+  }
+}
+
+export class NodePostgresDataSource implements DBOSDataSource<NodePostgresTransactionOptions> {
+  static get client(): ClientBase {
+    if (!DBOS.isInTransaction()) {
+      throw new Error('invalid use of NodePostgresDataSource.client outside of a DBOS transaction.');
+    }
+    const ctx = asyncLocalCtx.getStore();
+    if (!ctx) {
+      throw new Error('No async local context found.');
+    }
+    return ctx.client;
+  }
+
+  static async initializeInternalSchema(config: ClientConfig): Promise<void> {
+    const client = new Client(config);
+    try {
+      await client.connect();
+      await client.query(createTransactionCompletionSchemaPG);
+      await client.query(createTransactionCompletionTablePG);
+    } finally {
+      await client.end();
+    }
+  }
+
+  readonly name: string;
+  #provider: NodePGDSTH;
+
+  constructor(name: string, config: PoolConfig) {
+    this.name = name;
+    this.#provider = new NodePGDSTH(name, config);
+    registerDataSource(this.#provider);
+  }
+
+  async runTransaction<T>(callback: () => Promise<T>, funcName: string, config?: NodePostgresTransactionOptions) {
+    return await runTransaction(callback, funcName, { dsName: this.name, config });
+  }
+
+  registerTransaction<This, Args extends unknown[], Return>(
+    func: (this: This, ...args: Args) => Promise<Return>,
+    name: string,
+    config?: NodePostgresTransactionOptions,
+  ): (this: This, ...args: Args) => Promise<Return> {
+    return registerTransaction(this.name, func, { name }, config);
+  }
+
+  transaction(config?: NodePostgresTransactionOptions) {
+    // eslint-disable-next-line @typescript-eslint/no-this-alias
+    const ds = this;
+    return function decorator<This, Args extends unknown[], Return>(
+      _target: object,
+      propertyKey: string,
+      descriptor: TypedPropertyDescriptor<(this: This, ...args: Args) => Promise<Return>>,
+    ) {
+      if (!descriptor.value) {
+        throw Error('Use of decorator when original method is undefined');
+      }
+
+      descriptor.value = ds.registerTransaction(descriptor.value, propertyKey.toString(), config);
+
+      return descriptor;
+    };
+  }
+}

package/jest.config.js

@@ -0,0 +1,8 @@
+/** @type {import('ts-jest').JestConfigWithTsJest} */
+module.exports = {
+  preset: 'ts-jest',
+  testEnvironment: 'node',
+  testRegex: '((\\.|/)(test|spec))\\.ts?$',
+  moduleFileExtensions: ['ts', 'tsx', 'js', 'jsx', 'json', 'node'],
+  modulePaths: ['./'],
+};

package/package.json

@@ -0,0 +1,32 @@
+{
+  "name": "@dbos-inc/node-pg-datasource",
+  "version": "3.0.6-preview",
+  "description": "",
+  "license": "MIT",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/dbos-inc/dbos-transact-ts",
+    "directory": "packages/knex-datasource"
+  },
+  "homepage": "https://docs.dbos.dev/",
+  "main": "index.js",
+  "scripts": {
+    "build": "tsc --project tsconfig.json",
+    "test": "jest --detectOpenHandles"
+  },
+  "dependencies": {
+    "pg": "^8.11.3",
+    "pg-protocol": "^1.6.1",
+    "superjson": "^1.13"
+  },
+  "peerDependencies": {
+    "@dbos-inc/dbos-sdk": "*"
+  },
+  "devDependencies": {
+    "@types/jest": "^29.5.14",
+    "@types/node": "^20.11.25",
+    "@types/pg": "^8.11.2",
+    "jest": "^29.7.0",
+    "typescript": "^5.4.5"
+  }
+}

package/tests/config.test.ts

@@ -0,0 +1,31 @@
+import { Client } from 'pg';
+import { NodePostgresDataSource } from '../index';
+import { dropDB, ensureDB } from './test-helpers';
+
+describe('NodePostgresDataSource.configure', () => {
+  const config = { user: 'postgres', database: 'nodepg_ds_config_test' };
+
+  beforeAll(async () => {
+    const client = new Client({ ...config, database: 'postgres' });
+    try {
+      await client.connect();
+      await dropDB(client, config.database);
+      await ensureDB(client, config.database);
+    } finally {
+      await client.end();
+    }
+  });
+
+  test('configure creates tx outputs table', async () => {
+    await NodePostgresDataSource.initializeInternalSchema(config);
+
+    const client = new Client(config);
+    try {
+      await client.connect();
+      const result = await client.query('SELECT workflow_id, function_num, output FROM dbos.transaction_completion');
+      expect(result.rows.length).toBe(0);
+    } finally {
+      await client.end();
+    }
+  });
+});

package/tests/datasource.test.ts

@@ -0,0 +1,320 @@
+import { DBOS } from '@dbos-inc/dbos-sdk';
+import { Client, Pool } from 'pg';
+import { NodePostgresDataSource } from '..';
+import { dropDB, ensureDB } from './test-helpers';
+import { randomUUID } from 'crypto';
+import { SuperJSON } from 'superjson';
+
+const config = { user: 'postgres', database: 'nodepg_ds_test_userdb' };
+const dataSource = new NodePostgresDataSource('app-db', config);
+
+interface transaction_completion {
+  workflow_id: string;
+  function_num: number;
+  output: string | null;
+}
+
+describe('NodePostgresDataSource', () => {
+  const userDB = new Pool(config);
+
+  beforeAll(async () => {
+    {
+      const client = new Client({ ...config, database: 'postgres' });
+      try {
+        await client.connect();
+        await dropDB(client, 'knex_ds_test');
+        await dropDB(client, 'knex_ds_test_dbos_sys');
+        await dropDB(client, config.database);
+        await ensureDB(client, config.database);
+      } finally {
+        await client.end();
+      }
+    }
+
+    {
+      const client = await userDB.connect();
+      try {
+        await client.query(
+          'CREATE TABLE greetings(name text NOT NULL, greet_count integer DEFAULT 0, PRIMARY KEY(name))',
+        );
+      } finally {
+        client.release();
+      }
+    }
+
+    await NodePostgresDataSource.initializeInternalSchema(config);
+    DBOS.setConfig({ name: 'knex-ds-test' });
+    await DBOS.launch();
+  });
+
+  afterAll(async () => {
+    await DBOS.shutdown();
+    await userDB.end();
+  });
+
+  test('insert dataSource.register function', async () => {
+    const user = 'helloTest1';
+
+    await userDB.query('DELETE FROM greetings WHERE name = $1', [user]);
+    const workflowID = randomUUID();
+
+    await expect(DBOS.withNextWorkflowID(workflowID, () => regInsertWorfklowReg(user))).resolves.toEqual({
+      user,
+      greet_count: 1,
+    });
+
+    const { rows } = await userDB.query<transaction_completion>(
+      'SELECT * FROM dbos.transaction_completion WHERE workflow_id = $1',
+      [workflowID],
+    );
+    expect(rows.length).toBe(1);
+    expect(rows[0].workflow_id).toBe(workflowID);
+    expect(rows[0].function_num).toBe(0);
+    expect(rows[0].output).not.toBeNull();
+    expect(SuperJSON.parse(rows[0].output!)).toEqual({ user, greet_count: 1 });
+  });
+
+  test('insert dataSource.runAsTx function', async () => {
+    const user = 'helloTest2';
+
+    await userDB.query('DELETE FROM greetings WHERE name = $1', [user]);
+    const workflowID = randomUUID();
+
+    await expect(DBOS.withNextWorkflowID(workflowID, () => regInsertWorfklowRunTx(user))).resolves.toEqual({
+      user,
+      greet_count: 1,
+    });
+
+    const { rows } = await userDB.query<transaction_completion>(
+      'SELECT * FROM dbos.transaction_completion WHERE workflow_id = $1',
+      [workflowID],
+    );
+    expect(rows.length).toBe(1);
+    expect(rows[0].workflow_id).toBe(workflowID);
+    expect(rows[0].function_num).toBe(0);
+    expect(rows[0].output).not.toBeNull();
+    expect(SuperJSON.parse(rows[0].output!)).toEqual({ user, greet_count: 1 });
+  });
+
+  test('error dataSource.register function', async () => {
+    const user = 'errorTest1';
+
+    await userDB.query('DELETE FROM greetings WHERE name = $1', [user]);
+    await userDB.query('INSERT INTO greetings("name","greet_count") VALUES($1,10);', [user]);
+    const workflowID = randomUUID();
+
+    await expect(DBOS.withNextWorkflowID(workflowID, () => regErrorWorkflowReg(user))).rejects.toThrow('test error');
+
+    const { rows: txOutput } = await userDB.query('SELECT * FROM dbos.transaction_completion WHERE workflow_id = $1', [
+      workflowID,
+    ]);
+    expect(txOutput.length).toBe(0);
+
+    const { rows } = await userDB.query<greetings>('SELECT * FROM greetings WHERE name = $1', [user]);
+    expect(rows.length).toBe(1);
+    expect(rows[0].greet_count).toBe(10);
+  });
+
+  test('error dataSource.runAsTx function', async () => {
+    const user = 'errorTest2';
+
+    await userDB.query('DELETE FROM greetings WHERE name = $1', [user]);
+    await userDB.query('INSERT INTO greetings("name","greet_count") VALUES($1,10);', [user]);
+    const workflowID = randomUUID();
+
+    await expect(DBOS.withNextWorkflowID(workflowID, () => regErrorWorkflowRunTx(user))).rejects.toThrow('test error');
+
+    const { rows: txOutput } = await userDB.query<transaction_completion>(
+      'SELECT * FROM dbos.transaction_completion WHERE workflow_id = $1',
+      [workflowID],
+    );
+    expect(txOutput.length).toBe(0);
+
+    const { rows } = await userDB.query<greetings>('SELECT * FROM greetings WHERE name = $1', [user]);
+    expect(rows.length).toBe(1);
+    expect(rows[0].greet_count).toBe(10);
+  });
+
+  test('readonly dataSource.register function', async () => {
+    const user = 'readTest1';
+
+    await userDB.query('DELETE FROM greetings WHERE name = $1', [user]);
+    await userDB.query('INSERT INTO greetings("name","greet_count") VALUES($1,10);', [user]);
+
+    const workflowID = randomUUID();
+    await expect(DBOS.withNextWorkflowID(workflowID, () => regReadWorkflowReg(user))).resolves.toEqual({
+      user,
+      greet_count: 10,
+    });
+
+    const { rows } = await userDB.query<transaction_completion>(
+      'SELECT * FROM dbos.transaction_completion WHERE workflow_id = $1',
+      [workflowID],
+    );
+    expect(rows.length).toBe(0);
+  });
+
+  test('readonly dataSource.runAsTx function', async () => {
+    const user = 'readTest2';
+
+    await userDB.query('DELETE FROM greetings WHERE name = $1', [user]);
+    await userDB.query('INSERT INTO greetings("name","greet_count") VALUES($1,10);', [user]);
+
+    const workflowID = randomUUID();
+    await expect(DBOS.withNextWorkflowID(workflowID, () => regReadWorkflowRunTx(user))).resolves.toEqual({
+      user,
+      greet_count: 10,
+    });
+
+    const { rows } = await userDB.query<transaction_completion>(
+      'SELECT * FROM dbos.transaction_completion WHERE workflow_id = $1',
+      [workflowID],
+    );
+    expect(rows.length).toBe(0);
+  });
+
+  test('static dataSource.register methods', async () => {
+    const user = 'staticTest1';
+
+    await userDB.query('DELETE FROM greetings WHERE name = $1', [user]);
+
+    const workflowID = randomUUID();
+    await expect(DBOS.withNextWorkflowID(workflowID, () => regStaticWorkflow(user))).resolves.toEqual([
+      { user, greet_count: 1 },
+      { user, greet_count: 1 },
+    ]);
+  });
+
+  test('instance dataSource.register methods', async () => {
+    const user = 'instanceTest1';
+
+    await userDB.query('DELETE FROM greetings WHERE name = $1', [user]);
+
+    const workflowID = randomUUID();
+    await expect(DBOS.withNextWorkflowID(workflowID, () => regInstanceWorkflow(user))).resolves.toEqual([
+      { user, greet_count: 1 },
+      { user, greet_count: 1 },
+    ]);
+  });
+});
+
+export interface greetings {
+  name: string;
+  greet_count: number;
+}
+
+async function insertFunction(user: string) {
+  const { rows } = await NodePostgresDataSource.client.query<Pick<greetings, 'greet_count'>>(
+    `
+    INSERT INTO greetings(name, greet_count) 
+    VALUES($1, 1) 
+    ON CONFLICT(name)
+    DO UPDATE SET greet_count = greetings.greet_count + 1
+    RETURNING greet_count`,
+    [user],
+  );
+  const row = rows.length > 0 ? rows[0] : undefined;
+  return { user, greet_count: row?.greet_count };
+}
+
+async function errorFunction(user: string) {
+  const result = await insertFunction(user);
+  throw new Error('test error');
+  return result;
+}
+
+async function readFunction(user: string) {
+  const { rows } = await NodePostgresDataSource.client.query<Pick<greetings, 'greet_count'>>(
+    `
+    SELECT greet_count
+    FROM greetings
+    WHERE name = $1`,
+    [user],
+  );
+  const row = rows.length > 0 ? rows[0] : undefined;
+  return { user, greet_count: row?.greet_count };
+}
+
+const regInsertFunction = dataSource.registerTransaction(insertFunction, 'insertFunction');
+const regErrorFunction = dataSource.registerTransaction(errorFunction, 'errorFunction');
+const regReadFunction = dataSource.registerTransaction(readFunction, 'readFunction', { readOnly: true });
+
+class StaticClass {
+  static async insertFunction(user: string) {
+    return await insertFunction(user);
+  }
+
+  static async readFunction(user: string) {
+    return await readFunction(user);
+  }
+}
+
+StaticClass.insertFunction = dataSource.registerTransaction(StaticClass.insertFunction, 'insertFunction');
+StaticClass.readFunction = dataSource.registerTransaction(StaticClass.readFunction, 'readFunction');
+
+class InstanceClass {
+  async insertFunction(user: string) {
+    return await insertFunction(user);
+  }
+
+  async readFunction(user: string) {
+    return await readFunction(user);
+  }
+}
+
+InstanceClass.prototype.insertFunction = dataSource.registerTransaction(
+  // eslint-disable-next-line @typescript-eslint/unbound-method
+  InstanceClass.prototype.insertFunction,
+  'insertFunction',
+);
+InstanceClass.prototype.readFunction = dataSource.registerTransaction(
+  // eslint-disable-next-line @typescript-eslint/unbound-method
+  InstanceClass.prototype.readFunction,
+  'readFunction',
+);
+
+async function insertWorkflowReg(user: string) {
+  return await regInsertFunction(user);
+}
+
+async function insertWorkflowRunTx(user: string) {
+  return await dataSource.runTransaction(() => insertFunction(user), 'insertFunction');
+}
+
+async function errorWorkflowReg(user: string) {
+  return await regErrorFunction(user);
+}
+
+async function errorWorkflowRunTx(user: string) {
+  return await dataSource.runTransaction(() => errorFunction(user), 'errorFunction');
+}
+
+async function readWorkflowReg(user: string) {
+  return await regReadFunction(user);
+}
+
+async function readWorkflowRunTx(user: string) {
+  return await dataSource.runTransaction(() => readFunction(user), 'readFunction', { readOnly: true });
+}
+
+async function staticWorkflow(user: string) {
+  const result = await StaticClass.insertFunction(user);
+  const readResult = await StaticClass.readFunction(user);
+  return [result, readResult];
+}
+
+async function instanceWorkflow(user: string) {
+  const instance = new InstanceClass();
+  const result = await instance.insertFunction(user);
+  const readResult = await instance.readFunction(user);
+  return [result, readResult];
+}
+
+const regInsertWorfklowReg = DBOS.registerWorkflow(insertWorkflowReg, 'insertWorkflowReg');
+const regInsertWorfklowRunTx = DBOS.registerWorkflow(insertWorkflowRunTx, 'insertWorkflowRunTx');
+const regErrorWorkflowReg = DBOS.registerWorkflow(errorWorkflowReg, 'errorWorkflowReg');
+const regErrorWorkflowRunTx = DBOS.registerWorkflow(errorWorkflowRunTx, 'errorWorkflowRunTx');
+const regReadWorkflowReg = DBOS.registerWorkflow(readWorkflowReg, 'readWorkflowReg');
+const regReadWorkflowRunTx = DBOS.registerWorkflow(readWorkflowRunTx, 'readWorkflowRunTx');
+const regStaticWorkflow = DBOS.registerWorkflow(staticWorkflow, 'staticWorkflow');
+const regInstanceWorkflow = DBOS.registerWorkflow(instanceWorkflow, 'instanceWorkflow');

package/tests/test-helpers.ts

@@ -0,0 +1,12 @@
+import { Client } from 'pg';
+
+export async function ensureDB(client: Client, name: string) {
+  const results = await client.query('SELECT 1 FROM pg_database WHERE datname = $1', [name]);
+  if (results.rows.length === 0) {
+    await client.query(`CREATE DATABASE ${name}`);
+  }
+}
+
+export async function dropDB(client: Client, name: string) {
+  await client.query(`DROP DATABASE IF EXISTS ${name}`);
+}

package/tsconfig.json

@@ -0,0 +1,9 @@
+/* Visit https://aka.ms/tsconfig to read more about this file */
+{
+  "extends": "../../tsconfig.shared.json",
+  "compilerOptions": {
+    "outDir": "./dist"
+  },
+  "include": ["index.ts"],
+  "exclude": ["dist", "tests"]
+}