@dbos-inc/drizzle-datasource 3.0.7-preview → 3.0.8-preview

package/index.ts

@@ -0,0 +1,266 @@
+import { Client, ClientConfig, Pool, PoolConfig } from 'pg';
+import { DBOS, DBOSWorkflowConflictError } from '@dbos-inc/dbos-sdk';
+import {
+  type DataSourceTransactionHandler,
+  createTransactionCompletionSchemaPG,
+  createTransactionCompletionTablePG,
+  isPGRetriableTransactionError,
+  isPGKeyConflictError,
+  registerTransaction,
+  runTransaction,
+  DBOSDataSource,
+  registerDataSource,
+} from '@dbos-inc/dbos-sdk/datasource';
+import { drizzle, NodePgDatabase } from 'drizzle-orm/node-postgres';
+import { AsyncLocalStorage } from 'async_hooks';
+import { SuperJSON } from 'superjson';
+import { PgTransactionConfig } from 'drizzle-orm/pg-core';
+import { sql } from 'drizzle-orm';
+
+interface DrizzleLocalCtx {
+  client: NodePgDatabase<{ [key: string]: object }>;
+}
+
+export type TransactionConfig = Pick<PgTransactionConfig, 'isolationLevel' | 'accessMode'>;
+
+const asyncLocalCtx = new AsyncLocalStorage<DrizzleLocalCtx>();
+
+export interface transaction_completion {
+  workflow_id: string;
+  function_num: number;
+  output: string | null;
+  error: string | null;
+}
+
+interface DrizzleConnection {
+  readonly db: NodePgDatabase<{ [key: string]: object }>;
+  end(): Promise<void>;
+}
+
+class DrizzleTransactionHandler implements DataSourceTransactionHandler {
+  readonly dsType = 'drizzle';
+  #connection: DrizzleConnection | undefined;
+
+  constructor(
+    readonly name: string,
+    private readonly config: PoolConfig,
+    private readonly entities: { [key: string]: object } = {},
+  ) {}
+
+  async initialize(): Promise<void> {
+    const conn = this.#connection;
+
+    const driver = new Pool(this.config);
+    const db = drizzle(driver, { schema: this.entities });
+    this.#connection = { db, end: () => driver.end() };
+
+    await conn?.end();
+  }
+
+  async destroy(): Promise<void> {
+    const conn = this.#connection;
+
+    this.#connection = undefined;
+
+    await conn?.end();
+  }
+
+  get #drizzle(): NodePgDatabase<{ [key: string]: object }> {
+    if (!this.#connection) {
+      throw new Error(`DataSource ${this.name} is not initialized.`);
+    }
+    return this.#connection.db;
+  }
+
+  async #checkExecution(
+    workflowID: string,
+    stepID: number,
+  ): Promise<{ output: string | null } | { error: string } | undefined> {
+    type Result = { output: string | null; error: string | null };
+
+    const statement = sql`
+        SELECT output, error FROM dbos.transaction_completion
+        WHERE workflow_id = ${workflowID} AND function_num = ${stepID}`;
+    const result = await this.#drizzle.execute<Result>(statement);
+
+    if (result.rows.length !== 1) {
+      return undefined;
+    }
+
+    const { output, error } = result.rows[0];
+    return error !== null ? { error } : { output };
+  }
+
+  static async #recordOutput(
+    client: NodePgDatabase<{ [key: string]: object }>,
+    workflowID: string,
+    stepID: number,
+    output: string,
+  ): Promise<void> {
+    try {
+      const statement = sql`
+        INSERT INTO dbos.transaction_completion (workflow_id, function_num, output) 
+        VALUES (${workflowID}, ${stepID}, ${output})`;
+      await client.execute(statement);
+    } catch (error) {
+      if (isPGKeyConflictError(error)) {
+        throw new DBOSWorkflowConflictError(workflowID);
+      } else {
+        throw error;
+      }
+    }
+  }
+
+  async #recordError(workflowID: string, stepID: number, error: string): Promise<void> {
+    try {
+      const statement = sql`
+        INSERT INTO dbos.transaction_completion (workflow_id, function_num, error) 
+        VALUES (${workflowID}, ${stepID}, ${error})`;
+      await this.#drizzle.execute(statement);
+    } catch (error) {
+      if (isPGKeyConflictError(error)) {
+        throw new DBOSWorkflowConflictError(workflowID);
+      } else {
+        throw error;
+      }
+    }
+  }
+
+  /* Invoke a transaction function, called by the framework */
+  async invokeTransactionFunction<This, Args extends unknown[], Return>(
+    config: TransactionConfig | undefined,
+    target: This,
+    func: (this: This, ...args: Args) => Promise<Return>,
+    ...args: Args
+  ): Promise<Return> {
+    const workflowID = DBOS.workflowID;
+    const stepID = DBOS.stepID;
+    if (workflowID !== undefined && stepID === undefined) {
+      throw new Error('DBOS.stepID is undefined inside a workflow.');
+    }
+
+    const readOnly = config?.accessMode === 'read only' ? true : false;
+    const saveResults = !readOnly && workflowID !== undefined;
+
+    // Retry loop if appropriate
+    let retryWaitMS = 1;
+    const backoffFactor = 1.5;
+    const maxRetryWaitMS = 2000; // Maximum wait 2 seconds.
+
+    while (true) {
+      // 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;
+      }
+
+      try {
+        const result = await this.#drizzle.transaction(
+          async (client) => {
+            // execute user's transaction function
+            const result = await asyncLocalCtx.run({ client }, async () => {
+              return await func.call(target, ...args);
+            });
+
+            // save the output of read/write transactions
+            if (saveResults) {
+              await DrizzleTransactionHandler.#recordOutput(client, workflowID, stepID!, SuperJSON.stringify(result));
+            }
+
+            return result;
+          },
+          { accessMode: config?.accessMode, isolationLevel: config?.isolationLevel },
+        );
+
+        return result;
+      } catch (error) {
+        if (isPGRetriableTransactionError(error)) {
+          DBOS.span?.addEvent('TXN SERIALIZATION FAILURE', { retryWaitMillis: retryWaitMS }, performance.now());
+          // Retry serialization failures.
+          await new Promise((resolve) => setTimeout(resolve, retryWaitMS));
+          retryWaitMS = Math.min(retryWaitMS * backoffFactor, maxRetryWaitMS);
+          continue;
+        } else {
+          if (saveResults) {
+            const message = SuperJSON.stringify(error);
+            await this.#recordError(workflowID, stepID!, message);
+          }
+
+          throw error;
+        }
+      }
+    }
+  }
+}
+
+export class DrizzleDataSource implements DBOSDataSource<TransactionConfig> {
+  // User calls this... DBOS not directly involved...
+  static get client(): NodePgDatabase<{ [key: string]: object }> {
+    if (!DBOS.isInTransaction()) {
+      throw new Error('Invalid use of DrizzleDataSource.client outside of a DBOS transaction');
+    }
+    const ctx = asyncLocalCtx.getStore();
+    if (!ctx) {
+      throw new Error('Invalid use of DrizzleDataSource.client outside of a DBOS transaction');
+    }
+    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();
+    }
+  }
+
+  #provider: DrizzleTransactionHandler;
+
+  constructor(
+    readonly name: string,
+    config: PoolConfig,
+    entities: { [key: string]: object } = {},
+  ) {
+    this.#provider = new DrizzleTransactionHandler(name, config, entities);
+    registerDataSource(this.#provider);
+  }
+
+  async runTransaction<T>(callback: () => Promise<T>, funcName: string, config?: TransactionConfig) {
+    return await runTransaction(callback, funcName, { dsName: this.name, config });
+  }
+
+  registerTransaction<This, Args extends unknown[], Return>(
+    func: (this: This, ...args: Args) => Promise<Return>,
+    config?: TransactionConfig,
+    name?: string,
+  ): (this: This, ...args: Args) => Promise<Return> {
+    return registerTransaction(this.name, func, { name: name ?? func.name }, config);
+  }
+
+  // decorator
+  transaction(config?: TransactionConfig) {
+    // eslint-disable-next-line @typescript-eslint/no-this-alias
+    const ds = this;
+    return function decorator<This, Args extends unknown[], Return>(
+      _target: object,
+      propertyKey: PropertyKey,
+      descriptor: TypedPropertyDescriptor<(this: This, ...args: Args) => Promise<Return>>,
+    ) {
+      if (!descriptor.value) {
+        throw new Error('Use of decorator when original method is undefined');
+      }
+
+      descriptor.value = ds.registerTransaction(descriptor.value, config, String(propertyKey));
+
+      return descriptor;
+    };
+  }
+}

package/package.json

@@ -1,8 +1,11 @@
 {
   "name": "@dbos-inc/drizzle-datasource",
-  "version": "3.0.7-preview",
-  "main": "dist/src/index.js",
-  "types": "dist/src/index.d.ts",
+  "version": "3.0.8-preview",
+  "description": "DBOS DataSource library for Drizzle ORM with PostgreSQL support",
+  "license": "MIT",
+  "main": "dist/index.js",
+  "types": "dist/index.d.ts",
+  "homepage": "https://docs.dbos.dev/",
   "repository": {
     "type": "git",
     "url": "https://github.com/dbos-inc/dbos-transact-ts",
@@ -10,14 +13,9 @@
   },
   "scripts": {
     "build": "tsc --project tsconfig.json",
-    "test": "npm run build && jest --detectOpenHandles"
+    "test": "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"
@@ -26,12 +24,10 @@
     "@dbos-inc/dbos-sdk": "*"
   },
   "devDependencies": {
-    "@types/jest": "^29.5.12",
-    "@types/node": "^20.11.25",
-    "@types/supertest": "^6.0.2",
+    "@types/jest": "^29.5.14",
+    "@types/pg": "^8.15.2",
+    "drizzle-kit": "^0.31.1",
     "jest": "^29.7.0",
-    "supertest": "^7.0.0",
-    "ts-jest": "^29.1.4",
-    "typescript": "^5.3.3"
+    "typescript": "^5.4.5"
   }
 }

package/tests/config.test.ts

@@ -0,0 +1,31 @@
+import { Client } from 'pg';
+import { DrizzleDataSource } from '../index';
+import { dropDB, ensureDB } from './test-helpers';
+
+describe('DrizzleDataSource.configure', () => {
+  const config = { user: 'postgres', database: 'drizzle_ds_config_test' };
+
+  beforeAll(async () => {
+    const client = new Client({ ...config, database: 'postgres' });
+    try {
+      await client.connect();
+      await dropDB(client, config.database, true);
+      await ensureDB(client, config.database);
+    } finally {
+      await client.end();
+    }
+  });
+
+  test('configure creates tx outputs table', async () => {
+    await DrizzleDataSource.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();
+    }
+  });
+});