chore: move knex-specific loader logic out of EntityDataManager

wschurman · wschurman · commit 0be9f1d58fa8 · 2026-01-31T10:15:35.000-07:00
diff --git a/packages/entity/src/AuthorizationResultBasedKnexEntityLoader.ts b/packages/entity/src/AuthorizationResultBasedKnexEntityLoader.ts
@@ -11,7 +11,7 @@ import { EntityPrivacyPolicy } from './EntityPrivacyPolicy';
 import { EntityQueryContext } from './EntityQueryContext';
 import { ReadonlyEntity } from './ReadonlyEntity';
 import { ViewerContext } from './ViewerContext';
-import { EntityDataManager } from './internal/EntityDataManager';
+import { EntityKnexDataManager } from './internal/EntityKnexDataManager';
 import { IEntityMetricsAdapter } from './metrics/IEntityMetricsAdapter';
 
 /**
@@ -36,7 +36,7 @@ export class AuthorizationResultBasedKnexEntityLoader<
 > {
   constructor(
     private readonly queryContext: EntityQueryContext,
-    private readonly dataManager: EntityDataManager<TFields, TIDField>,
+    private readonly knexDataManager: EntityKnexDataManager<TFields, TIDField>,
     protected readonly metricsAdapter: IEntityMetricsAdapter,
     public readonly utils: EntityLoaderUtils<
       TFields,
@@ -80,7 +80,7 @@ export class AuthorizationResultBasedKnexEntityLoader<
       this.utils.validateFieldAndValues(fieldEqualityOperand.fieldName, fieldValues);
     }
 
-    const fieldObjects = await this.dataManager.loadManyByFieldEqualityConjunctionAsync(
+    const fieldObjects = await this.knexDataManager.loadManyByFieldEqualityConjunctionAsync(
       this.queryContext,
       fieldEqualityOperands,
       querySelectionModifiers,
@@ -98,7 +98,7 @@ export class AuthorizationResultBasedKnexEntityLoader<
     bindings: any[] | object,
     querySelectionModifiers: QuerySelectionModifiersWithOrderByRaw<TFields> = {},
   ): Promise<readonly Result<TEntity>[]> {
-    const fieldObjects = await this.dataManager.loadManyByRawWhereClauseAsync(
+    const fieldObjects = await this.knexDataManager.loadManyByRawWhereClauseAsync(
       this.queryContext,
       rawWhereClause,
       bindings,
diff --git a/packages/entity/src/EntityCompanion.ts b/packages/entity/src/EntityCompanion.ts
@@ -87,7 +87,7 @@ export class EntityCompanion<
       TEntity,
       TPrivacyPolicy,
       TSelectedFields
-    >(this, tableDataCoordinator.dataManager, metricsAdapter);
+    >(this, tableDataCoordinator.dataManager, tableDataCoordinator.knexDataManager, metricsAdapter);
     this.entityMutatorFactory = new EntityMutatorFactory(
       entityCompanionProvider,
       tableDataCoordinator.entityConfiguration,
diff --git a/packages/entity/src/KnexEntityLoaderFactory.ts b/packages/entity/src/KnexEntityLoaderFactory.ts
@@ -6,6 +6,7 @@ import { EntityQueryContext } from './EntityQueryContext';
 import { ReadonlyEntity } from './ReadonlyEntity';
 import { ViewerContext } from './ViewerContext';
 import { EntityDataManager } from './internal/EntityDataManager';
+import { EntityKnexDataManager } from './internal/EntityKnexDataManager';
 import { IEntityMetricsAdapter } from './metrics/IEntityMetricsAdapter';
 
 /**
@@ -35,6 +36,7 @@ export class KnexEntityLoaderFactory<
       TSelectedFields
     >,
     private readonly dataManager: EntityDataManager<TFields, TIDField>,
+    private readonly knexDataManager: EntityKnexDataManager<TFields, TIDField>,
     protected readonly metricsAdapter: IEntityMetricsAdapter,
   ) {}
 
@@ -75,7 +77,7 @@ export class KnexEntityLoaderFactory<
 
     return new AuthorizationResultBasedKnexEntityLoader(
       queryContext,
-      this.dataManager,
+      this.knexDataManager,
       this.metricsAdapter,
       utils,
     );
diff --git a/packages/entity/src/__tests__/AuthorizationResultBasedKnexEntityLoader-test.ts b/packages/entity/src/__tests__/AuthorizationResultBasedKnexEntityLoader-test.ts
@@ -9,6 +9,7 @@ import { EntityPrivacyPolicyEvaluationContext } from '../EntityPrivacyPolicy';
 import { ViewerContext } from '../ViewerContext';
 import { enforceResultsAsync } from '../entityUtils';
 import { EntityDataManager } from '../internal/EntityDataManager';
+import { EntityKnexDataManager } from '../internal/EntityKnexDataManager';
 import { ReadThroughEntityCache } from '../internal/ReadThroughEntityCache';
 import { IEntityMetricsAdapter } from '../metrics/IEntityMetricsAdapter';
 import { NoCacheStubCacheAdapterProvider } from '../utils/__testfixtures__/StubCacheAdapter';
@@ -90,6 +91,11 @@ describe(AuthorizationResultBasedKnexEntityLoader, () => {
       instance(mock<IEntityMetricsAdapter>()),
       TestEntity.name,
     );
+    const knexDataManager = new EntityKnexDataManager(
+      databaseAdapter,
+      metricsAdapter,
+      TestEntity.name,
+    );
     const utils = new EntityLoaderUtils(
       viewerContext,
       queryContext,
@@ -103,7 +109,7 @@ describe(AuthorizationResultBasedKnexEntityLoader, () => {
     );
     const knexEntityLoader = new AuthorizationResultBasedKnexEntityLoader(
       queryContext,
-      dataManager,
+      knexDataManager,
       metricsAdapter,
       utils,
     );
@@ -205,6 +211,11 @@ describe(AuthorizationResultBasedKnexEntityLoader, () => {
       instance(mock<IEntityMetricsAdapter>()),
       TestEntity.name,
     );
+    const knexDataManager = new EntityKnexDataManager(
+      databaseAdapter,
+      metricsAdapter,
+      TestEntity.name,
+    );
     const utils = new EntityLoaderUtils(
       viewerContext,
       queryContext,
@@ -218,7 +229,7 @@ describe(AuthorizationResultBasedKnexEntityLoader, () => {
     );
     const knexEntityLoader = new AuthorizationResultBasedKnexEntityLoader(
       queryContext,
-      dataManager,
+      knexDataManager,
       metricsAdapter,
       utils,
     );
@@ -271,9 +282,10 @@ describe(AuthorizationResultBasedKnexEntityLoader, () => {
     const metricsAdapter = instance(mock<IEntityMetricsAdapter>());
     const queryContext = new StubQueryContextProvider().getQueryContext();
 
-    const dataManagerMock = mock<EntityDataManager<TestFields, 'customIdField'>>(EntityDataManager);
+    const knexDataManagerMock =
+      mock<EntityKnexDataManager<TestFields, 'customIdField'>>(EntityKnexDataManager);
     when(
-      dataManagerMock.loadManyByRawWhereClauseAsync(
+      knexDataManagerMock.loadManyByRawWhereClauseAsync(
         queryContext,
         anything(),
         anything(),
@@ -289,6 +301,10 @@ describe(AuthorizationResultBasedKnexEntityLoader, () => {
         nullableField: null,
       },
     ]);
+    const knexDataManager = instance(knexDataManagerMock);
+
+    // Create a real dataManager for the EntityLoaderUtils
+    const dataManagerMock = mock<EntityDataManager<TestFields, 'customIdField'>>(EntityDataManager);
     const dataManager = instance(dataManagerMock);
 
     const utils = new EntityLoaderUtils(
@@ -304,7 +320,7 @@ describe(AuthorizationResultBasedKnexEntityLoader, () => {
     );
     const knexEntityLoader = new AuthorizationResultBasedKnexEntityLoader(
       queryContext,
-      dataManager,
+      knexDataManager,
       metricsAdapter,
       utils,
     );
diff --git a/packages/entity/src/index.ts b/packages/entity/src/index.ts
@@ -64,6 +64,7 @@ export * from './internal/CompositeFieldHolder';
 export * from './internal/CompositeFieldValueMap';
 export * from './internal/EntityDataManager';
 export * from './internal/EntityFieldTransformationUtils';
+export * from './internal/EntityKnexDataManager';
 export * from './internal/EntityLoadInterfaces';
 export * from './internal/EntityTableDataCoordinator';
 export * from './internal/ReadThroughEntityCache';
diff --git a/packages/entity/src/internal/EntityDataManager.ts b/packages/entity/src/internal/EntityDataManager.ts
@@ -1,12 +1,7 @@
 import DataLoader from 'dataloader';
 import invariant from 'invariant';
 
-import {
-  EntityDatabaseAdapter,
-  FieldEqualityCondition,
-  QuerySelectionModifiers,
-  QuerySelectionModifiersWithOrderByRaw,
-} from '../EntityDatabaseAdapter';
+import { EntityDatabaseAdapter } from '../EntityDatabaseAdapter';
 import {
   EntityQueryContext,
   EntityTransactionalQueryContext,
@@ -16,10 +11,7 @@ import { EntityQueryContextProvider } from '../EntityQueryContextProvider';
 import { partitionErrors } from '../entityUtils';
 import { IEntityLoadKey, IEntityLoadValue, LoadPair } from './EntityLoadInterfaces';
 import { ReadThroughEntityCache } from './ReadThroughEntityCache';
-import {
-  timeAndLogLoadEventAsync,
-  timeAndLogLoadMapEventAsync,
-} from '../metrics/EntityMetricsUtils';
+import { timeAndLogLoadMapEventAsync } from '../metrics/EntityMetricsUtils';
 import {
   EntityMetricsLoadType,
   IEntityMetricsAdapter,
@@ -251,64 +243,6 @@ export class EntityDataManager<
     return mapToReturn;
   }
 
-  /**
-   * Loads many objects matching the conjunction of where clauses constructed from
-   * specified field equality operands.
-   *
-   * @param queryContext - query context in which to perform the load
-   * @param fieldEqualityOperands - list of field equality where clause operand specifications
-   * @param querySelectionModifiers - limit, offset, and orderBy for the query
-   * @returns array of objects matching the query
-   */
-  async loadManyByFieldEqualityConjunctionAsync<N extends keyof TFields>(
-    queryContext: EntityQueryContext,
-    fieldEqualityOperands: FieldEqualityCondition<TFields, N>[],
-    querySelectionModifiers: QuerySelectionModifiers<TFields>,
-  ): Promise<readonly Readonly<TFields>[]> {
-    return await timeAndLogLoadEventAsync(
-      this.metricsAdapter,
-      EntityMetricsLoadType.LOAD_MANY_EQUALITY_CONJUNCTION,
-      this.entityClassName,
-      queryContext,
-    )(
-      this.databaseAdapter.fetchManyByFieldEqualityConjunctionAsync(
-        queryContext,
-        fieldEqualityOperands,
-        querySelectionModifiers,
-      ),
-    );
-  }
-
-  /**
-   * Loads many objects matching the raw WHERE clause.
-   *
-   * @param queryContext - query context in which to perform the load
-   * @param rawWhereClause - parameterized SQL WHERE clause with positional binding placeholders or named binding placeholders
-   * @param bindings - array of positional bindings or object of named bindings
-   * @param querySelectionModifiers - limit, offset, orderBy, and orderByRaw for the query
-   * @returns array of objects matching the query
-   */
-  async loadManyByRawWhereClauseAsync(
-    queryContext: EntityQueryContext,
-    rawWhereClause: string,
-    bindings: any[] | object,
-    querySelectionModifiers: QuerySelectionModifiersWithOrderByRaw<TFields>,
-  ): Promise<readonly Readonly<TFields>[]> {
-    return await timeAndLogLoadEventAsync(
-      this.metricsAdapter,
-      EntityMetricsLoadType.LOAD_MANY_RAW,
-      this.entityClassName,
-      queryContext,
-    )(
-      this.databaseAdapter.fetchManyByRawWhereClauseAsync(
-        queryContext,
-        rawWhereClause,
-        bindings,
-        querySelectionModifiers,
-      ),
-    );
-  }
-
   private async invalidateOneAsync<
     TLoadKey extends IEntityLoadKey<TFields, TIDField, TSerializedLoadValue, TLoadValue>,
     TSerializedLoadValue,
diff --git a/packages/entity/src/internal/EntityKnexDataManager.ts b/packages/entity/src/internal/EntityKnexDataManager.ts
@@ -0,0 +1,84 @@
+import {
+  EntityDatabaseAdapter,
+  FieldEqualityCondition,
+  QuerySelectionModifiers,
+  QuerySelectionModifiersWithOrderByRaw,
+} from '../EntityDatabaseAdapter';
+import { EntityQueryContext } from '../EntityQueryContext';
+import { timeAndLogLoadEventAsync } from '../metrics/EntityMetricsUtils';
+import { EntityMetricsLoadType, IEntityMetricsAdapter } from '../metrics/IEntityMetricsAdapter';
+
+/**
+ * A knex data manager is responsible for handling non-dataloader-based
+ * database operations.
+ *
+ * @internal
+ */
+export class EntityKnexDataManager<
+  TFields extends Record<string, any>,
+  TIDField extends keyof TFields,
+> {
+  constructor(
+    private readonly databaseAdapter: EntityDatabaseAdapter<TFields, TIDField>,
+    private readonly metricsAdapter: IEntityMetricsAdapter,
+    private readonly entityClassName: string,
+  ) {}
+
+  /**
+   * Loads many objects matching the conjunction of where clauses constructed from
+   * specified field equality operands.
+   *
+   * @param queryContext - query context in which to perform the load
+   * @param fieldEqualityOperands - list of field equality where clause operand specifications
+   * @param querySelectionModifiers - limit, offset, and orderBy for the query
+   * @returns array of objects matching the query
+   */
+  async loadManyByFieldEqualityConjunctionAsync<N extends keyof TFields>(
+    queryContext: EntityQueryContext,
+    fieldEqualityOperands: FieldEqualityCondition<TFields, N>[],
+    querySelectionModifiers: QuerySelectionModifiers<TFields>,
+  ): Promise<readonly Readonly<TFields>[]> {
+    return await timeAndLogLoadEventAsync(
+      this.metricsAdapter,
+      EntityMetricsLoadType.LOAD_MANY_EQUALITY_CONJUNCTION,
+      this.entityClassName,
+      queryContext,
+    )(
+      this.databaseAdapter.fetchManyByFieldEqualityConjunctionAsync(
+        queryContext,
+        fieldEqualityOperands,
+        querySelectionModifiers,
+      ),
+    );
+  }
+
+  /**
+   * Loads many objects matching the raw WHERE clause.
+   *
+   * @param queryContext - query context in which to perform the load
+   * @param rawWhereClause - parameterized SQL WHERE clause with positional binding placeholders or named binding placeholders
+   * @param bindings - array of positional bindings or object of named bindings
+   * @param querySelectionModifiers - limit, offset, orderBy, and orderByRaw for the query
+   * @returns array of objects matching the query
+   */
+  async loadManyByRawWhereClauseAsync(
+    queryContext: EntityQueryContext,
+    rawWhereClause: string,
+    bindings: any[] | object,
+    querySelectionModifiers: QuerySelectionModifiersWithOrderByRaw<TFields>,
+  ): Promise<readonly Readonly<TFields>[]> {
+    return await timeAndLogLoadEventAsync(
+      this.metricsAdapter,
+      EntityMetricsLoadType.LOAD_MANY_RAW,
+      this.entityClassName,
+      queryContext,
+    )(
+      this.databaseAdapter.fetchManyByRawWhereClauseAsync(
+        queryContext,
+        rawWhereClause,
+        bindings,
+        querySelectionModifiers,
+      ),
+    );
+  }
+}
diff --git a/packages/entity/src/internal/EntityTableDataCoordinator.ts b/packages/entity/src/internal/EntityTableDataCoordinator.ts
@@ -5,6 +5,7 @@ import { IEntityCacheAdapter } from '../IEntityCacheAdapter';
 import { IEntityCacheAdapterProvider } from '../IEntityCacheAdapterProvider';
 import { IEntityDatabaseAdapterProvider } from '../IEntityDatabaseAdapterProvider';
 import { EntityDataManager } from './EntityDataManager';
+import { EntityKnexDataManager } from './EntityKnexDataManager';
 import { ReadThroughEntityCache } from './ReadThroughEntityCache';
 import { IEntityMetricsAdapter } from '../metrics/IEntityMetricsAdapter';
 
@@ -22,6 +23,7 @@ export class EntityTableDataCoordinator<
   readonly databaseAdapter: EntityDatabaseAdapter<TFields, TIDField>;
   readonly cacheAdapter: IEntityCacheAdapter<TFields, TIDField>;
   readonly dataManager: EntityDataManager<TFields, TIDField>;
+  readonly knexDataManager: EntityKnexDataManager<TFields, TIDField>;
 
   constructor(
     readonly entityConfiguration: EntityConfiguration<TFields, TIDField>,
@@ -40,6 +42,11 @@ export class EntityTableDataCoordinator<
       metricsAdapter,
       entityClassName,
     );
+    this.knexDataManager = new EntityKnexDataManager(
+      this.databaseAdapter,
+      metricsAdapter,
+      entityClassName,
+    );
   }
 
   getQueryContextProvider(): EntityQueryContextProvider {
diff --git a/packages/entity/src/internal/__tests__/EntityDataManager-test.ts b/packages/entity/src/internal/__tests__/EntityDataManager-test.ts
diff --git a/packages/entity/src/internal/__tests__/EntityKnexDataManager-test.ts b/packages/entity/src/internal/__tests__/EntityKnexDataManager-test.ts
