@dbos-inc/knex-datasource 3.0.8-preview → 3.0.8-preview.g493d2d1c2b

package/index.ts

@@ -9,8 +9,6 @@ import {
   runTransaction,
   DBOSDataSource,
   registerDataSource,
-  createTransactionCompletionSchemaPG,
-  createTransactionCompletionTablePG,
 } from '@dbos-inc/dbos-sdk/datasource';
 import { AsyncLocalStorage } from 'async_hooks';
 import knex, { type Knex } from 'knex';
@@ -20,7 +18,6 @@ interface transaction_completion {
   workflow_id: string;
   function_num: number;
   output: string | null;
-  error: string | null;
 }
 
 interface KnexDataSourceContext {
@@ -31,79 +28,50 @@ export type TransactionConfig = Pick<Knex.TransactionConfig, 'isolationLevel' |
 
 const asyncLocalCtx = new AsyncLocalStorage<KnexDataSourceContext>();
 
-class KnexTransactionHandler implements DataSourceTransactionHandler {
+class KnexDSTH implements DataSourceTransactionHandler {
+  readonly name: string;
   readonly dsType = 'KnexDataSource';
-  #knexDBField: Knex | undefined;
+  readonly #knexDB: Knex;
 
-  constructor(
-    readonly name: string,
-    private readonly config: Knex.Config,
-  ) {}
-
-  async initialize(): Promise<void> {
-    const knexDB = this.#knexDBField;
-    this.#knexDBField = knex(this.config);
-    await knexDB?.destroy();
+  constructor(name: string, config: Knex.Config) {
+    this.name = name;
+    this.#knexDB = knex(config);
   }
 
-  async destroy(): Promise<void> {
-    const knexDB = this.#knexDBField;
-    this.#knexDBField = undefined;
-    await knexDB?.destroy();
+  initialize(): Promise<void> {
+    return Promise.resolve();
   }
 
-  get #knexDB() {
-    if (!this.#knexDBField) {
-      throw new Error(`DataSource ${this.name} is not initialized.`);
-    }
-    return this.#knexDBField;
+  destroy(): Promise<void> {
+    return this.#knexDB.destroy();
   }
 
-  async #checkExecution(
+  static async #checkExecution(
+    client: Knex.Transaction,
     workflowID: string,
-    stepID: number,
-  ): Promise<{ output: string | null } | { error: string } | undefined> {
-    const result = await this.#knexDB<transaction_completion>('transaction_completion')
+    functionNum: number,
+  ): Promise<{ output: string | null } | undefined> {
+    const result = await client<transaction_completion>('transaction_completion')
       .withSchema('dbos')
-      .select('output', 'error')
+      .select('output')
       .where({
         workflow_id: workflowID,
-        function_num: stepID,
+        function_num: functionNum,
       })
       .first();
-    if (result === undefined) {
-      return undefined;
-    }
-    const { output, error } = result;
-    return error !== null ? { error } : { output };
-  }
-
-  async #recordError(workflowID: string, stepID: number, error: string): Promise<void> {
-    try {
-      await this.#knexDB<transaction_completion>('transaction_completion').withSchema('dbos').insert({
-        workflow_id: workflowID,
-        function_num: stepID,
-        error,
-      });
-    } catch (error) {
-      if (isPGKeyConflictError(error)) {
-        throw new DBOSWorkflowConflictError(workflowID);
-      } else {
-        throw error;
-      }
-    }
+    return result === undefined ? undefined : { output: result.output };
   }
 
   static async #recordOutput(
     client: Knex.Transaction,
     workflowID: string,
-    stepID: number,
+    functionNum: number,
     output: string | null,
   ): Promise<void> {
     try {
       await client<transaction_completion>('transaction_completion').withSchema('dbos').insert({
         workflow_id: workflowID,
-        function_num: stepID,
+        function_num: functionNum,
         output,
       });
     } catch (error) {
@@ -122,61 +90,67 @@ class KnexTransactionHandler implements DataSourceTransactionHandler {
     ...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.');
+    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;
-    const saveResults = !readOnly && workflowID !== undefined;
-
-    // Retry loop if appropriate
     let retryWaitMS = 1;
     const backoffFactor = 1.5;
-    const maxRetryWaitMS = 2000; // Maximum wait 2 seconds.
+    const maxRetryWaitMS = 2000;
 
     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.#knexDB.transaction<Return>(
           async (client) => {
+            // Check to see if this tx has already been executed
+            const previousResult =
+              readOnly || !workflowID ? undefined : await KnexDSTH.#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 (saveResults) {
-              await KnexTransactionHandler.#recordOutput(client, workflowID, stepID!, SuperJSON.stringify(result));
+            if (!readOnly && workflowID) {
+              await KnexDSTH.#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)) {
-          DBOS.span?.addEvent('TXN SERIALIZATION FAILURE', { retryWaitMillis: retryWaitMS }, performance.now());
+          // 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 {
-          if (saveResults) {
-            const message = SuperJSON.stringify(error);
-            await this.#recordError(workflowID, stepID!, message);
-          }
+          // 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;
         }
@@ -192,7 +166,7 @@ export class KnexDataSource implements DBOSDataSource<TransactionConfig> {
     }
     const ctx = asyncLocalCtx.getStore();
     if (!ctx) {
-      throw new Error('invalid use of KnexDataSource.client outside of a DBOS transaction.');
+      throw new Error('No async local context found.');
     }
     return ctx.client;
   }
@@ -200,20 +174,27 @@ export class KnexDataSource implements DBOSDataSource<TransactionConfig> {
   static async initializeInternalSchema(config: Knex.Config) {
     const knexDB = knex(config);
     try {
-      await knexDB.raw(createTransactionCompletionSchemaPG);
-      await knexDB.raw(createTransactionCompletionTablePG);
+      await knexDB.schema.createSchemaIfNotExists('dbos');
+      const exists = await knexDB.schema.withSchema('dbos').hasTable('transaction_completion');
+      if (!exists) {
+        await knexDB.schema.withSchema('dbos').createTable('transaction_completion', (table) => {
+          table.string('workflow_id').notNullable();
+          table.integer('function_num').notNullable();
+          table.string('output').nullable();
+          table.primary(['workflow_id', 'function_num']);
+        });
+      }
     } finally {
       await knexDB.destroy();
     }
   }
 
-  #provider: KnexTransactionHandler;
+  readonly name: string;
+  #provider: KnexDSTH;
 
-  constructor(
-    readonly name: string,
-    config: Knex.Config,
-  ) {
-    this.#provider = new KnexTransactionHandler(name, config);
+  constructor(name: string, config: Knex.Config) {
+    this.name = name;
+    this.#provider = new KnexDSTH(name, config);
     registerDataSource(this.#provider);
   }
 
@@ -223,10 +204,10 @@ export class KnexDataSource implements DBOSDataSource<TransactionConfig> {
 
   registerTransaction<This, Args extends unknown[], Return>(
     func: (this: This, ...args: Args) => Promise<Return>,
+    name: string,
     config?: TransactionConfig,
-    name?: string,
   ): (this: This, ...args: Args) => Promise<Return> {
-    return registerTransaction(this.name, func, { name: name ?? func.name }, config);
+    return registerTransaction(this.name, func, { name }, config);
   }
 
   transaction(config?: TransactionConfig) {
@@ -234,14 +215,14 @@ export class KnexDataSource implements DBOSDataSource<TransactionConfig> {
     const ds = this;
     return function decorator<This, Args extends unknown[], Return>(
       _target: object,
-      propertyKey: PropertyKey,
+      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, config, String(propertyKey));
+      descriptor.value = ds.registerTransaction(descriptor.value, propertyKey.toString(), config);
 
       return descriptor;
     };

package/package.json

@@ -1,23 +1,21 @@
 {
   "name": "@dbos-inc/knex-datasource",
-  "version": "3.0.8-preview",
-  "description": "DBOS DataSource library for Knex ORM with PostgreSQL support",
+  "version": "3.0.8-preview.g493d2d1c2b",
+  "description": "",
   "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",
     "directory": "packages/knex-datasource"
   },
+  "homepage": "https://docs.dbos.dev/",
+  "main": "index.js",
   "scripts": {
     "build": "tsc --project tsconfig.json",
     "test": "jest --detectOpenHandles"
   },
   "dependencies": {
     "knex": "^3.1.0",
-    "pg": "^8.11.3",
     "superjson": "^1.13"
   },
   "peerDependencies": {
@@ -27,6 +25,7 @@
     "@types/jest": "^29.5.14",
     "@types/pg": "^8.15.2",
     "jest": "^29.7.0",
+    "pg": "^8.11.3",
     "typescript": "^5.4.5"
   }
 }

package/tests/config.test.ts

@@ -9,7 +9,7 @@ describe('KnexDataSource.configure', () => {
     const client = new Client({ ...config, database: 'postgres' });
     try {
       await client.connect();
-      await dropDB(client, config.database, true);
+      await dropDB(client, config.database);
       await ensureDB(client, config.database);
     } finally {
       await client.end();