@dbos-inc/knex-datasource 3.0.8-preview → 3.0.10-preview

package/index.ts

@@ -185,6 +185,10 @@ class KnexTransactionHandler implements DataSourceTransactionHandler {
   }
 }
 
+function isKnex(value: Knex | Knex.Config): value is Knex {
+  return 'raw' in value;
+}
+
 export class KnexDataSource implements DBOSDataSource<TransactionConfig> {
   static get client(): Knex.Transaction {
     if (!DBOS.isInTransaction()) {
@@ -197,13 +201,38 @@ export class KnexDataSource implements DBOSDataSource<TransactionConfig> {
     return ctx.client;
   }
 
-  static async initializeInternalSchema(config: Knex.Config) {
-    const knexDB = knex(config);
-    try {
+  static async initializeSchema(knexOrConfig: Knex.Config) {
+    if (isKnex(knexOrConfig)) {
+      await $initSchema(knexOrConfig);
+    } else {
+      const knexDB = knex(knexOrConfig);
+      try {
+        await $initSchema(knexDB);
+      } finally {
+        await knexDB.destroy();
+      }
+    }
+
+    async function $initSchema(knexDB: Knex) {
       await knexDB.raw(createTransactionCompletionSchemaPG);
       await knexDB.raw(createTransactionCompletionTablePG);
-    } finally {
-      await knexDB.destroy();
+    }
+  }
+
+  static async uninitializeSchema(knexOrConfig: Knex.Config) {
+    if (isKnex(knexOrConfig)) {
+      await $uninitSchema(knexOrConfig);
+    } else {
+      const knexDB = knex(knexOrConfig);
+      try {
+        await $uninitSchema(knexDB);
+      } finally {
+        await knexDB.destroy();
+      }
+    }
+
+    async function $uninitSchema(knexDB: Knex) {
+      await knexDB.raw('DROP TABLE IF EXISTS dbos.transaction_completion; DROP SCHEMA IF EXISTS dbos;');
     }
   }
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@dbos-inc/knex-datasource",
-  "version": "3.0.8-preview",
+  "version": "3.0.10-preview",
   "description": "DBOS DataSource library for Knex ORM with PostgreSQL support",
   "license": "MIT",
   "main": "dist/index.js",

package/tests/config.test.ts

@@ -1,11 +1,12 @@
 import { Client } from 'pg';
 import { KnexDataSource } from '../index';
 import { dropDB, ensureDB } from './test-helpers';
+import knex from 'knex';
 
-describe('KnexDataSource.configure', () => {
+describe('KnexDataSource.initializeSchema', () => {
   const config = { user: 'postgres', database: 'knex_ds_config_test' };
 
-  beforeAll(async () => {
+  beforeEach(async () => {
     const client = new Client({ ...config, database: 'postgres' });
     try {
       await client.connect();
@@ -16,16 +17,55 @@ describe('KnexDataSource.configure', () => {
     }
   });
 
-  test('configure creates tx outputs table', async () => {
-    await KnexDataSource.initializeInternalSchema({ client: 'pg', connection: config });
+  async function queryTxCompletionTable(client: Client) {
+    const result = await client.query(
+      'SELECT workflow_id, function_num, output, error FROM dbos.transaction_completion',
+    );
+    return result.rowCount;
+  }
+
+  async function txCompletionTableExists(client: Client) {
+    const result = await client.query<{ exists: boolean }>(
+      "SELECT EXISTS (SELECT 1 FROM information_schema.tables  WHERE table_schema = 'dbos' AND table_name = 'transaction_completion');",
+    );
+    if (result.rowCount !== 1) throw new Error(`unexpected rowcount ${result.rowCount}`);
+    return result.rows[0].exists;
+  }
+
+  test('initializeSchema-with-config', async () => {
+    const knexConfig = { client: 'pg', connection: config };
+    await KnexDataSource.initializeSchema(knexConfig);
 
     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);
+      await expect(txCompletionTableExists(client)).resolves.toBe(true);
+      await expect(queryTxCompletionTable(client)).resolves.toEqual(0);
+
+      await KnexDataSource.uninitializeSchema(knexConfig);
+
+      await expect(txCompletionTableExists(client)).resolves.toBe(false);
+    } finally {
+      await client.end();
+    }
+  });
+
+  test('initializeSchema-with-knex', async () => {
+    const knexDB = knex({ client: 'pg', connection: config });
+    const client = new Client(config);
+    try {
+      await KnexDataSource.initializeSchema(knexDB);
+      await client.connect();
+
+      await expect(txCompletionTableExists(client)).resolves.toBe(true);
+      await expect(queryTxCompletionTable(client)).resolves.toEqual(0);
+
+      await KnexDataSource.uninitializeSchema(knexDB);
+
+      await expect(txCompletionTableExists(client)).resolves.toBe(false);
     } finally {
       await client.end();
+      await knexDB.destroy();
     }
   });
 });

package/tests/datasource.test.ts

@@ -43,7 +43,7 @@ describe('KnexDataSource', () => {
       }
     }
 
-    await KnexDataSource.initializeInternalSchema(config);
+    await KnexDataSource.initializeSchema(config);
   });
 
   afterAll(async () => {