feat: add pre-commit callbacks on EntityQueryContext (#147)

packages/entity-database-adapter-knex/src/__integration-tests__/PostgresEntityIntegration-test.ts

@@ -600,4 +600,37 @@ describe('postgres entity integration', () => {
       });
     });
   });
+
+  describe('queryContext callback behavior', () => {
+    it('calls callbacks correctly', async () => {
+      const vc1 = new ViewerContext(createKnexIntegrationTestEntityCompanionProvider(knexInstance));
+
+      let preCommitCallCount = 0;
+      let postCommitCallCount = 0;
+
+      await vc1.runInTransactionForDatabaseAdaptorFlavorAsync('postgres', async (queryContext) => {
+        queryContext.appendPostCommitCallback(async () => {
+          postCommitCallCount++;
+        });
+        queryContext.appendPreCommitCallback(async () => {
+          preCommitCallCount++;
+        }, 0);
+      });
+
+      await expect(
+        vc1.runInTransactionForDatabaseAdaptorFlavorAsync('postgres', async (queryContext) => {
+          queryContext.appendPostCommitCallback(async () => {
+            postCommitCallCount++;
+          });
+          queryContext.appendPreCommitCallback(async () => {
+            preCommitCallCount++;
+            throw Error('wat');
+          }, 0);
+        })
+      ).rejects.toThrowError('wat');
+
+      expect(preCommitCallCount).toBe(2);
+      expect(postCommitCallCount).toBe(1);
+    });
+  });
 });

packages/entity/src/EntityQueryContext.ts

@@ -1,6 +1,12 @@
+import assert from 'assert';
+
 import EntityQueryContextProvider from './EntityQueryContextProvider';
 
 export type PostCommitCallback = (...args: any) => Promise<any>;
+export type PreCommitCallback = (
+  queryContext: EntityTransactionalQueryContext,
+  ...args: any
+) => Promise<any>;
 
 /**
  * Entity framework representation of transactional and non-transactional database
@@ -46,10 +52,27 @@ export class EntityTransactionalQueryContext extends EntityQueryContext {
   private readonly postCommitInvalidationCallbacks: PostCommitCallback[] = [];
   private readonly postCommitCallbacks: PostCommitCallback[] = [];
 
+  private readonly preCommitCallbacks: { callback: PreCommitCallback; order: number }[] = [];
+
+  /**
+   * Schedule a pre-commit callback. These will be run within the transaction right before it is
+   * committed, and will be run in the order specified. Ordering of callbacks scheduled with the
+   * same value for the order parameter is undefined within that ordering group.
+   * @param callback - callback to schedule
+   * @param order - order in which this should be run relative to other scheduled pre-commit callbacks,
+   *                with higher numbers running later than lower numbers.
+   */
+  public appendPreCommitCallback(callback: PreCommitCallback, order: number): void {
+    assert(
+      order >= Number.MIN_SAFE_INTEGER && order <= Number.MAX_SAFE_INTEGER,
+      `Invalid order specified: ${order}`
+    );
+    this.preCommitCallbacks.push({ callback, order });
+  }
+
   /**
    * Schedule a post-commit cache invalidation callback. These are run before normal
    * post-commit callbacks in order to have cache consistency in normal post-commit callbacks.
-   *
    * @param callback - callback to schedule
    */
   public appendPostCommitInvalidationCallback(callback: PostCommitCallback): void {
@@ -65,6 +88,17 @@ export class EntityTransactionalQueryContext extends EntityQueryContext {
     this.postCommitCallbacks.push(callback);
   }
 
+  public async runPreCommitCallbacksAsync(): Promise<void> {
+    const callbacks = [...this.preCommitCallbacks]
+      .sort((a, b) => a.order - b.order)
+      .map((c) => c.callback);
+    this.preCommitCallbacks.length = 0;
+
+    for (const callback of callbacks) {
+      await callback(this);
+    }
+  }
+
   public async runPostCommitCallbacksAsync(): Promise<void> {
     const invalidationCallbacks = [...this.postCommitInvalidationCallbacks];
     this.postCommitInvalidationCallbacks.length = 0;

packages/entity/src/EntityQueryContextProvider.ts

@@ -38,6 +38,7 @@ export default abstract class EntityQueryContextProvider {
     >()(async (queryInterface) => {
       const queryContext = new EntityTransactionalQueryContext(queryInterface);
       const result = await transactionScope(queryContext);
+      await queryContext.runPreCommitCallbacksAsync();
       return [result, queryContext];
     });
     await queryContext.runPostCommitCallbacksAsync();

packages/entity/src/__tests__/EntityQueryContext-test.ts

@@ -0,0 +1,82 @@
+import invariant from 'invariant';
+
+import { EntityQueryContext } from '../EntityQueryContext';
+import ViewerContext from '../ViewerContext';
+import { createUnitTestEntityCompanionProvider } from '../utils/testing/createUnitTestEntityCompanionProvider';
+
+describe(EntityQueryContext, () => {
+  describe('callbacks', () => {
+    it('calls all callbacks, and calls invalidation first', async () => {
+      const companionProvider = createUnitTestEntityCompanionProvider();
+      const viewerContext = new ViewerContext(companionProvider);
+
+      const preCommitFirstCallback = jest.fn(async (): Promise<void> => {});
+      const preCommitSecondCallback = jest.fn(async (): Promise<void> => {});
+      const postCommitInvalidationCallback = jest.fn(async (): Promise<void> => {
+        invariant(
+          preCommitFirstCallback.mock.calls.length === 1,
+          'preCommit should be called before postCommitInvalidation'
+        );
+        invariant(
+          preCommitSecondCallback.mock.calls.length === 1,
+          'preCommit should be called before postCommitInvalidation'
+        );
+      });
+      const postCommitCallback = jest.fn(async (): Promise<void> => {
+        invariant(
+          preCommitFirstCallback.mock.calls.length === 1,
+          'preCommit should be called before postCommit'
+        );
+        invariant(
+          preCommitSecondCallback.mock.calls.length === 1,
+          'preCommit should be called before postCommit'
+        );
+        invariant(
+          postCommitInvalidationCallback.mock.calls.length === 1,
+          'postCommitInvalidation should be called before postCommit'
+        );
+      });
+
+      await viewerContext.runInTransactionForDatabaseAdaptorFlavorAsync(
+        'postgres',
+        async (queryContext) => {
+          queryContext.appendPostCommitCallback(postCommitCallback);
+          queryContext.appendPostCommitInvalidationCallback(postCommitInvalidationCallback);
+          queryContext.appendPreCommitCallback(preCommitSecondCallback, 2);
+          queryContext.appendPreCommitCallback(preCommitFirstCallback, 1);
+        }
+      );
+
+      expect(preCommitFirstCallback).toHaveBeenCalledTimes(1);
+      expect(preCommitSecondCallback).toHaveBeenCalledTimes(1);
+      expect(postCommitCallback).toHaveBeenCalledTimes(1);
+      expect(postCommitInvalidationCallback).toHaveBeenCalledTimes(1);
+    });
+
+    it('prevents transaction from finishing when precommit throws (post commit callbacks are not called)', async () => {
+      const companionProvider = createUnitTestEntityCompanionProvider();
+      const viewerContext = new ViewerContext(companionProvider);
+
+      const preCommitCallback = jest.fn(async (): Promise<void> => {
+        throw new Error('wat');
+      });
+      const postCommitInvalidationCallback = jest.fn(async (): Promise<void> => {});
+      const postCommitCallback = jest.fn(async (): Promise<void> => {});
+
+      await expect(
+        viewerContext.runInTransactionForDatabaseAdaptorFlavorAsync(
+          'postgres',
+          async (queryContext) => {
+            queryContext.appendPostCommitCallback(postCommitCallback);
+            queryContext.appendPostCommitInvalidationCallback(postCommitInvalidationCallback);
+            queryContext.appendPreCommitCallback(preCommitCallback, 0);
+          }
+        )
+      ).rejects.toThrowError('wat');
+
+      expect(preCommitCallback).toHaveBeenCalledTimes(1);
+      expect(postCommitCallback).toHaveBeenCalledTimes(0);
+      expect(postCommitInvalidationCallback).toHaveBeenCalledTimes(0);
+    });
+  });
+});