feat: dataloaders within transactions

packages/entity/src/EntityLoader.ts

@@ -284,6 +284,13 @@ export default class EntityLoader<
     await this.dataManager.invalidateObjectFieldsAsync(objectFields);
   }
 
+  async invalidateFieldsTransactionalAsync(objectFields: Readonly<TFields>): Promise<void> {
+    await this.dataManager.invalidateObjectFieldsTransactionalAsync(
+      this.queryContext,
+      objectFields
+    );
+  }
+
   /**
    * Invalidate all caches for an entity. One potential use case would be to keep the entity
    * framework in sync with changes made to data outside of the framework.

packages/entity/src/EntityMutator.ts

@@ -205,6 +205,7 @@ export class CreateMutator<
     queryContext.appendPostCommitCallback(
       entityLoader.invalidateFieldsAsync.bind(entityLoader, insertResult)
     );
+    await entityLoader.invalidateFieldsTransactionalAsync(insertResult);
 
     const unauthorizedEntityAfterInsert = new this.entityClass(this.viewerContext, insertResult);
     const newEntity = await entityLoader
@@ -395,6 +396,11 @@ export class UpdateMutator<
       entityLoader.invalidateFieldsAsync.bind(entityLoader, this.fieldsForEntity)
     );
 
+    await Promise.all([
+      entityLoader.invalidateFieldsTransactionalAsync(this.originalFieldsForEntity),
+      entityLoader.invalidateFieldsTransactionalAsync(this.fieldsForEntity),
+    ]);
+
     const unauthorizedEntityAfterUpdate = new this.entityClass(this.viewerContext, updateResult);
     const updatedEntity = await entityLoader
       .enforcing()
@@ -566,6 +572,7 @@ export class DeleteMutator<
     queryContext.appendPostCommitCallback(
       entityLoader.invalidateFieldsAsync.bind(entityLoader, this.entity.getAllDatabaseFields())
     );
+    await entityLoader.invalidateFieldsTransactionalAsync(this.entity.getAllDatabaseFields());
 
     await this.executeMutationTriggersOrValidatorsAsync(
       this.mutationTriggers.afterDelete,

packages/entity/src/EntityQueryContext.ts

@@ -1,3 +1,5 @@
+import { v4 as uuidv4 } from 'uuid';
+
 import EntityQueryContextProvider from './EntityQueryContextProvider';
 
 export type PostCommitCallback = (...args: any) => Promise<any>;
@@ -12,7 +14,7 @@ export type PostCommitCallback = (...args: any) => Promise<any>;
 export abstract class EntityQueryContext {
   constructor(private readonly queryInterface: any) {}
 
-  abstract isInTransaction(): boolean;
+  abstract isInTransaction(): this is EntityTransactionalQueryContext;
 
   getQueryInterface(): any {
     return this.queryInterface;
@@ -31,7 +33,7 @@ export class EntityNonTransactionalQueryContext extends EntityQueryContext {
     super(queryInterface);
   }
 
-  isInTransaction(): boolean {
+  isInTransaction(): this is EntityTransactionalQueryContext {
     return false;
   }
 
@@ -43,6 +45,8 @@ export class EntityNonTransactionalQueryContext extends EntityQueryContext {
 }
 
 export class EntityTransactionalQueryContext extends EntityQueryContext {
+  public readonly transactionID = uuidv4();
+
   private readonly postCommitCallbacks: PostCommitCallback[] = [];
 
   public appendPostCommitCallback(callback: PostCommitCallback): void {
@@ -55,7 +59,7 @@ export class EntityTransactionalQueryContext extends EntityQueryContext {
     await Promise.all(callbacks.map((callback) => callback()));
   }
 
-  isInTransaction(): boolean {
+  isInTransaction(): this is EntityTransactionalQueryContext {
     return true;
   }
 

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

@@ -287,6 +287,7 @@ describe(EnforcingEntityLoader, () => {
       'enforcing',
       'invalidateFieldsAsync',
       'invalidateEntityAsync',
+      'invalidateFieldsTransactionalAsync',
       'tryConstructEntities',
     ];
     expect(loaderProperties).toEqual(expect.arrayContaining(knownLoaderOnlyDifferences));

packages/entity/src/internal/EntityDataManager.ts

@@ -4,7 +4,7 @@ import EntityDatabaseAdapter, {
   FieldEqualityCondition,
   QuerySelectionModifiers,
 } from '../EntityDatabaseAdapter';
-import { EntityQueryContext } from '../EntityQueryContext';
+import { EntityQueryContext, EntityTransactionalQueryContext } from '../EntityQueryContext';
 import EntityQueryContextProvider from '../EntityQueryContextProvider';
 import { partitionErrors } from '../entityUtils';
 import {
@@ -15,6 +15,13 @@ import IEntityMetricsAdapter, { EntityMetricsLoadType } from '../metrics/IEntity
 import { computeIfAbsent, zipToMap } from '../utils/collections/maps';
 import ReadThroughEntityCache from './ReadThroughEntityCache';
 
+type FieldDataLoader<TFields> = DataLoader<
+  NonNullable<TFields[keyof TFields]>,
+  readonly Readonly<TFields>[]
+>;
+
+type FieldDataLoadersMap<TFields> = Map<keyof TFields, FieldDataLoader<TFields>>;
+
 /**
  * A data manager is responsible for orchestrating multiple sources of entity
  * data including local caches, {@link EntityCacheAdapter}, and {@link EntityDatabaseAdapter}.
@@ -22,9 +29,11 @@ import ReadThroughEntityCache from './ReadThroughEntityCache';
  * It is also responsible for invalidating all sources of data when mutated using {@link EntityMutator}.
  */
 export default class EntityDataManager<TFields> {
-  private readonly fieldDataLoaders: Map<
-    keyof TFields,
-    DataLoader<NonNullable<TFields[keyof TFields]>, readonly Readonly<TFields>[]>
+  private readonly regularFieldDataLoaders: FieldDataLoadersMap<TFields> = new Map();
+
+  private readonly transactionalFieldDataLoaders: Map<
+    string,
+    FieldDataLoadersMap<TFields>
   > = new Map();
 
   constructor(
@@ -35,10 +44,10 @@ export default class EntityDataManager<TFields> {
     private readonly entityClassName: string
   ) {}
 
-  private getFieldDataLoaderForFieldName<N extends keyof TFields>(
+  private getRegularFieldDataLoderForFieldName<N extends keyof TFields>(
     fieldName: N
-  ): DataLoader<NonNullable<TFields[N]>, readonly Readonly<TFields>[]> {
-    return computeIfAbsent(this.fieldDataLoaders, fieldName, () => {
+  ): FieldDataLoader<TFields> {
+    return computeIfAbsent(this.regularFieldDataLoaders, fieldName, () => {
       return new DataLoader(
         async (
           fieldValues: readonly NonNullable<TFields[N]>[]
@@ -53,6 +62,43 @@ export default class EntityDataManager<TFields> {
     });
   }
 
+  private getTransactionalFieldDataLoaderForFieldName<N extends keyof TFields>(
+    queryContext: EntityTransactionalQueryContext,
+    fieldName: N
+  ): FieldDataLoader<TFields> {
+    const fieldDataLoaderMap = computeIfAbsent(
+      this.transactionalFieldDataLoaders,
+      queryContext.transactionID,
+      () => new Map<keyof TFields, FieldDataLoader<TFields>>()
+    );
+    return computeIfAbsent(fieldDataLoaderMap, fieldName, () => {
+      return new DataLoader(
+        async (
+          fieldValues: readonly NonNullable<TFields[N]>[]
+        ): Promise<readonly (readonly TFields[])[]> => {
+          const objectMap = await this.loadManyForTransactionalDataLoaderByFieldEqualingAsync(
+            fieldName,
+            fieldValues,
+            queryContext
+          );
+          return fieldValues.map((fv) => objectMap.get(fv) ?? []);
+        }
+      );
+    });
+  }
+
+  private async loadManyForTransactionalDataLoaderByFieldEqualingAsync<N extends keyof TFields>(
+    fieldName: N,
+    fieldValues: readonly NonNullable<TFields[N]>[],
+    queryContext: EntityTransactionalQueryContext
+  ): Promise<ReadonlyMap<NonNullable<TFields[N]>, readonly Readonly<TFields>[]>> {
+    this.metricsAdapter.incrementDataManagerDatabaseLoadCount({
+      fieldValueCount: fieldValues.length,
+      entityClassName: this.entityClassName,
+    });
+    return await this.databaseAdapter.fetchManyWhereAsync(queryContext, fieldName, fieldValues);
+  }
+
   private async loadManyForDataLoaderByFieldEqualingAsync<N extends keyof TFields>(
     fieldName: N,
     fieldValues: readonly NonNullable<TFields[N]>[]
@@ -61,10 +107,11 @@ export default class EntityDataManager<TFields> {
       fieldValueCount: fieldValues.length,
       entityClassName: this.entityClassName,
     });
+
     return await this.entityCache.readManyThroughAsync(
       fieldName,
       fieldValues,
-      async (fetcherValues) => {
+      async (fetcherValues: readonly NonNullable<TFields[N]>[]) => {
         this.metricsAdapter.incrementDataManagerDatabaseLoadCount({
           fieldValueCount: fieldValues.length,
           entityClassName: this.entityClassName,
@@ -112,16 +159,13 @@ export default class EntityDataManager<TFields> {
       );
     }
 
-    // don't cache when in transaction, as rollbacks complicate things significantly
-    if (queryContext.isInTransaction()) {
-      return await this.databaseAdapter.fetchManyWhereAsync(queryContext, fieldName, fieldValues);
-    }
-
     this.metricsAdapter.incrementDataManagerDataloaderLoadCount({
       fieldValueCount: fieldValues.length,
       entityClassName: this.entityClassName,
     });
-    const dataLoader = this.getFieldDataLoaderForFieldName(fieldName);
+    const dataLoader = queryContext.isInTransaction()
+      ? this.getTransactionalFieldDataLoaderForFieldName(queryContext, fieldName)
+      : this.getRegularFieldDataLoderForFieldName(fieldName);
     const results = await dataLoader.loadMany(fieldValues);
     const [values, errors] = partitionErrors(results);
     if (errors.length > 0) {
@@ -192,7 +236,20 @@ export default class EntityDataManager<TFields> {
     fieldValues: readonly NonNullable<TFields[N]>[]
   ): Promise<void> {
     await this.entityCache.invalidateManyAsync(fieldName, fieldValues);
-    const dataLoader = this.getFieldDataLoaderForFieldName(fieldName);
+    const dataLoader = this.getRegularFieldDataLoderForFieldName(fieldName);
+    fieldValues.forEach((fieldValue) => dataLoader.clear(fieldValue));
+  }
+
+  private async invalidateManyByFieldEqualingTransactionalAsync<N extends keyof TFields>(
+    queryContext: EntityQueryContext,
+    fieldName: N,
+    fieldValues: readonly NonNullable<TFields[N]>[]
+  ): Promise<void> {
+    if (!queryContext.isInTransaction()) {
+      return;
+    }
+
+    const dataLoader = this.getTransactionalFieldDataLoaderForFieldName(queryContext, fieldName);
     fieldValues.forEach((fieldValue) => dataLoader.clear(fieldValue));
   }
 
@@ -215,4 +272,21 @@ export default class EntityDataManager<TFields> {
       })
     );
   }
+
+  async invalidateObjectFieldsTransactionalAsync(
+    queryContext: EntityQueryContext,
+    objectFields: Readonly<TFields>
+  ): Promise<void> {
+    const keys = Object.keys(objectFields) as (keyof TFields)[];
+    await Promise.all(
+      keys.map(async (fieldName: keyof TFields) => {
+        const value = objectFields[fieldName];
+        if (value !== undefined) {
+          await this.invalidateManyByFieldEqualingTransactionalAsync(queryContext, fieldName, [
+            value as NonNullable<TFields[keyof TFields]>,
+          ]);
+        }
+      })
+    );
+  }
 }