feat: add includeTotalCount to pagination

Author: wschurman

packages/entity-database-adapter-knex-testing-utils/src/StubPostgresDatabaseAdapter.ts

@@ -298,4 +298,12 @@ export class StubPostgresDatabaseAdapter<
     objectCollection.splice(objectIndex, 1);
     return 1;
   }
+
+  protected async fetchCountBySQLFragmentInternalAsync(
+    _queryInterface: any,
+    _tableName: string,
+    _sqlFragment: any,
+  ): Promise<number> {
+    throw new Error('SQL fragments not supported for StubDatabaseAdapter');
+  }
 }

packages/entity-database-adapter-knex-testing-utils/src/__tests__/StubPostgresDatabaseAdapter-test.ts

@@ -520,6 +520,19 @@ describe(StubPostgresDatabaseAdapter, () => {
     });
   });
 
+  describe('fetchCountBySQLFragmentAsync', () => {
+    it('throws because it is unsupported', async () => {
+      const queryContext = instance(mock(EntityQueryContext));
+      const databaseAdapter = new StubPostgresDatabaseAdapter<TestFields, 'customIdField'>(
+        testEntityConfiguration,
+        new Map(),
+      );
+      await expect(
+        databaseAdapter.fetchCountBySQLFragmentAsync(queryContext, sql``),
+      ).rejects.toThrow();
+    });
+  });
+
   describe('insertAsync', () => {
     it('inserts a record', async () => {
       const queryContext = instance(mock(EntityQueryContext));

packages/entity-database-adapter-knex/src/AuthorizationResultBasedKnexEntityLoader.ts

@@ -97,6 +97,13 @@ interface EntityLoaderBasePaginationArgs<
    * Cursor fields to use for pagination. If not provided, the order by fields will be used as the cursor fields.
    */
   cursorFields?: readonly (keyof TFields)[];
+
+  /**
+   * Whether to include the total count of matching entities in the result.
+   * When true, uses a window function to count all matching rows, which may impact performance on large datasets.
+   * The total count is available in the connection's totalCount field.
+   */
+  includeTotal?: boolean;
 }
 
 /**
@@ -309,6 +316,7 @@ export class AuthorizationResultBasedKnexEntityLoader<
     return {
       edges,
       pageInfo,
+      ...(pageResult.totalCount !== undefined && { totalCount: pageResult.totalCount }),
     };
   }
 }

packages/entity-database-adapter-knex/src/BasePostgresEntityDatabaseAdapter.ts

@@ -268,6 +268,23 @@ export abstract class BasePostgresEntityDatabaseAdapter<
     querySelectionModifiers: TableQuerySelectionModifiersWithOrderByFragment,
   ): Promise<object[]>;
 
+  async fetchCountBySQLFragmentAsync(
+    queryContext: EntityQueryContext,
+    sqlFragment: SQLFragment,
+  ): Promise<number> {
+    return await this.fetchCountBySQLFragmentInternalAsync(
+      queryContext.getQueryInterface(),
+      this.entityConfiguration.tableName,
+      sqlFragment,
+    );
+  }
+
+  protected abstract fetchCountBySQLFragmentInternalAsync(
+    queryInterface: Knex,
+    tableName: string,
+    sqlFragment: SQLFragment,
+  ): Promise<number>;
+
   private convertToTableQueryModifiersWithOrderByRaw(
     querySelectionModifiers: PostgresQuerySelectionModifiersWithOrderByRaw<TFields>,
   ): TableQuerySelectionModifiersWithOrderByRaw {

packages/entity-database-adapter-knex/src/EnforcingKnexEntityLoader.ts

@@ -238,6 +238,7 @@ export class EnforcingKnexEntityLoader<
     return {
       edges,
       pageInfo: pageResult.pageInfo,
+      ...(pageResult.totalCount !== undefined && { totalCount: pageResult.totalCount }),
     };
   }
 }

packages/entity-database-adapter-knex/src/PostgresEntityDatabaseAdapter.ts

@@ -224,6 +224,20 @@ export class PostgresEntityDatabaseAdapter<
     return await wrapNativePostgresCallAsync(() => query);
   }
 
+  protected async fetchCountBySQLFragmentInternalAsync(
+    queryInterface: Knex,
+    tableName: string,
+    sqlFragment: SQLFragment,
+  ): Promise<number> {
+    const query = queryInterface
+      .count('* as count')
+      .from(tableName)
+      .whereRaw(sqlFragment.sql, sqlFragment.getKnexBindings())
+      .first();
+    const result = await wrapNativePostgresCallAsync(() => query);
+    return parseInt(result?.count ?? '0', 10);
+  }
+
   protected async insertInternalAsync(
     queryInterface: Knex,
     tableName: string,

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

@@ -1774,5 +1774,93 @@ describe('postgres entity integration', () => {
       expect(page2.edges).toHaveLength(1);
       expect(page2.pageInfo.hasNextPage).toBe(false);
     });
+
+    it('includes total count when includeTotal is true', async () => {
+      const vc = new ViewerContext(createKnexIntegrationTestEntityCompanionProvider(knexInstance));
+
+      await PostgresTestEntity.dropPostgresTableAsync(knexInstance);
+      await PostgresTestEntity.createOrTruncatePostgresTableAsync(knexInstance);
+
+      // Create test data
+      const names = ['Alice', 'Bob', 'Charlie', 'David', 'Eve', 'Frank', 'Grace', 'Henry'];
+      for (let i = 0; i < names.length; i++) {
+        await PostgresTestEntity.creator(vc)
+          .setField('name', names[i]!)
+          .setField('hasACat', i % 2 === 0)
+          .setField('hasADog', i % 3 === 0)
+          .createAsync();
+      }
+
+      // Test includeTotal with no filter
+      const pageWithTotal = await PostgresTestEntity.knexLoader(vc).loadPageBySQLAsync({
+        first: 3,
+        orderBy: [{ fieldName: 'name', order: OrderByOrdering.ASCENDING }],
+        includeTotal: true,
+      });
+
+      expect(pageWithTotal.edges).toHaveLength(3);
+      expect(pageWithTotal.totalCount).toBe(8);
+      expect(pageWithTotal.edges[0]?.node.getField('name')).toBe('Alice');
+      expect(pageWithTotal.edges[1]?.node.getField('name')).toBe('Bob');
+      expect(pageWithTotal.edges[2]?.node.getField('name')).toBe('Charlie');
+
+      // Test includeTotal with where clause
+      const pageWithFilter = await PostgresTestEntity.knexLoader(vc).loadPageBySQLAsync({
+        first: 2,
+        where: sql`has_a_cat = ${true}`,
+        orderBy: [{ fieldName: 'name', order: OrderByOrdering.ASCENDING }],
+        includeTotal: true,
+      });
+
+      expect(pageWithFilter.edges).toHaveLength(2);
+      expect(pageWithFilter.totalCount).toBe(4); // Alice, Charlie, Eve, Grace have cats
+      expect(pageWithFilter.edges[0]?.node.getField('name')).toBe('Alice');
+      expect(pageWithFilter.edges[1]?.node.getField('name')).toBe('Charlie');
+
+      // Test that totalCount is not affected by pagination cursor
+      const nextPageWithTotal = await PostgresTestEntity.knexLoader(vc).loadPageBySQLAsync({
+        first: 2,
+        after: pageWithFilter.pageInfo.endCursor!,
+        where: sql`has_a_cat = ${true}`,
+        orderBy: [{ fieldName: 'name', order: OrderByOrdering.ASCENDING }],
+        includeTotal: true,
+      });
+
+      expect(nextPageWithTotal.edges).toHaveLength(2);
+      expect(nextPageWithTotal.totalCount).toBe(4); // Total count remains the same
+      expect(nextPageWithTotal.edges[0]?.node.getField('name')).toBe('Eve');
+      expect(nextPageWithTotal.edges[1]?.node.getField('name')).toBe('Grace');
+
+      // Test backward pagination with includeTotal
+      const backwardPageWithTotal = await PostgresTestEntity.knexLoader(vc).loadPageBySQLAsync({
+        last: 2,
+        orderBy: [{ fieldName: 'name', order: OrderByOrdering.ASCENDING }],
+        includeTotal: true,
+      });
+
+      expect(backwardPageWithTotal.edges).toHaveLength(2);
+      expect(backwardPageWithTotal.totalCount).toBe(8);
+      expect(backwardPageWithTotal.edges[0]?.node.getField('name')).toBe('Grace');
+      expect(backwardPageWithTotal.edges[1]?.node.getField('name')).toBe('Henry');
+
+      // Test that totalCount is undefined when includeTotal is false
+      const pageWithoutTotal = await PostgresTestEntity.knexLoader(vc).loadPageBySQLAsync({
+        first: 3,
+        orderBy: [{ fieldName: 'name', order: OrderByOrdering.ASCENDING }],
+        includeTotal: false,
+      });
+
+      expect(pageWithoutTotal.edges).toHaveLength(3);
+      expect(pageWithoutTotal.totalCount).toBeUndefined();
+
+      // Test that totalCount is undefined when includeTotal is not specified
+      const pageDefaultNoTotal = await PostgresTestEntity.knexLoader(vc).loadPageBySQLAsync({
+        first: 3,
+        orderBy: [{ fieldName: 'name', order: OrderByOrdering.ASCENDING }],
+      });
+
+      expect(pageDefaultNoTotal.edges).toHaveLength(3);
+      expect(pageDefaultNoTotal.totalCount).toBeUndefined();
+    });
   });
 });

packages/entity-database-adapter-knex/src/__tests__/BasePostgresEntityDatabaseAdapter-test.ts

@@ -126,6 +126,14 @@ class TestEntityDatabaseAdapter extends BasePostgresEntityDatabaseAdapter<
   ): Promise<number> {
     return this.deleteCount;
   }
+
+  protected async fetchCountBySQLFragmentInternalAsync(
+    _queryInterface: any,
+    _tableName: string,
+    _sqlFragment: any,
+  ): Promise<number> {
+    return 0;
+  }
 }
 
 describe(BasePostgresEntityDatabaseAdapter, () => {

packages/entity-database-adapter-knex/src/__tests__/fixtures/StubPostgresDatabaseAdapter.ts

@@ -299,4 +299,12 @@ export class StubPostgresDatabaseAdapter<
     objectCollection.splice(objectIndex, 1);
     return 1;
   }
+
+  protected async fetchCountBySQLFragmentInternalAsync(
+    _queryInterface: any,
+    _tableName: string,
+    _sqlFragment: any,
+  ): Promise<number> {
+    throw new Error('SQL fragments not supported for StubDatabaseAdapter');
+  }
 }

packages/entity-database-adapter-knex/src/internal/EntityKnexDataManager.ts

@@ -27,6 +27,7 @@ interface BasePaginationArgs<TFields extends Record<string, any>> {
   where?: SQLFragment;
   orderBy?: PostgresOrderByClause<TFields>[];
   cursorFields?: readonly (keyof TFields)[];
+  includeTotal?: boolean;
 }
 
 /**
@@ -80,6 +81,11 @@ export interface PageInfo {
 export interface Connection<TNode> {
   edges: Edge<TNode>[];
   pageInfo: PageInfo;
+  /**
+   * Total count of all matching entities (not just those in the current page).
+   * Only present when includeTotal is true in the request.
+   */
+  totalCount?: number;
 }
 
 /**
@@ -204,7 +210,7 @@ export class EntityKnexDataManager<
     }
 
     const isForward = 'first' in args;
-    const { where, orderBy, cursorFields } = args;
+    const { where, orderBy, cursorFields, includeTotal } = args;
 
     let limit: number;
     let cursor: string | undefined;
@@ -248,7 +254,8 @@ export class EntityKnexDataManager<
         }));
 
     // Fetch data with limit + 1 to check for more pages
-    const fieldObjects = await timeAndLogLoadEventAsync(
+    // When includeTotal is true, also run a count query in parallel
+    const fetchPromise = timeAndLogLoadEventAsync(
       this.metricsAdapter,
       EntityMetricsLoadType.LOAD_PAGE,
       this.entityClassName,
@@ -260,6 +267,13 @@ export class EntityKnexDataManager<
       }),
     );
 
+    const [fieldObjects, totalCount] = await Promise.all([
+      fetchPromise,
+      includeTotal
+        ? this.databaseAdapter.fetchCountBySQLFragmentAsync(queryContext, where ?? sql`1 = 1`)
+        : Promise.resolve(undefined),
+    ]);
+
     // Process results
     const hasMore = fieldObjects.length > limit;
     const pageFieldObjects = hasMore ? fieldObjects.slice(0, limit) : [...fieldObjects];
@@ -284,6 +298,7 @@ export class EntityKnexDataManager<
     return {
       edges,
       pageInfo,
+      ...(totalCount !== undefined && { totalCount }),
     };
   }