@dbos-inc/drizzle-datasource 3.0.6-preview

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,37 @@
+{
+  "name": "@dbos-inc/drizzle-datasource",
+  "version": "3.0.6-preview",
+  "main": "dist/src/index.js",
+  "types": "dist/src/index.d.ts",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/dbos-inc/dbos-transact-ts",
+    "directory": "packages/drizzle-datasource"
+  },
+  "scripts": {
+    "build": "tsc --project tsconfig.json",
+    "test": "npm run build && jest --detectOpenHandles"
+  },
+  "keywords": [],
+  "author": "",
+  "license": "MIT",
+  "description": "",
+  "dependencies": {
+    "drizzle-kit": "^0.31.1",
+    "drizzle-orm": "^0.40.1",
+    "pg": "^8.11.3",
+    "superjson": "^1.13"
+  },
+  "peerDependencies": {
+    "@dbos-inc/dbos-sdk": "*"
+  },
+  "devDependencies": {
+    "@types/jest": "^29.5.12",
+    "@types/node": "^20.11.25",
+    "@types/supertest": "^6.0.2",
+    "jest": "^29.7.0",
+    "supertest": "^7.0.0",
+    "ts-jest": "^29.1.4",
+    "typescript": "^5.3.3"
+  }
+}

package/src/drizzle_datasource.ts

@@ -0,0 +1,332 @@
+import { Pool, PoolConfig } from 'pg';
+import { DBOS, Error } from '@dbos-inc/dbos-sdk';
+import {
+  type DataSourceTransactionHandler,
+  createTransactionCompletionSchemaPG,
+  createTransactionCompletionTablePG,
+  isPGRetriableTransactionError,
+  isPGKeyConflictError,
+  isPGFailedSqlTransactionError,
+  registerTransaction,
+  runTransaction,
+  PGIsolationLevel as IsolationLevel,
+  PGTransactionConfig as DrizzleTransactionConfig,
+  DBOSDataSource,
+  registerDataSource,
+} from '@dbos-inc/dbos-sdk/datasource';
+import { drizzle, NodePgDatabase } from 'drizzle-orm/node-postgres';
+import { pushSchema } from 'drizzle-kit/api';
+import { AsyncLocalStorage } from 'async_hooks';
+import { SuperJSON } from 'superjson';
+
+export { IsolationLevel, DrizzleTransactionConfig };
+
+interface DrizzleLocalCtx {
+  drizzleClient: NodePgDatabase<{ [key: string]: object }>;
+}
+const asyncLocalCtx = new AsyncLocalStorage<DrizzleLocalCtx>();
+
+function getCurrentDSContextStore(): DrizzleLocalCtx | undefined {
+  return asyncLocalCtx.getStore();
+}
+
+function assertCurrentDSContextStore(): DrizzleLocalCtx {
+  const ctx = getCurrentDSContextStore();
+  if (!ctx) throw new Error.DBOSInvalidWorkflowTransitionError('Invalid use of TypeOrmDs outside of a `transaction`');
+  return ctx;
+}
+
+export interface transaction_completion {
+  workflow_id: string;
+  function_num: number;
+  output: string | null;
+  error: string | null;
+}
+
+class DrizzleDSTH implements DataSourceTransactionHandler {
+  readonly dsType = 'drizzle';
+  dataSource: NodePgDatabase<{ [key: string]: object }> | undefined;
+  drizzlePool: Pool | undefined;
+
+  constructor(
+    readonly name: string,
+    readonly config: PoolConfig,
+    readonly entities: { [key: string]: object } = {},
+  ) {}
+
+  async initialize(): Promise<void> {
+    this.drizzlePool = new Pool(this.config);
+    this.dataSource = drizzle(this.drizzlePool, { schema: this.entities });
+
+    return Promise.resolve();
+  }
+
+  async destroy(): Promise<void> {
+    await this.drizzlePool?.end();
+  }
+
+  async #checkExecution<R>(
+    client: Pool,
+    workflowID: string,
+    funcNum: number,
+  ): Promise<
+    | {
+        res: R;
+      }
+    | undefined
+  > {
+    const result = await client.query<{
+      output: string;
+    }>(
+      `SELECT output
+       FROM dbos.transaction_completion
+        WHERE workflow_id = $1 AND function_num = $2`,
+      [workflowID, funcNum],
+    );
+
+    if (result.rows.length !== 1) {
+      return undefined;
+    }
+
+    return { res: SuperJSON.parse(result.rows[0].output) };
+  }
+
+  async #recordOutput<R>(client: Pool, workflowID: string, funcNum: number, output: R): Promise<void> {
+    const serialOutput = SuperJSON.stringify(output);
+    await client.query<{ rows: transaction_completion[] }>(
+      `INSERT INTO dbos.transaction_completion (
+        workflow_id, 
+        function_num,
+        output,
+        created_at
+      ) VALUES ($1, $2, $3, $4)`,
+      [workflowID, funcNum, serialOutput, Date.now()],
+    );
+  }
+
+  async #recordError<R>(client: Pool, workflowID: string, funcNum: number, error: R): Promise<void> {
+    const serialError = SuperJSON.stringify(error);
+    await client.query<{ rows: transaction_completion[] }>(
+      `INSERT INTO dbos.transaction_completion (
+        workflow_id, 
+        function_num,
+        error,
+        created_at
+      ) VALUES ($1, $2, $3, $4)`,
+      [workflowID, funcNum, serialError, Date.now()],
+    );
+  }
+
+  /* Invoke a transaction function, called by the framework */
+  async invokeTransactionFunction<This, Args extends unknown[], Return>(
+    config: DrizzleTransactionConfig,
+    target: This,
+    func: (this: This, ...args: Args) => Promise<Return>,
+    ...args: Args
+  ): Promise<Return> {
+    let isolationLevel: 'read uncommitted' | 'read committed' | 'repeatable read' | 'serializable';
+
+    if (config === undefined || config.isolationLevel === undefined) {
+      isolationLevel = 'serializable'; // Default isolation level
+    } else if (config.isolationLevel === IsolationLevel.ReadUncommitted) {
+      isolationLevel = 'read uncommitted';
+    } else if (config.isolationLevel === IsolationLevel.ReadCommitted) {
+      isolationLevel = 'read committed';
+    } else if (config.isolationLevel === IsolationLevel.RepeatableRead) {
+      isolationLevel = 'repeatable read';
+    } else {
+      isolationLevel = 'serializable';
+    }
+
+    const accessMode = 'read write';
+
+    const readOnly = config?.readOnly ? true : false;
+
+    const wfid = DBOS.workflowID!;
+    const funcnum = DBOS.stepID!;
+    const funcname = func.name;
+
+    // Retry loop if appropriate
+    let retryWaitMillis = 1;
+    const backoffFactor = 1.5;
+    const maxRetryWaitMs = 2000; // Maximum wait 2 seconds.
+    const shouldCheckOutput = false;
+
+    if (this.drizzlePool === undefined) {
+      throw new Error.DBOSInvalidWorkflowTransitionError('Invalid use of Datasource');
+    }
+
+    if (this.dataSource === undefined) {
+      throw new Error.DBOSInvalidWorkflowTransitionError('Invalid use of Datasource');
+    }
+
+    while (true) {
+      let failedForRetriableReasons = false;
+
+      try {
+        const result = await this.dataSource.transaction(
+          async (drizzleClient: NodePgDatabase<{ [key: string]: object }>) => {
+            if (this.drizzlePool === undefined) {
+              throw new Error.DBOSInvalidWorkflowTransitionError('Invalid use of Datasource');
+            }
+
+            if (shouldCheckOutput && !readOnly && wfid) {
+              const executionResult = await this.#checkExecution<Return>(this.drizzlePool, wfid, funcnum);
+
+              if (executionResult) {
+                DBOS.span?.setAttribute('cached', true);
+                return executionResult.res;
+              }
+            }
+
+            const result = await asyncLocalCtx.run({ drizzleClient }, async () => {
+              return await func.call(target, ...args);
+            });
+
+            // Save result
+            try {
+              if (!readOnly && wfid) {
+                await this.#recordOutput(this.drizzlePool, wfid, funcnum, result);
+              }
+            } catch (e) {
+              const error = e as Error;
+              await this.#recordError(this.drizzlePool, wfid, funcnum, error);
+
+              // Aside from a connectivity error, two kinds of error are anticipated here:
+              //  1. The transaction is marked failed, but the user code did not throw.
+              //      Bad on them.  We will throw an error (this will get recorded) and not retry.
+              //  2. There was a key conflict in the statement, and we need to use the fetched output
+              if (isPGFailedSqlTransactionError(error)) {
+                DBOS.logger.error(
+                  `In workflow ${wfid}, Postgres aborted a transaction but the function '${funcname}' did not raise an exception.  Please ensure that the transaction method raises an exception if the database transaction is aborted.`,
+                );
+                failedForRetriableReasons = false;
+                throw new Error.DBOSFailedSqlTransactionError(wfid, funcname);
+              } else if (isPGKeyConflictError(error)) {
+                throw new Error.DBOSWorkflowConflictError(
+                  `In workflow ${wfid}, Postgres raised a key conflict error in transaction '${funcname}'.  This is not retriable, but the output will be fetched from the database.`,
+                );
+              } else {
+                DBOS.logger.error(`Unexpected error raised in transaction '${funcname}: ${error}`);
+                failedForRetriableReasons = false;
+                throw error;
+              }
+            }
+
+            return result;
+          },
+          { isolationLevel, accessMode },
+        );
+
+        return result;
+      } catch (e) {
+        const err = e as Error;
+        if (failedForRetriableReasons || isPGRetriableTransactionError(err)) {
+          DBOS.span?.addEvent('TXN SERIALIZATION FAILURE', { retryWaitMillis: retryWaitMillis }, performance.now());
+          // Retry serialization failures.
+          await DBOS.sleepms(retryWaitMillis);
+          retryWaitMillis *= backoffFactor;
+          retryWaitMillis = retryWaitMillis < maxRetryWaitMs ? retryWaitMillis : maxRetryWaitMs;
+          continue;
+        } else {
+          throw err;
+        }
+      }
+    }
+  }
+
+  createInstance(): NodePgDatabase<{ [key: string]: object }> {
+    const drizzlePool = new Pool(this.config);
+    const ds = drizzle(drizzlePool, { schema: this.entities });
+    return ds;
+  }
+}
+
+export class DrizzleDataSource implements DBOSDataSource<DrizzleTransactionConfig> {
+  #provider: DrizzleDSTH;
+
+  constructor(
+    readonly name: string,
+    readonly config: PoolConfig,
+    readonly entities: { [key: string]: object } = {},
+  ) {
+    this.#provider = new DrizzleDSTH(name, config, entities);
+    registerDataSource(this.#provider);
+  }
+
+  // User calls this... DBOS not directly involved...
+  static get drizzleClient(): NodePgDatabase<{ [key: string]: object }> {
+    const ctx = assertCurrentDSContextStore();
+    if (!DBOS.isInTransaction())
+      throw new Error.DBOSInvalidWorkflowTransitionError(
+        'Invalid use of `DrizzleDataSource.drizzleClient` outside of a `transaction`',
+      );
+    return ctx.drizzleClient;
+  }
+
+  get dataSource() {
+    return this.#provider.dataSource;
+  }
+
+  async initializeInternalSchema(): Promise<void> {
+    const drizzlePool = new Pool(this.config);
+    const ds = drizzle(drizzlePool, { schema: this.entities });
+
+    try {
+      await ds.execute(createTransactionCompletionSchemaPG);
+      await ds.execute(createTransactionCompletionTablePG);
+    } catch (e) {
+      const error = e as Error;
+      throw new Error.DBOSError(`Unexpected error initializing schema: ${error.message}`);
+    } finally {
+      try {
+        await drizzlePool.end();
+      } catch (e) {}
+    }
+  }
+
+  async runTransaction<T>(callback: () => Promise<T>, funcName: string, config?: DrizzleTransactionConfig) {
+    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?: DrizzleTransactionConfig,
+  ): (this: This, ...args: Args) => Promise<Return> {
+    return registerTransaction(this.name, func, { name }, config);
+  }
+
+  // decorator
+  transaction(config?: DrizzleTransactionConfig) {
+    // 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 new Error.DBOSError('Use of decorator when original method is undefined');
+      }
+
+      descriptor.value = ds.registerTransaction(descriptor.value, propertyKey.toString(), config);
+
+      return descriptor;
+    };
+  }
+
+  /**
+   * Create user schema in database (for testing)
+   */
+  async createSchema() {
+    const drizzlePool = new Pool(this.config);
+    const db = drizzle(drizzlePool);
+    try {
+      const res = await pushSchema(this.entities, db);
+      await res.apply();
+    } finally {
+      await drizzlePool.end();
+    }
+  }
+}

package/src/index.ts

@@ -0,0 +1 @@
+export { DrizzleDataSource, DrizzleTransactionConfig, IsolationLevel } from './drizzle_datasource';

package/tests/drizzle.test.ts

@@ -0,0 +1,187 @@
+/* eslint-disable */
+import { DBOS } from '@dbos-inc/dbos-sdk';
+import { DrizzleDataSource } from '../src';
+import { randomUUID } from 'node:crypto';
+import { setUpDBOSTestDb } from './testutils';
+import { pgTable, text } from 'drizzle-orm/pg-core';
+import { eq } from 'drizzle-orm/expressions';
+
+const kv = pgTable('kv', {
+  id: text('id').primaryKey().default('t'),
+  value: text('value').default('v'),
+});
+
+const dbPassword: string | undefined = process.env.DB_PASSWORD || process.env.PGPASSWORD;
+if (!dbPassword) {
+  throw new Error('DB_PASSWORD or PGPASSWORD environment variable not set');
+}
+
+const databaseUrl = `postgresql://postgres:${dbPassword}@localhost:5432/drizzle_testdb?sslmode=disable`;
+
+const poolconfig = {
+  connectionString: databaseUrl,
+  user: 'postgres',
+  password: dbPassword,
+  database: 'drizzle_testdb',
+
+  host: 'localhost',
+  port: 5432,
+};
+
+const drizzleDS = new DrizzleDataSource('app-db', poolconfig, { kv });
+
+const dbosConfig = {
+  name: 'dbos_drizzle_test',
+  databaseUrl: databaseUrl,
+  poolConfig: poolconfig,
+  system_database: 'drizzle_testdb_dbos_sys',
+  telemetry: {
+    logs: {
+      silent: true,
+    },
+  },
+};
+
+async function txFunctionGuts() {
+  expect(DBOS.isInTransaction()).toBe(true);
+  expect(DBOS.isWithinWorkflow()).toBe(true);
+  const res = await DrizzleDataSource.drizzleClient.execute("SELECT 'Tx2 result' as a");
+  return res.rows[0].a as string;
+}
+
+const txFunc = drizzleDS.registerTransaction(txFunctionGuts, 'MySecondTx', {});
+
+async function wfFunctionGuts() {
+  // Transaction variant 2: Let DBOS run a code snippet as a step
+  const p1 = await drizzleDS.runTransaction(
+    async () => {
+      return (await DrizzleDataSource.drizzleClient.execute("SELECT 'My first tx result' as a")).rows[0].a;
+    },
+    'MyFirstTx',
+    { readOnly: true },
+  );
+
+  // Transaction variant 1: Use a registered DBOS transaction function
+  const p2 = await txFunc();
+
+  return p1 + '|' + p2;
+}
+
+// Workflow functions must always be registered before launch; this
+//  allows recovery to occur.
+const wfFunction = DBOS.registerWorkflow(wfFunctionGuts, 'workflow');
+
+class DBWFI {
+  @drizzleDS.transaction({ readOnly: true })
+  static async tx(): Promise<string> {
+    let res = await DrizzleDataSource.drizzleClient.execute("SELECT 'My decorated tx result' as a");
+    return res.rows[0].a as string;
+  }
+
+  @DBOS.workflow()
+  static async wf(): Promise<string> {
+    return await DBWFI.tx();
+  }
+}
+
+describe('decoratorless-api-tests', () => {
+  beforeAll(() => {
+    DBOS.setConfig(dbosConfig);
+  });
+
+  beforeEach(async () => {
+    await setUpDBOSTestDb(dbosConfig);
+    await drizzleDS.initializeInternalSchema();
+    await drizzleDS.createSchema();
+    await DBOS.launch();
+  });
+
+  afterEach(async () => {
+    await DBOS.shutdown();
+  });
+
+  test('bare-tx-wf-functions', async () => {
+    const wfid = randomUUID();
+
+    await DBOS.withNextWorkflowID(wfid, async () => {
+      const res = await wfFunction();
+      expect(res).toBe('My first tx result|Tx2 result');
+    });
+
+    const wfsteps = (await DBOS.listWorkflowSteps(wfid))!;
+    expect(wfsteps.length).toBe(2);
+    expect(wfsteps[0].functionID).toBe(0);
+    expect(wfsteps[0].name).toBe('MyFirstTx');
+    expect(wfsteps[1].functionID).toBe(1);
+    expect(wfsteps[1].name).toBe('MySecondTx');
+  });
+
+  test('decorated-tx-wf-functions', async () => {
+    const wfid = randomUUID();
+
+    await DBOS.withNextWorkflowID(wfid, async () => {
+      const res = await DBWFI.wf();
+      expect(res).toBe('My decorated tx result');
+    });
+
+    const wfsteps = (await DBOS.listWorkflowSteps(wfid))!;
+    expect(wfsteps.length).toBe(1);
+    expect(wfsteps[0].functionID).toBe(0);
+    expect(wfsteps[0].name).toBe('tx');
+  });
+});
+
+class KVController {
+  @drizzleDS.transaction()
+  static async testTxn(id: string, value: string) {
+    await drizzleDS.dataSource?.insert(kv).values({ id: id, value: value }).onConflictDoNothing().execute();
+
+    return id;
+  }
+
+  static async readTxn(id: string): Promise<string> {
+    const kvp = await drizzleDS.dataSource?.select().from(kv).where(eq(kv.id, id)).limit(1).execute();
+
+    return kvp?.[0]?.value ?? '<Not Found>';
+  }
+
+  @DBOS.workflow()
+  static async wf(id: string, value: string) {
+    return await KVController.testTxn(id, value);
+  }
+}
+
+const txFunc2 = drizzleDS.registerTransaction(KVController.readTxn, 'explicitRegister', {});
+async function explicitWf(id: string): Promise<string> {
+  return await txFunc2(id);
+}
+const wfFunction2 = DBOS.registerWorkflow(explicitWf, 'explicitworkflow');
+
+describe('drizzle-tests', () => {
+  beforeAll(() => {
+    DBOS.setConfig(dbosConfig);
+  });
+
+  beforeEach(async () => {
+    await setUpDBOSTestDb(dbosConfig);
+    await drizzleDS.initializeInternalSchema();
+    await drizzleDS.createSchema();
+    await DBOS.launch();
+  });
+
+  afterEach(async () => {
+    await DBOS.shutdown();
+  });
+
+  test('simple-drizzle', async () => {
+    await KVController.wf('test', 'value');
+    let read = await KVController.readTxn('test');
+    expect(read).toBe('value');
+  });
+
+  test('drizzle-register', async () => {
+    await expect(wfFunction2('test')).resolves.toBe('<Not Found>');
+    await KVController.wf('test', 'value');
+    await expect(wfFunction2('test')).resolves.toBe('value');
+  });
+});

package/tests/testutils.ts

@@ -0,0 +1,30 @@
+import { Client } from 'pg';
+import { DBOSConfig } from '@dbos-inc/dbos-sdk';
+
+export async function setUpDBOSTestDb(config: DBOSConfig) {
+  const pgSystemClient = new Client({
+    user: config.poolConfig?.user,
+    port: config.poolConfig?.port,
+    host: config.poolConfig?.host,
+    password: config.poolConfig?.password,
+    database: 'postgres',
+  });
+
+  try {
+    await pgSystemClient.connect();
+    await pgSystemClient.query(`DROP DATABASE IF EXISTS ${config.poolConfig?.database};`);
+    await pgSystemClient.query(`CREATE DATABASE ${config.poolConfig?.database};`);
+    await pgSystemClient.query(`DROP DATABASE IF EXISTS ${config.system_database};`);
+    await pgSystemClient.end();
+  } catch (e) {
+    if (e instanceof AggregateError) {
+      console.error(`Test database setup failed: AggregateError containing ${e.errors.length} errors:`);
+      e.errors.forEach((err, index) => {
+        console.error(`  Error ${index + 1}:`, err);
+      });
+    } else {
+      console.error(`Test database setup failed:`, e);
+    }
+    throw e;
+  }
+}

package/tsconfig.json

@@ -0,0 +1,8 @@
+{
+  "extends": "../../tsconfig.shared.json",
+  "compilerOptions": {
+    "outDir": "./dist"
+  },
+  "exclude": ["dist", "*.test.ts", "testutils.ts"],
+  "include": ["src/"]
+}