@expo/entity 0.31.1 → 0.32.0

package/src/EntityCompanion.ts

@@ -1,7 +1,5 @@
-import { IEntityClass } from './Entity';
+import EntityCompanionProvider, { EntityCompanionDefinition } from './EntityCompanionProvider';
 import EntityLoaderFactory from './EntityLoaderFactory';
-import EntityMutationTriggerConfiguration from './EntityMutationTriggerConfiguration';
-import EntityMutationValidator from './EntityMutationValidator';
 import EntityMutatorFactory from './EntityMutatorFactory';
 import EntityPrivacyPolicy from './EntityPrivacyPolicy';
 import EntityQueryContextProvider from './EntityQueryContextProvider';
@@ -18,7 +16,7 @@ export interface IPrivacyPolicyClass<TPrivacyPolicy> {
  * Composition root responsible for orchestrating setup of Entity mutators and loaders.
  */
 export default class EntityCompanion<
-  TFields,
+  TFields extends object,
   TID extends NonNullable<TFields[TSelectedFields]>,
   TViewerContext extends ViewerContext,
   TEntity extends ReadonlyEntity<TFields, TID, TViewerContext, TSelectedFields>,
@@ -31,6 +29,8 @@ export default class EntityCompanion<
   >,
   TSelectedFields extends keyof TFields
 > {
+  public readonly privacyPolicy: TPrivacyPolicy;
+
   private readonly entityLoaderFactory: EntityLoaderFactory<
     TFields,
     TID,
@@ -49,7 +49,8 @@ export default class EntityCompanion<
   >;
 
   constructor(
-    entityClass: IEntityClass<
+    public readonly entityCompanionProvider: EntityCompanionProvider,
+    public readonly entityCompanionDefinition: EntityCompanionDefinition<
       TFields,
       TID,
       TViewerContext,
@@ -58,24 +59,9 @@ export default class EntityCompanion<
       TSelectedFields
     >,
     private readonly tableDataCoordinator: EntityTableDataCoordinator<TFields>,
-    PrivacyPolicyClass: IPrivacyPolicyClass<TPrivacyPolicy>,
-    mutationValidators: EntityMutationValidator<
-      TFields,
-      TID,
-      TViewerContext,
-      TEntity,
-      TSelectedFields
-    >[],
-    mutationTriggers: EntityMutationTriggerConfiguration<
-      TFields,
-      TID,
-      TViewerContext,
-      TEntity,
-      TSelectedFields
-    >,
     private readonly metricsAdapter: IEntityMetricsAdapter
   ) {
-    const privacyPolicy = new PrivacyPolicyClass();
+    this.privacyPolicy = new entityCompanionDefinition.privacyPolicyClass();
     this.entityLoaderFactory = new EntityLoaderFactory<
       TFields,
       TID,
@@ -83,19 +69,14 @@ export default class EntityCompanion<
       TEntity,
       TPrivacyPolicy,
       TSelectedFields
-    >(
-      tableDataCoordinator.entityConfiguration,
-      entityClass,
-      privacyPolicy,
-      tableDataCoordinator.dataManager,
-      metricsAdapter
-    );
+    >(this, tableDataCoordinator.dataManager, metricsAdapter);
     this.entityMutatorFactory = new EntityMutatorFactory(
+      entityCompanionProvider,
       tableDataCoordinator.entityConfiguration,
-      entityClass,
-      privacyPolicy,
-      mutationValidators,
-      mutationTriggers,
+      entityCompanionDefinition.entityClass,
+      this.privacyPolicy,
+      entityCompanionDefinition.mutationValidators ?? [],
+      entityCompanionDefinition.mutationTriggers ?? {},
       this.entityLoaderFactory,
       tableDataCoordinator.databaseAdapter,
       metricsAdapter

package/src/EntityCompanionProvider.ts

@@ -49,8 +49,8 @@ export interface CacheAdapterFlavorDefinition {
  * Definition for constructing a companion for an entity. Defines the core set of objects
  * used to power the entity framework for a particular type of entity.
  */
-export class EntityCompanionDefinition<
-  TFields,
+export interface EntityCompanionDefinition<
+  TFields extends object,
   TID extends NonNullable<TFields[TSelectedFields]>,
   TViewerContext extends ViewerContext,
   TEntity extends ReadonlyEntity<TFields, TID, TViewerContext, TSelectedFields>,
@@ -63,6 +63,9 @@ export class EntityCompanionDefinition<
   >,
   TSelectedFields extends keyof TFields = keyof TFields
 > {
+  /**
+   * The concrete Entity class for which this is the definition.
+   */
   readonly entityClass: IEntityClass<
     TFields,
     TID,
@@ -71,85 +74,45 @@ export class EntityCompanionDefinition<
     TPrivacyPolicy,
     TSelectedFields
   >;
+
+  /**
+   * The EntityConfiguration for this entity.
+   */
   readonly entityConfiguration: EntityConfiguration<TFields>;
+
+  /**
+   * The EntityPrivacyPolicy class for this entity.
+   */
   readonly privacyPolicyClass: IPrivacyPolicyClass<TPrivacyPolicy>;
-  readonly mutationValidators: () => EntityMutationValidator<
+
+  /**
+   * An optional list of EntityMutationValidator for this entity.
+   */
+  readonly mutationValidators?: EntityMutationValidator<
     TFields,
     TID,
     TViewerContext,
     TEntity,
     TSelectedFields
   >[];
-  readonly mutationTriggers: () => EntityMutationTriggerConfiguration<
+
+  /**
+   * An optional list of EntityMutationTrigger for this entity.
+   */
+  readonly mutationTriggers?: EntityMutationTriggerConfiguration<
     TFields,
     TID,
     TViewerContext,
     TEntity,
     TSelectedFields
   >;
-  readonly entitySelectedFields: TSelectedFields[];
 
-  constructor({
-    entityClass,
-    entityConfiguration,
-    privacyPolicyClass,
-    mutationValidators = () => [],
-    mutationTriggers = () => ({}),
-    entitySelectedFields = Array.from(entityConfiguration.schema.keys()) as TSelectedFields[],
-  }: {
-    /**
-     * The concrete Entity class for which this is the definition.
-     */
-    entityClass: IEntityClass<
-      TFields,
-      TID,
-      TViewerContext,
-      TEntity,
-      TPrivacyPolicy,
-      TSelectedFields
-    >;
-    /**
-     * The EntityConfiguration for this entity.
-     */
-    entityConfiguration: EntityConfiguration<TFields>;
-    /**
-     * The EntityPrivacyPolicy class for this entity.
-     */
-    privacyPolicyClass: IPrivacyPolicyClass<TPrivacyPolicy>;
-    /**
-     * An optional list of EntityMutationValidator for this entity.
-     */
-    mutationValidators?: () => EntityMutationValidator<
-      TFields,
-      TID,
-      TViewerContext,
-      TEntity,
-      TSelectedFields
-    >[];
-    /**
-     * An optional list of EntityMutationTrigger for this entity.
-     */
-    mutationTriggers?: () => EntityMutationTriggerConfiguration<
-      TFields,
-      TID,
-      TViewerContext,
-      TEntity,
-      TSelectedFields
-    >;
-    /**
-     * An optional subset of fields defined in the EntityConfiguration which belong to this entity.
-     * For use when multiple types of entities are backed by a single table (EntityConfiguration) yet
-     * only expose a subset of the fields.
-     */
-    entitySelectedFields?: TSelectedFields[];
-  }) {
-    this.entityClass = entityClass;
-    this.entityConfiguration = entityConfiguration;
-    this.privacyPolicyClass = privacyPolicyClass;
-    this.mutationValidators = mutationValidators;
-    this.mutationTriggers = mutationTriggers;
-    this.entitySelectedFields = entitySelectedFields;
-  }
+  /**
+   * An optional subset of fields defined in the EntityConfiguration which belong to this entity.
+   * For use when multiple types of entities are backed by a single table (EntityConfiguration) yet
+   * only expose a subset of the fields.
+   */
+  readonly entitySelectedFields?: TSelectedFields[];
 }
 
 /**
@@ -162,6 +125,10 @@ export class EntityCompanionDefinition<
  * EntityCompanion for each type of Entity.
  */
 export default class EntityCompanionProvider {
+  private readonly companionDefinitionMap: Map<
+    string,
+    EntityCompanionDefinition<any, any, any, any, any, any>
+  > = new Map();
   private readonly companionMap: Map<string, EntityCompanion<any, any, any, any, any, any>> =
     new Map();
   private readonly tableDataCoordinatorMap: Map<string, EntityTableDataCoordinator<any>> =
@@ -187,10 +154,9 @@ export default class EntityCompanionProvider {
    * companion is constructed using the configuration provided by the factory.
    *
    * @param entityClass - entity class to load
-   * @param factory - entity companion factory
    */
   getCompanionForEntity<
-    TFields,
+    TFields extends object,
     TID extends NonNullable<TFields[TSelectedFields]>,
     TViewerContext extends ViewerContext,
     TEntity extends ReadonlyEntity<TFields, TID, TViewerContext, TSelectedFields>,
@@ -210,27 +176,22 @@ export default class EntityCompanionProvider {
       TEntity,
       TPrivacyPolicy,
       TSelectedFields
-    >,
-    entityCompanionDefinition: EntityCompanionDefinition<
-      TFields,
-      TID,
-      TViewerContext,
-      TEntity,
-      TPrivacyPolicy,
-      TSelectedFields
     >
   ): EntityCompanion<TFields, TID, TViewerContext, TEntity, TPrivacyPolicy, TSelectedFields> {
+    const entityCompanionDefinition = computeIfAbsent(
+      this.companionDefinitionMap,
+      entityClass.name,
+      () => entityClass.defineCompanionDefinition()
+    );
     const tableDataCoordinator = this.getTableDataCoordinatorForEntity(
       entityCompanionDefinition.entityConfiguration,
       entityClass.name
     );
     return computeIfAbsent(this.companionMap, entityClass.name, () => {
       return new EntityCompanion(
-        entityCompanionDefinition.entityClass,
+        this,
+        entityCompanionDefinition,
         tableDataCoordinator,
-        entityCompanionDefinition.privacyPolicyClass,
-        entityCompanionDefinition.mutationValidators(),
-        entityCompanionDefinition.mutationTriggers(),
         this.metricsAdapter
       );
     });

package/src/EntityConfiguration.ts

@@ -13,7 +13,7 @@ export default class EntityConfiguration<TFields> {
   readonly cacheableKeys: ReadonlySet<keyof TFields>;
   readonly cacheKeyVersion: number;
 
-  readonly getInboundEdges: () => IEntityClass<any, any, any, any, any, any>[];
+  readonly inboundEdges: IEntityClass<any, any, any, any, any, any>[];
   readonly schema: ReadonlyMap<keyof TFields, EntityFieldDefinition<any>>;
   readonly entityToDBFieldsKeyMapping: ReadonlyMap<keyof TFields, string>;
   readonly dbToEntityFieldsKeyMapping: ReadonlyMap<string, keyof TFields>;
@@ -25,7 +25,7 @@ export default class EntityConfiguration<TFields> {
     idField,
     tableName,
     schema,
-    getInboundEdges = () => [],
+    inboundEdges = [],
     cacheKeyVersion = 0,
     databaseAdapterFlavor,
     cacheAdapterFlavor,
@@ -48,7 +48,7 @@ export default class EntityConfiguration<TFields> {
     /**
      * List of other entity types that reference this type in EntityFieldDefinition associations.
      */
-    getInboundEdges?: () => IEntityClass<any, any, any, any, any, any>[];
+    inboundEdges?: IEntityClass<any, any, any, any, any, any>[];
 
     /**
      * Cache key version for this entity type. Should be bumped when a field is added to, removed from, or changed
@@ -71,8 +71,7 @@ export default class EntityConfiguration<TFields> {
     this.cacheKeyVersion = cacheKeyVersion;
     this.databaseAdapterFlavor = databaseAdapterFlavor;
     this.cacheAdapterFlavor = cacheAdapterFlavor;
-
-    this.getInboundEdges = getInboundEdges;
+    this.inboundEdges = inboundEdges;
 
     // external schema is a Record to typecheck that all fields have FieldDefinitions,
     // but internally the most useful representation is a map for lookups

package/src/EntityFieldDefinition.ts

@@ -41,7 +41,7 @@ export enum EntityEdgeDeletionBehavior {
  */
 export interface EntityAssociationDefinition<
   TViewerContext extends ViewerContext,
-  TAssociatedFields,
+  TAssociatedFields extends object,
   TAssociatedID extends NonNullable<TAssociatedFields[TAssociatedSelectedFields]>,
   TAssociatedEntity extends ReadonlyEntity<
     TAssociatedFields,
@@ -61,7 +61,7 @@ export interface EntityAssociationDefinition<
   /**
    * Class of entity on the other end of this edge.
    */
-  getAssociatedEntityClass: () => IEntityClass<
+  associatedEntityClass: IEntityClass<
     TAssociatedFields,
     TAssociatedID,
     TViewerContext,

package/src/EntityLoader.ts

@@ -1,5 +1,6 @@
 import { Result, asyncResult, result } from '@expo/results';
 import invariant from 'invariant';
+import nullthrows from 'nullthrows';
 
 import EnforcingEntityLoader from './EnforcingEntityLoader';
 import { IEntityClass } from './Entity';
@@ -14,6 +15,7 @@ import EntityPrivacyPolicy, { EntityPrivacyPolicyEvaluationContext } from './Ent
 import { EntityQueryContext } from './EntityQueryContext';
 import ReadonlyEntity from './ReadonlyEntity';
 import ViewerContext from './ViewerContext';
+import { pick } from './entityUtils';
 import EntityInvalidFieldValueError from './errors/EntityInvalidFieldValueError';
 import EntityNotFoundError from './errors/EntityNotFoundError';
 import EntityDataManager from './internal/EntityDataManager';
@@ -25,7 +27,7 @@ import { mapMap, mapMapAsync } from './utils/collections/maps';
  * cached, and authorized against the entity's EntityPrivacyPolicy.
  */
 export default class EntityLoader<
-  TFields,
+  TFields extends object,
   TID extends NonNullable<TFields[TSelectedFields]>,
   TViewerContext extends ViewerContext,
   TEntity extends ReadonlyEntity<TFields, TID, TViewerContext, TSelectedFields>,
@@ -51,6 +53,7 @@ export default class EntityLoader<
       TPrivacyPolicy,
       TSelectedFields
     >,
+    private readonly entitySelectedFields: TSelectedFields[] | undefined,
     private readonly privacyPolicy: TPrivacyPolicy,
     private readonly dataManager: EntityDataManager<TFields>,
     protected readonly metricsAdapter: IEntityMetricsAdapter
@@ -332,7 +335,7 @@ export default class EntityLoader<
   private tryConstructEntities(fieldsObjects: readonly TFields[]): readonly Result<TEntity>[] {
     return fieldsObjects.map((fieldsObject) => {
       try {
-        return result(new this.entityClass(this.viewerContext, fieldsObject));
+        return result(this.constructEntity(fieldsObject));
       } catch (e) {
         if (!(e instanceof Error)) {
           throw e;
@@ -342,6 +345,20 @@ export default class EntityLoader<
     });
   }
 
+  public constructEntity(fieldsObject: TFields): TEntity {
+    const idField = this.entityConfiguration.idField;
+    const id = nullthrows(fieldsObject[idField], 'must provide ID to create an entity');
+    const entitySelectedFields =
+      this.entitySelectedFields ?? Array.from(this.entityConfiguration.schema.keys());
+    const selectedFields = pick(fieldsObject, entitySelectedFields);
+    return new this.entityClass({
+      viewerContext: this.viewerContext,
+      id: id as TID,
+      databaseFields: fieldsObject,
+      selectedFields,
+    });
+  }
+
   /**
    * Construct and authorize entities from fields map, returning error results for entities that fail
    * to construct or fail to authorize.

package/src/EntityLoaderFactory.ts

@@ -1,5 +1,4 @@
-import { IEntityClass } from './Entity';
-import EntityConfiguration from './EntityConfiguration';
+import EntityCompanion from './EntityCompanion';
 import EntityLoader from './EntityLoader';
 import EntityPrivacyPolicy, { EntityPrivacyPolicyEvaluationContext } from './EntityPrivacyPolicy';
 import { EntityQueryContext } from './EntityQueryContext';
@@ -12,7 +11,7 @@ import IEntityMetricsAdapter from './metrics/IEntityMetricsAdapter';
  * The primary entry point for loading entities.
  */
 export default class EntityLoaderFactory<
-  TFields,
+  TFields extends object,
   TID extends NonNullable<TFields[TSelectedFields]>,
   TViewerContext extends ViewerContext,
   TEntity extends ReadonlyEntity<TFields, TID, TViewerContext, TSelectedFields>,
@@ -26,8 +25,7 @@ export default class EntityLoaderFactory<
   TSelectedFields extends keyof TFields
 > {
   constructor(
-    private readonly entityConfiguration: EntityConfiguration<TFields>,
-    private readonly entityClass: IEntityClass<
+    private readonly entityCompanion: EntityCompanion<
       TFields,
       TID,
       TViewerContext,
@@ -35,7 +33,6 @@ export default class EntityLoaderFactory<
       TPrivacyPolicy,
       TSelectedFields
     >,
-    private readonly privacyPolicyClass: TPrivacyPolicy,
     private readonly dataManager: EntityDataManager<TFields>,
     protected readonly metricsAdapter: IEntityMetricsAdapter
   ) {}
@@ -54,9 +51,10 @@ export default class EntityLoaderFactory<
       viewerContext,
       queryContext,
       privacyPolicyEvaluationContext,
-      this.entityConfiguration,
-      this.entityClass,
-      this.privacyPolicyClass,
+      this.entityCompanion.entityCompanionDefinition.entityConfiguration,
+      this.entityCompanion.entityCompanionDefinition.entityClass,
+      this.entityCompanion.entityCompanionDefinition.entitySelectedFields,
+      this.entityCompanion.privacyPolicy,
       this.dataManager,
       this.metricsAdapter
     );

package/src/EntityMutationInfo.ts

@@ -8,7 +8,7 @@ export enum EntityMutationType {
 }
 
 export type EntityValidatorMutationInfo<
-  TFields,
+  TFields extends object,
   TID extends NonNullable<TFields[TSelectedFields]>,
   TViewerContext extends ViewerContext,
   TEntity extends Entity<TFields, TID, TViewerContext, TSelectedFields>,
@@ -39,7 +39,7 @@ export type EntityCascadingDeletionInfo = {
 };
 
 export type EntityTriggerMutationInfo<
-  TFields,
+  TFields extends object,
   TID extends NonNullable<TFields[TSelectedFields]>,
   TViewerContext extends ViewerContext,
   TEntity extends Entity<TFields, TID, TViewerContext, TSelectedFields>,

package/src/EntityMutationTriggerConfiguration.ts

@@ -7,7 +7,7 @@ import ViewerContext from './ViewerContext';
  * Interface to define trigger behavior for entities.
  */
 export default interface EntityMutationTriggerConfiguration<
-  TFields,
+  TFields extends object,
   TID extends NonNullable<TFields[TSelectedFields]>,
   TViewerContext extends ViewerContext,
   TEntity extends ReadonlyEntity<TFields, TID, TViewerContext, TSelectedFields>,
@@ -70,7 +70,7 @@ export default interface EntityMutationTriggerConfiguration<
  * the transaction if a transaction is supplied.
  */
 export abstract class EntityMutationTrigger<
-  TFields,
+  TFields extends object,
   TID extends NonNullable<TFields[TSelectedFields]>,
   TViewerContext extends ViewerContext,
   TEntity extends ReadonlyEntity<TFields, TID, TViewerContext, TSelectedFields>,
@@ -89,7 +89,7 @@ export abstract class EntityMutationTrigger<
  * since they explicitly occur outside of the transaction.
  */
 export abstract class EntityNonTransactionalMutationTrigger<
-  TFields,
+  TFields extends object,
   TID extends NonNullable<TFields[TSelectedFields]>,
   TViewerContext extends ViewerContext,
   TEntity extends ReadonlyEntity<TFields, TID, TViewerContext, TSelectedFields>,

package/src/EntityMutationValidator.ts

@@ -8,7 +8,7 @@ import ViewerContext from './ViewerContext';
  * same transaction as the mutation itself before creating or updating an entity.
  */
 export default abstract class EntityMutationValidator<
-  TFields,
+  TFields extends object,
   TID extends NonNullable<TFields[TSelectedFields]>,
   TViewerContext extends ViewerContext,
   TEntity extends ReadonlyEntity<TFields, TID, TViewerContext, TSelectedFields>,

package/src/EntityMutator.ts

@@ -2,7 +2,7 @@ import { Result, asyncResult, result, enforceAsyncResult } from '@expo/results';
 import invariant from 'invariant';
 
 import Entity, { IEntityClass } from './Entity';
-import { EntityCompanionDefinition } from './EntityCompanionProvider';
+import EntityCompanionProvider from './EntityCompanionProvider';
 import EntityConfiguration from './EntityConfiguration';
 import EntityDatabaseAdapter from './EntityDatabaseAdapter';
 import { EntityEdgeDeletionBehavior } from './EntityFieldDefinition';
@@ -28,7 +28,7 @@ import IEntityMetricsAdapter, { EntityMetricsMutationType } from './metrics/IEnt
 import { mapMapAsync } from './utils/collections/maps';
 
 abstract class BaseMutator<
-  TFields,
+  TFields extends object,
   TID extends NonNullable<TFields[TSelectedFields]>,
   TViewerContext extends ViewerContext,
   TEntity extends Entity<TFields, TID, TViewerContext, TSelectedFields>,
@@ -42,6 +42,7 @@ abstract class BaseMutator<
   TSelectedFields extends keyof TFields
 > {
   constructor(
+    protected readonly companionProvider: EntityCompanionProvider,
     protected readonly viewerContext: TViewerContext,
     protected readonly queryContext: EntityQueryContext,
     protected readonly entityConfiguration: EntityConfiguration<TFields>,
@@ -155,7 +156,7 @@ abstract class BaseMutator<
  * Mutator for creating a new entity.
  */
 export class CreateMutator<
-  TFields,
+  TFields extends object,
   TID extends NonNullable<TFields[TSelectedFields]>,
   TViewerContext extends ViewerContext,
   TEntity extends Entity<TFields, TID, TViewerContext, TSelectedFields>,
@@ -211,7 +212,11 @@ export class CreateMutator<
   ): Promise<Result<TEntity>> {
     this.validateFields(this.fieldsForEntity);
 
-    const temporaryEntityForPrivacyCheck = new this.entityClass(this.viewerContext, {
+    const entityLoader = this.entityLoaderFactory.forLoad(this.viewerContext, queryContext, {
+      cascadingDeleteCause: null,
+    });
+
+    const temporaryEntityForPrivacyCheck = entityLoader.constructEntity({
       [this.entityConfiguration.idField]: '00000000-0000-0000-0000-000000000000', // zero UUID
       ...this.fieldsForEntity,
     } as unknown as TFields);
@@ -250,14 +255,11 @@ export class CreateMutator<
 
     const insertResult = await this.databaseAdapter.insertAsync(queryContext, this.fieldsForEntity);
 
-    const entityLoader = this.entityLoaderFactory.forLoad(this.viewerContext, queryContext, {
-      cascadingDeleteCause: null,
-    });
     queryContext.appendPostCommitInvalidationCallback(
       entityLoader.invalidateFieldsAsync.bind(entityLoader, insertResult)
     );
 
-    const unauthorizedEntityAfterInsert = new this.entityClass(this.viewerContext, insertResult);
+    const unauthorizedEntityAfterInsert = entityLoader.constructEntity(insertResult);
     const newEntity = await entityLoader
       .enforcing()
       .loadByIDAsync(unauthorizedEntityAfterInsert.getID());
@@ -292,7 +294,7 @@ export class CreateMutator<
  * Mutator for updating an existing entity.
  */
 export class UpdateMutator<
-  TFields,
+  TFields extends object,
   TID extends NonNullable<TFields[TSelectedFields]>,
   TViewerContext extends ViewerContext,
   TEntity extends Entity<TFields, TID, TViewerContext, TSelectedFields>,
@@ -310,6 +312,7 @@ export class UpdateMutator<
   private readonly updatedFields: Partial<TFields> = {};
 
   constructor(
+    companionProvider: EntityCompanionProvider,
     viewerContext: TViewerContext,
     queryContext: EntityQueryContext,
     entityConfiguration: EntityConfiguration<TFields>,
@@ -349,6 +352,7 @@ export class UpdateMutator<
     originalEntity: TEntity
   ) {
     super(
+      companionProvider,
       viewerContext,
       queryContext,
       entityConfiguration,
@@ -419,7 +423,11 @@ export class UpdateMutator<
   ): Promise<Result<TEntity>> {
     this.validateFields(this.updatedFields);
 
-    const entityAboutToBeUpdated = new this.entityClass(this.viewerContext, this.fieldsForEntity);
+    const entityLoader = this.entityLoaderFactory.forLoad(this.viewerContext, queryContext, {
+      cascadingDeleteCause,
+    });
+
+    const entityAboutToBeUpdated = entityLoader.constructEntity(this.fieldsForEntity);
     const authorizeUpdateResult = await asyncResult(
       this.privacyPolicy.authorizeUpdateAsync(
         this.viewerContext,
@@ -462,10 +470,6 @@ export class UpdateMutator<
           this.updatedFields
         );
 
-    const entityLoader = this.entityLoaderFactory.forLoad(this.viewerContext, queryContext, {
-      cascadingDeleteCause,
-    });
-
     queryContext.appendPostCommitInvalidationCallback(
       entityLoader.invalidateFieldsAsync.bind(
         entityLoader,
@@ -481,7 +485,7 @@ export class UpdateMutator<
     const updatedEntity = updateResult
       ? await entityLoader
           .enforcing()
-          .loadByIDAsync(new this.entityClass(this.viewerContext, updateResult).getID())
+          .loadByIDAsync(entityLoader.constructEntity(updateResult).getID())
       : entityAboutToBeUpdated;
 
     await this.executeMutationTriggersAsync(
@@ -518,7 +522,7 @@ export class UpdateMutator<
  * Mutator for deleting an existing entity.
  */
 export class DeleteMutator<
-  TFields,
+  TFields extends object,
   TID extends NonNullable<TFields[TSelectedFields]>,
   TViewerContext extends ViewerContext,
   TEntity extends Entity<TFields, TID, TViewerContext, TSelectedFields>,
@@ -532,6 +536,7 @@ export class DeleteMutator<
   TSelectedFields extends keyof TFields
 > extends BaseMutator<TFields, TID, TViewerContext, TEntity, TPrivacyPolicy, TSelectedFields> {
   constructor(
+    companionProvider: EntityCompanionProvider,
     viewerContext: TViewerContext,
     queryContext: EntityQueryContext,
     entityConfiguration: EntityConfiguration<TFields>,
@@ -571,6 +576,7 @@ export class DeleteMutator<
     private readonly entity: TEntity
   ) {
     super(
+      companionProvider,
       viewerContext,
       queryContext,
       entityConfiguration,
@@ -723,31 +729,32 @@ export class DeleteMutator<
     }
     processedEntityIdentifiers.add(entity.getUniqueIdentifier());
 
-    const companionDefinition = (
-      entity.constructor as any
-    ).getCompanionDefinition() as EntityCompanionDefinition<
-      TFields,
-      TID,
-      TViewerContext,
-      TEntity,
-      TPrivacyPolicy,
-      TSelectedFields
-    >;
+    const companionDefinition = this.companionProvider.getCompanionForEntity(
+      entity.constructor as IEntityClass<
+        TFields,
+        TID,
+        TViewerContext,
+        TEntity,
+        TPrivacyPolicy,
+        TSelectedFields
+      >
+    ).entityCompanionDefinition;
     const entityConfiguration = companionDefinition.entityConfiguration;
-    const inboundEdges = entityConfiguration.getInboundEdges();
+    const inboundEdges = entityConfiguration.inboundEdges;
     await Promise.all(
       inboundEdges.map(async (entityClass) => {
         return await mapMapAsync(
-          entityClass.getCompanionDefinition().entityConfiguration.schema,
+          this.companionProvider.getCompanionForEntity(entityClass).entityCompanionDefinition
+            .entityConfiguration.schema,
           async (fieldDefinition, fieldName) => {
             const association = fieldDefinition.association;
             if (!association) {
               return;
             }
 
-            const associatedConfiguration = association
-              .getAssociatedEntityClass()
-              .getCompanionDefinition().entityConfiguration;
+            const associatedConfiguration = this.companionProvider.getCompanionForEntity(
+              association.associatedEntityClass
+            ).entityCompanionDefinition.entityConfiguration;
             if (associatedConfiguration !== entityConfiguration) {
               return;
             }