@dbos-inc/drizzle-datasource 3.0.11-preview.gc9233b8190 → 3.0.16-preview

package/dist/index.d.ts

@@ -0,0 +1,31 @@
+import { ClientConfig, PoolConfig } from 'pg';
+import { DBOSDataSource } from '@dbos-inc/dbos-sdk/datasource';
+import { NodePgDatabase } from 'drizzle-orm/node-postgres';
+import { PgTransactionConfig } from 'drizzle-orm/pg-core';
+export type TransactionConfig = Pick<PgTransactionConfig, 'isolationLevel' | 'accessMode'> & {
+    name?: string;
+};
+export interface transaction_completion {
+    workflow_id: string;
+    function_num: number;
+    output: string | null;
+    error: string | null;
+}
+export declare class DrizzleDataSource implements DBOSDataSource<TransactionConfig> {
+    #private;
+    readonly name: string;
+    static get client(): NodePgDatabase<{
+        [key: string]: object;
+    }>;
+    get client(): NodePgDatabase<{
+        [key: string]: object;
+    }>;
+    static initializeDBOSSchema(config: ClientConfig): Promise<void>;
+    constructor(name: string, config: PoolConfig, entities?: {
+        [key: string]: object;
+    });
+    runTransaction<T>(func: () => Promise<T>, config?: TransactionConfig): Promise<T>;
+    registerTransaction<This, Args extends unknown[], Return>(func: (this: This, ...args: Args) => Promise<Return>, config?: TransactionConfig): (this: This, ...args: Args) => Promise<Return>;
+    transaction(config?: TransactionConfig): <This, Args extends unknown[], Return>(_target: object, propertyKey: PropertyKey, descriptor: TypedPropertyDescriptor<(this: This, ...args: Args) => Promise<Return>>) => TypedPropertyDescriptor<(this: This, ...args: Args) => Promise<Return>>;
+}
+//# sourceMappingURL=index.d.ts.map

package/dist/index.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../index.ts"],"names":[],"mappings":"AAAA,OAAO,EAAU,YAAY,EAAQ,UAAU,EAAE,MAAM,IAAI,CAAC;AAE5D,OAAO,EAQL,cAAc,EAEf,MAAM,+BAA+B,CAAC;AACvC,OAAO,EAAW,cAAc,EAAE,MAAM,2BAA2B,CAAC;AAGpE,OAAO,EAAE,mBAAmB,EAAE,MAAM,qBAAqB,CAAC;AAQ1D,MAAM,MAAM,iBAAiB,GAAG,IAAI,CAAC,mBAAmB,EAAE,gBAAgB,GAAG,YAAY,CAAC,GAAG;IAAE,IAAI,CAAC,EAAE,MAAM,CAAA;CAAE,CAAC;AAI/G,MAAM,WAAW,sBAAsB;IACrC,WAAW,EAAE,MAAM,CAAC;IACpB,YAAY,EAAE,MAAM,CAAC;IACrB,MAAM,EAAE,MAAM,GAAG,IAAI,CAAC;IACtB,KAAK,EAAE,MAAM,GAAG,IAAI,CAAC;CACtB;AAwKD,qBAAa,iBAAkB,YAAW,cAAc,CAAC,iBAAiB,CAAC;;IAsCvE,QAAQ,CAAC,IAAI,EAAE,MAAM;IAtBvB,MAAM,KAAK,MAAM;;OAEhB;IAED,IAAI,MAAM;;OAET;WAEY,oBAAoB,CAAC,MAAM,EAAE,YAAY,GAAG,OAAO,CAAC,IAAI,CAAC;gBAc3D,IAAI,EAAE,MAAM,EACrB,MAAM,EAAE,UAAU,EAClB,QAAQ,GAAE;QAAE,CAAC,GAAG,EAAE,MAAM,GAAG,MAAM,CAAA;KAAO;IAMpC,cAAc,CAAC,CAAC,EAAE,IAAI,EAAE,MAAM,OAAO,CAAC,CAAC,CAAC,EAAE,MAAM,CAAC,EAAE,iBAAiB;IAI1E,mBAAmB,CAAC,IAAI,EAAE,IAAI,SAAS,OAAO,EAAE,EAAE,MAAM,EACtD,IAAI,EAAE,CAAC,IAAI,EAAE,IAAI,EAAE,GAAG,IAAI,EAAE,IAAI,KAAK,OAAO,CAAC,MAAM,CAAC,EACpD,MAAM,CAAC,EAAE,iBAAiB,GACzB,CAAC,IAAI,EAAE,IAAI,EAAE,GAAG,IAAI,EAAE,IAAI,KAAK,OAAO,CAAC,MAAM,CAAC;IAKjD,WAAW,CAAC,MAAM,CAAC,EAAE,iBAAiB,mDAIzB,MAAM,eACF,WAAW,cACZ,wBAAwB,CAAC,IAAI,EAAE,IAAI,EAAE,GAAG,IAAI,EAAE,IAAI,KAAK,QAAQ,MAAM,CAAC,CAAC,oCAAxC,IAAI,WAAW,IAAI,KAAK,QAAQ,MAAM,CAAC;CAcvF"}

package/dist/index.js

@@ -0,0 +1,202 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.DrizzleDataSource = void 0;
+const pg_1 = require("pg");
+const dbos_sdk_1 = require("@dbos-inc/dbos-sdk");
+const datasource_1 = require("@dbos-inc/dbos-sdk/datasource");
+const node_postgres_1 = require("drizzle-orm/node-postgres");
+const async_hooks_1 = require("async_hooks");
+const superjson_1 = require("superjson");
+const drizzle_orm_1 = require("drizzle-orm");
+const asyncLocalCtx = new async_hooks_1.AsyncLocalStorage();
+class DrizzleTransactionHandler {
+    name;
+    config;
+    entities;
+    dsType = 'drizzle';
+    #connection;
+    constructor(name, config, entities = {}) {
+        this.name = name;
+        this.config = config;
+        this.entities = entities;
+    }
+    async initialize() {
+        const conn = this.#connection;
+        const driver = new pg_1.Pool(this.config);
+        const db = (0, node_postgres_1.drizzle)(driver, { schema: this.entities });
+        this.#connection = { db, end: () => driver.end() };
+        await conn?.end();
+    }
+    async destroy() {
+        const conn = this.#connection;
+        this.#connection = undefined;
+        await conn?.end();
+    }
+    get #drizzle() {
+        if (!this.#connection) {
+            throw new Error(`DataSource ${this.name} is not initialized.`);
+        }
+        return this.#connection.db;
+    }
+    async #checkExecution(workflowID, stepID) {
+        const statement = (0, drizzle_orm_1.sql) `
+        SELECT output, error FROM dbos.transaction_completion
+        WHERE workflow_id = ${workflowID} AND function_num = ${stepID}`;
+        const result = await this.#drizzle.execute(statement);
+        if (result.rows.length !== 1) {
+            return undefined;
+        }
+        const { output, error } = result.rows[0];
+        return error !== null ? { error } : { output };
+    }
+    static async #recordOutput(client, workflowID, stepID, output) {
+        try {
+            const statement = (0, drizzle_orm_1.sql) `
+        INSERT INTO dbos.transaction_completion (workflow_id, function_num, output) 
+        VALUES (${workflowID}, ${stepID}, ${output})`;
+            await client.execute(statement);
+        }
+        catch (error) {
+            if ((0, datasource_1.isPGKeyConflictError)(error)) {
+                throw new dbos_sdk_1.DBOSWorkflowConflictError(workflowID);
+            }
+            else {
+                throw error;
+            }
+        }
+    }
+    async #recordError(workflowID, stepID, error) {
+        try {
+            const statement = (0, drizzle_orm_1.sql) `
+        INSERT INTO dbos.transaction_completion (workflow_id, function_num, error) 
+        VALUES (${workflowID}, ${stepID}, ${error})`;
+            await this.#drizzle.execute(statement);
+        }
+        catch (error) {
+            if ((0, datasource_1.isPGKeyConflictError)(error)) {
+                throw new dbos_sdk_1.DBOSWorkflowConflictError(workflowID);
+            }
+            else {
+                throw error;
+            }
+        }
+    }
+    /* Invoke a transaction function, called by the framework */
+    async invokeTransactionFunction(config, target, func, ...args) {
+        const workflowID = dbos_sdk_1.DBOS.workflowID;
+        const stepID = dbos_sdk_1.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_sdk_1.DBOS.span?.setAttribute('cached', true);
+                if ('error' in previousResult) {
+                    throw superjson_1.SuperJSON.parse(previousResult.error);
+                }
+                return (previousResult.output ? superjson_1.SuperJSON.parse(previousResult.output) : null);
+            }
+            try {
+                const result = await this.#drizzle.transaction(async (client) => {
+                    // execute user's transaction function
+                    const result = await asyncLocalCtx.run({ client, owner: this }, async () => {
+                        return await func.call(target, ...args);
+                    });
+                    // save the output of read/write transactions
+                    if (saveResults) {
+                        await DrizzleTransactionHandler.#recordOutput(client, workflowID, stepID, superjson_1.SuperJSON.stringify(result));
+                    }
+                    return result;
+                }, { accessMode: config?.accessMode, isolationLevel: config?.isolationLevel });
+                return result;
+            }
+            catch (error) {
+                if ((0, datasource_1.isPGRetriableTransactionError)(error)) {
+                    dbos_sdk_1.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_1.SuperJSON.stringify(error);
+                        await this.#recordError(workflowID, stepID, message);
+                    }
+                    throw error;
+                }
+            }
+        }
+    }
+}
+class DrizzleDataSource {
+    name;
+    // User calls this... DBOS not directly involved...
+    static #getClient(p) {
+        if (!dbos_sdk_1.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');
+        }
+        if (p && p !== ctx.owner) {
+            throw new Error('Invalid retrieval of `DrizzleDataSource.client` from the incorrect object');
+        }
+        return ctx.client;
+    }
+    static get client() {
+        return DrizzleDataSource.#getClient(undefined);
+    }
+    get client() {
+        return DrizzleDataSource.#getClient(this.#provider);
+    }
+    static async initializeDBOSSchema(config) {
+        const client = new pg_1.Client(config);
+        try {
+            await client.connect();
+            await client.query(datasource_1.createTransactionCompletionSchemaPG);
+            await client.query(datasource_1.createTransactionCompletionTablePG);
+        }
+        finally {
+            await client.end();
+        }
+    }
+    #provider;
+    constructor(name, config, entities = {}) {
+        this.name = name;
+        this.#provider = new DrizzleTransactionHandler(name, config, entities);
+        (0, datasource_1.registerDataSource)(this.#provider);
+    }
+    async runTransaction(func, config) {
+        return await (0, datasource_1.runTransaction)(func, config?.name ?? func.name, { dsName: this.name, config });
+    }
+    registerTransaction(func, config) {
+        return (0, datasource_1.registerTransaction)(this.name, func, { name: config?.name ?? func.name }, config);
+    }
+    // decorator
+    transaction(config) {
+        // eslint-disable-next-line @typescript-eslint/no-this-alias
+        const ds = this;
+        return function decorator(_target, propertyKey, descriptor) {
+            if (!descriptor.value) {
+                throw new Error('Use of decorator when original method is undefined');
+            }
+            descriptor.value = ds.registerTransaction(descriptor.value, {
+                ...config,
+                name: config?.name ?? String(propertyKey),
+            });
+            return descriptor;
+        };
+    }
+}
+exports.DrizzleDataSource = DrizzleDataSource;
+//# sourceMappingURL=index.js.map

package/dist/index.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.js","sourceRoot":"","sources":["../index.ts"],"names":[],"mappings":";;;AAAA,2BAA4D;AAC5D,iDAAqE;AACrE,8DAUuC;AACvC,6DAAoE;AACpE,6CAAgD;AAChD,yCAAsC;AAEtC,6CAAkC;AASlC,MAAM,aAAa,GAAG,IAAI,+BAAiB,EAAmB,CAAC;AAc/D,MAAM,yBAAyB;IAKlB;IACQ;IACA;IANV,MAAM,GAAG,SAAS,CAAC;IAC5B,WAAW,CAAgC;IAE3C,YACW,IAAY,EACJ,MAAkB,EAClB,WAAsC,EAAE;QAFhD,SAAI,GAAJ,IAAI,CAAQ;QACJ,WAAM,GAAN,MAAM,CAAY;QAClB,aAAQ,GAAR,QAAQ,CAAgC;IACxD,CAAC;IAEJ,KAAK,CAAC,UAAU;QACd,MAAM,IAAI,GAAG,IAAI,CAAC,WAAW,CAAC;QAE9B,MAAM,MAAM,GAAG,IAAI,SAAI,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;QACrC,MAAM,EAAE,GAAG,IAAA,uBAAO,EAAC,MAAM,EAAE,EAAE,MAAM,EAAE,IAAI,CAAC,QAAQ,EAAE,CAAC,CAAC;QACtD,IAAI,CAAC,WAAW,GAAG,EAAE,EAAE,EAAE,GAAG,EAAE,GAAG,EAAE,CAAC,MAAM,CAAC,GAAG,EAAE,EAAE,CAAC;QAEnD,MAAM,IAAI,EAAE,GAAG,EAAE,CAAC;IACpB,CAAC;IAED,KAAK,CAAC,OAAO;QACX,MAAM,IAAI,GAAG,IAAI,CAAC,WAAW,CAAC;QAE9B,IAAI,CAAC,WAAW,GAAG,SAAS,CAAC;QAE7B,MAAM,IAAI,EAAE,GAAG,EAAE,CAAC;IACpB,CAAC;IAED,IAAI,QAAQ;QACV,IAAI,CAAC,IAAI,CAAC,WAAW,EAAE,CAAC;YACtB,MAAM,IAAI,KAAK,CAAC,cAAc,IAAI,CAAC,IAAI,sBAAsB,CAAC,CAAC;QACjE,CAAC;QACD,OAAO,IAAI,CAAC,WAAW,CAAC,EAAE,CAAC;IAC7B,CAAC;IAED,KAAK,CAAC,eAAe,CACnB,UAAkB,EAClB,MAAc;QAId,MAAM,SAAS,GAAG,IAAA,iBAAG,EAAA;;8BAEK,UAAU,uBAAuB,MAAM,EAAE,CAAC;QACpE,MAAM,MAAM,GAAG,MAAM,IAAI,CAAC,QAAQ,CAAC,OAAO,CAAS,SAAS,CAAC,CAAC;QAE9D,IAAI,MAAM,CAAC,IAAI,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;YAC7B,OAAO,SAAS,CAAC;QACnB,CAAC;QAED,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,GAAG,MAAM,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;QACzC,OAAO,KAAK,KAAK,IAAI,CAAC,CAAC,CAAC,EAAE,KAAK,EAAE,CAAC,CAAC,CAAC,EAAE,MAAM,EAAE,CAAC;IACjD,CAAC;IAED,MAAM,CAAC,KAAK,CAAC,aAAa,CACxB,MAAiD,EACjD,UAAkB,EAClB,MAAc,EACd,MAAc;QAEd,IAAI,CAAC;YACH,MAAM,SAAS,GAAG,IAAA,iBAAG,EAAA;;kBAET,UAAU,KAAK,MAAM,KAAK,MAAM,GAAG,CAAC;YAChD,MAAM,MAAM,CAAC,OAAO,CAAC,SAAS,CAAC,CAAC;QAClC,CAAC;QAAC,OAAO,KAAK,EAAE,CAAC;YACf,IAAI,IAAA,iCAAoB,EAAC,KAAK,CAAC,EAAE,CAAC;gBAChC,MAAM,IAAI,oCAAyB,CAAC,UAAU,CAAC,CAAC;YAClD,CAAC;iBAAM,CAAC;gBACN,MAAM,KAAK,CAAC;YACd,CAAC;QACH,CAAC;IACH,CAAC;IAED,KAAK,CAAC,YAAY,CAAC,UAAkB,EAAE,MAAc,EAAE,KAAa;QAClE,IAAI,CAAC;YACH,MAAM,SAAS,GAAG,IAAA,iBAAG,EAAA;;kBAET,UAAU,KAAK,MAAM,KAAK,KAAK,GAAG,CAAC;YAC/C,MAAM,IAAI,CAAC,QAAQ,CAAC,OAAO,CAAC,SAAS,CAAC,CAAC;QACzC,CAAC;QAAC,OAAO,KAAK,EAAE,CAAC;YACf,IAAI,IAAA,iCAAoB,EAAC,KAAK,CAAC,EAAE,CAAC;gBAChC,MAAM,IAAI,oCAAyB,CAAC,UAAU,CAAC,CAAC;YAClD,CAAC;iBAAM,CAAC;gBACN,MAAM,KAAK,CAAC;YACd,CAAC;QACH,CAAC;IACH,CAAC;IAED,4DAA4D;IAC5D,KAAK,CAAC,yBAAyB,CAC7B,MAAqC,EACrC,MAAY,EACZ,IAAoD,EACpD,GAAG,IAAU;QAEb,MAAM,UAAU,GAAG,eAAI,CAAC,UAAU,CAAC;QACnC,MAAM,MAAM,GAAG,eAAI,CAAC,MAAM,CAAC;QAC3B,IAAI,UAAU,KAAK,SAAS,IAAI,MAAM,KAAK,SAAS,EAAE,CAAC;YACrD,MAAM,IAAI,KAAK,CAAC,6CAA6C,CAAC,CAAC;QACjE,CAAC;QAED,MAAM,QAAQ,GAAG,MAAM,EAAE,UAAU,KAAK,WAAW,CAAC,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,KAAK,CAAC;QACnE,MAAM,WAAW,GAAG,CAAC,QAAQ,IAAI,UAAU,KAAK,SAAS,CAAC;QAE1D,4BAA4B;QAC5B,IAAI,WAAW,GAAG,CAAC,CAAC;QACpB,MAAM,aAAa,GAAG,GAAG,CAAC;QAC1B,MAAM,cAAc,GAAG,IAAI,CAAC,CAAC,0BAA0B;QAEvD,OAAO,IAAI,EAAE,CAAC;YACZ,oDAAoD;YACpD,MAAM,cAAc,GAAG,WAAW,CAAC,CAAC,CAAC,MAAM,IAAI,CAAC,eAAe,CAAC,UAAU,EAAE,MAAO,CAAC,CAAC,CAAC,CAAC,SAAS,CAAC;YACjG,IAAI,cAAc,EAAE,CAAC;gBACnB,eAAI,CAAC,IAAI,EAAE,YAAY,CAAC,QAAQ,EAAE,IAAI,CAAC,CAAC;gBAExC,IAAI,OAAO,IAAI,cAAc,EAAE,CAAC;oBAC9B,MAAM,qBAAS,CAAC,KAAK,CAAC,cAAc,CAAC,KAAK,CAAC,CAAC;gBAC9C,CAAC;gBACD,OAAO,CAAC,cAAc,CAAC,MAAM,CAAC,CAAC,CAAC,qBAAS,CAAC,KAAK,CAAC,cAAc,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,IAAI,CAAW,CAAC;YAC3F,CAAC;YAED,IAAI,CAAC;gBACH,MAAM,MAAM,GAAG,MAAM,IAAI,CAAC,QAAQ,CAAC,WAAW,CAC5C,KAAK,EAAE,MAAM,EAAE,EAAE;oBACf,sCAAsC;oBACtC,MAAM,MAAM,GAAG,MAAM,aAAa,CAAC,GAAG,CAAC,EAAE,MAAM,EAAE,KAAK,EAAE,IAAI,EAAE,EAAE,KAAK,IAAI,EAAE;wBACzE,OAAO,MAAM,IAAI,CAAC,IAAI,CAAC,MAAM,EAAE,GAAG,IAAI,CAAC,CAAC;oBAC1C,CAAC,CAAC,CAAC;oBAEH,6CAA6C;oBAC7C,IAAI,WAAW,EAAE,CAAC;wBAChB,MAAM,yBAAyB,CAAC,aAAa,CAAC,MAAM,EAAE,UAAU,EAAE,MAAO,EAAE,qBAAS,CAAC,SAAS,CAAC,MAAM,CAAC,CAAC,CAAC;oBAC1G,CAAC;oBAED,OAAO,MAAM,CAAC;gBAChB,CAAC,EACD,EAAE,UAAU,EAAE,MAAM,EAAE,UAAU,EAAE,cAAc,EAAE,MAAM,EAAE,cAAc,EAAE,CAC3E,CAAC;gBAEF,OAAO,MAAM,CAAC;YAChB,CAAC;YAAC,OAAO,KAAK,EAAE,CAAC;gBACf,IAAI,IAAA,0CAA6B,EAAC,KAAK,CAAC,EAAE,CAAC;oBACzC,eAAI,CAAC,IAAI,EAAE,QAAQ,CAAC,2BAA2B,EAAE,EAAE,eAAe,EAAE,WAAW,EAAE,EAAE,WAAW,CAAC,GAAG,EAAE,CAAC,CAAC;oBACtG,gCAAgC;oBAChC,MAAM,IAAI,OAAO,CAAC,CAAC,OAAO,EAAE,EAAE,CAAC,UAAU,CAAC,OAAO,EAAE,WAAW,CAAC,CAAC,CAAC;oBACjE,WAAW,GAAG,IAAI,CAAC,GAAG,CAAC,WAAW,GAAG,aAAa,EAAE,cAAc,CAAC,CAAC;oBACpE,SAAS;gBACX,CAAC;qBAAM,CAAC;oBACN,IAAI,WAAW,EAAE,CAAC;wBAChB,MAAM,OAAO,GAAG,qBAAS,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC;wBAC3C,MAAM,IAAI,CAAC,YAAY,CAAC,UAAU,EAAE,MAAO,EAAE,OAAO,CAAC,CAAC;oBACxD,CAAC;oBAED,MAAM,KAAK,CAAC;gBACd,CAAC;YACH,CAAC;QACH,CAAC;IACH,CAAC;CACF;AAED,MAAa,iBAAiB;IAsCjB;IArCX,mDAAmD;IACnD,MAAM,CAAC,UAAU,CAAC,CAA6B;QAC7C,IAAI,CAAC,eAAI,CAAC,eAAe,EAAE,EAAE,CAAC;YAC5B,MAAM,IAAI,KAAK,CAAC,uEAAuE,CAAC,CAAC;QAC3F,CAAC;QACD,MAAM,GAAG,GAAG,aAAa,CAAC,QAAQ,EAAE,CAAC;QACrC,IAAI,CAAC,GAAG,EAAE,CAAC;YACT,MAAM,IAAI,KAAK,CAAC,uEAAuE,CAAC,CAAC;QAC3F,CAAC;QACD,IAAI,CAAC,IAAI,CAAC,KAAK,GAAG,CAAC,KAAK,EAAE,CAAC;YACzB,MAAM,IAAI,KAAK,CAAC,2EAA2E,CAAC,CAAC;QAC/F,CAAC;QACD,OAAO,GAAG,CAAC,MAAM,CAAC;IACpB,CAAC;IAED,MAAM,KAAK,MAAM;QACf,OAAO,iBAAiB,CAAC,UAAU,CAAC,SAAS,CAAC,CAAC;IACjD,CAAC;IAED,IAAI,MAAM;QACR,OAAO,iBAAiB,CAAC,UAAU,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC;IACtD,CAAC;IAED,MAAM,CAAC,KAAK,CAAC,oBAAoB,CAAC,MAAoB;QACpD,MAAM,MAAM,GAAG,IAAI,WAAM,CAAC,MAAM,CAAC,CAAC;QAClC,IAAI,CAAC;YACH,MAAM,MAAM,CAAC,OAAO,EAAE,CAAC;YACvB,MAAM,MAAM,CAAC,KAAK,CAAC,gDAAmC,CAAC,CAAC;YACxD,MAAM,MAAM,CAAC,KAAK,CAAC,+CAAkC,CAAC,CAAC;QACzD,CAAC;gBAAS,CAAC;YACT,MAAM,MAAM,CAAC,GAAG,EAAE,CAAC;QACrB,CAAC;IACH,CAAC;IAED,SAAS,CAA4B;IAErC,YACW,IAAY,EACrB,MAAkB,EAClB,WAAsC,EAAE;QAF/B,SAAI,GAAJ,IAAI,CAAQ;QAIrB,IAAI,CAAC,SAAS,GAAG,IAAI,yBAAyB,CAAC,IAAI,EAAE,MAAM,EAAE,QAAQ,CAAC,CAAC;QACvE,IAAA,+BAAkB,EAAC,IAAI,CAAC,SAAS,CAAC,CAAC;IACrC,CAAC;IAED,KAAK,CAAC,cAAc,CAAI,IAAsB,EAAE,MAA0B;QACxE,OAAO,MAAM,IAAA,2BAAc,EAAC,IAAI,EAAE,MAAM,EAAE,IAAI,IAAI,IAAI,CAAC,IAAI,EAAE,EAAE,MAAM,EAAE,IAAI,CAAC,IAAI,EAAE,MAAM,EAAE,CAAC,CAAC;IAC9F,CAAC;IAED,mBAAmB,CACjB,IAAoD,EACpD,MAA0B;QAE1B,OAAO,IAAA,gCAAmB,EAAC,IAAI,CAAC,IAAI,EAAE,IAAI,EAAE,EAAE,IAAI,EAAE,MAAM,EAAE,IAAI,IAAI,IAAI,CAAC,IAAI,EAAE,EAAE,MAAM,CAAC,CAAC;IAC3F,CAAC;IAED,YAAY;IACZ,WAAW,CAAC,MAA0B;QACpC,4DAA4D;QAC5D,MAAM,EAAE,GAAG,IAAI,CAAC;QAChB,OAAO,SAAS,SAAS,CACvB,OAAe,EACf,WAAwB,EACxB,UAAmF;YAEnF,IAAI,CAAC,UAAU,CAAC,KAAK,EAAE,CAAC;gBACtB,MAAM,IAAI,KAAK,CAAC,oDAAoD,CAAC,CAAC;YACxE,CAAC;YAED,UAAU,CAAC,KAAK,GAAG,EAAE,CAAC,mBAAmB,CAAC,UAAU,CAAC,KAAK,EAAE;gBAC1D,GAAG,MAAM;gBACT,IAAI,EAAE,MAAM,EAAE,IAAI,IAAI,MAAM,CAAC,WAAW,CAAC;aAC1C,CAAC,CAAC;YAEH,OAAO,UAAU,CAAC;QACpB,CAAC,CAAC;IACJ,CAAC;CACF;AA9ED,8CA8EC"}

package/index.ts

@@ -0,0 +1,280 @@
+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 }>;
+  owner: DrizzleTransactionHandler;
+}
+
+export type TransactionConfig = Pick<PgTransactionConfig, 'isolationLevel' | 'accessMode'> & { name?: string };
+
+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, owner: this }, 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 #getClient(p?: DrizzleTransactionHandler): 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');
+    }
+    if (p && p !== ctx.owner) {
+      throw new Error('Invalid retrieval of `DrizzleDataSource.client` from the incorrect object');
+    }
+    return ctx.client;
+  }
+
+  static get client() {
+    return DrizzleDataSource.#getClient(undefined);
+  }
+
+  get client() {
+    return DrizzleDataSource.#getClient(this.#provider);
+  }
+
+  static async initializeDBOSSchema(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>(func: () => Promise<T>, config?: TransactionConfig) {
+    return await runTransaction(func, config?.name ?? func.name, { dsName: this.name, config });
+  }
+
+  registerTransaction<This, Args extends unknown[], Return>(
+    func: (this: This, ...args: Args) => Promise<Return>,
+    config?: TransactionConfig,
+  ): (this: This, ...args: Args) => Promise<Return> {
+    return registerTransaction(this.name, func, { name: config?.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,
+        name: config?.name ?? String(propertyKey),
+      });
+
+      return descriptor;
+    };
+  }
+}

package/package.json

@@ -1,8 +1,11 @@
 {
   "name": "@dbos-inc/drizzle-datasource",
-  "version": "3.0.11-preview.gc9233b8190",
-  "main": "dist/src/index.js",
-  "types": "dist/src/index.d.ts",
+  "version": "3.0.16-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.initializeDBOSSchema(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();
+    }
+  });
+});