feat: add includeTotalCount to pagination

Author: wschurman

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

@@ -298,4 +298,21 @@ 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');
+  }
+
+  protected async fetchManyBySQLFragmentWithCountInternalAsync(
+    _queryInterface: any,
+    _tableName: string,
+    _sqlFragment: SQLFragment,
+    _querySelectionModifiers: TableQuerySelectionModifiersWithOrderByFragment,
+  ): Promise<{ results: object[]; totalCount: 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,59 @@ 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>;
+
+  /**
+   * Fetch many objects matching the SQL fragment along with the total count using a window function.
+   * This is more efficient than running separate fetch and count queries when both are needed.
+   *
+   * @param queryContext - query context with which to perform the fetch
+   * @param sqlFragment - SQLFragment for the WHERE clause of the query
+   * @param querySelectionModifiers - limit, offset, and orderByFragment for the query
+   * @returns Object containing both the results array and the total count
+   */
+  async fetchManyBySQLFragmentWithCountAsync(
+    queryContext: EntityQueryContext,
+    sqlFragment: SQLFragment,
+    querySelectionModifiers: PostgresQuerySelectionModifiersWithOrderByFragment<TFields>,
+  ): Promise<{ results: readonly Readonly<TFields>[]; totalCount: number }> {
+    const { results, totalCount } = await this.fetchManyBySQLFragmentWithCountInternalAsync(
+      queryContext.getQueryInterface(),
+      this.entityConfiguration.tableName,
+      sqlFragment,
+      this.convertToTableQueryModifiersWithOrderByFragment(querySelectionModifiers),
+    );
+
+    return {
+      results: results.map((result) =>
+        transformDatabaseObjectToFields(this.entityConfiguration, this.fieldTransformerMap, result),
+      ),
+      totalCount,
+    };
+  }
+
+  protected abstract fetchManyBySQLFragmentWithCountInternalAsync(
+    queryInterface: Knex,
+    tableName: string,
+    sqlFragment: SQLFragment,
+    querySelectionModifiers: TableQuerySelectionModifiersWithOrderByFragment,
+  ): Promise<{ results: object[]; totalCount: 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,65 @@ 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);
+  }
+
+  /**
+   * Fetches paginated results with total count using a window function.
+   * This executes a single query that returns both the result set and the total count,
+   * avoiding the need for two separate queries.
+   *
+   * The query uses COUNT(*) OVER() to calculate the total count as a window function,
+   * which runs over the entire result set before LIMIT/OFFSET are applied.
+   */
+  protected async fetchManyBySQLFragmentWithCountInternalAsync(
+    queryInterface: Knex,
+    tableName: string,
+    sqlFragment: SQLFragment,
+    querySelectionModifiers: TableQuerySelectionModifiersWithOrderByFragment,
+  ): Promise<{ results: object[]; totalCount: number }> {
+    // Build the base query with window function for total count
+    let query = queryInterface
+      .select('*')
+      .select(queryInterface.raw('COUNT(*) OVER() as __total_count'))
+      .from(tableName)
+      .whereRaw(sqlFragment.sql, sqlFragment.getKnexBindings());
+
+    // Apply order by modifiers
+    query = this.applyQueryModifiersToQuery(query, querySelectionModifiers);
+    const { orderByFragment } = querySelectionModifiers;
+    if (orderByFragment !== undefined) {
+      query = query.orderByRaw(orderByFragment.sql, orderByFragment.getKnexBindings());
+    }
+
+    const rows = await wrapNativePostgresCallAsync(() => query);
+
+    // If no rows, return empty results with 0 count
+    if (rows.length === 0) {
+      return { results: [], totalCount: 0 };
+    }
+
+    // Extract total count from the first row and remove it from all rows
+    const totalCount = parseInt(rows[0].__total_count ?? '0', 10);
+    const results = rows.map((row) => {
+      const { __total_count, ...rest } = row;
+      return rest;
+    });
+
+    return { results, totalCount };
+  }
+
   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,23 @@ class TestEntityDatabaseAdapter extends BasePostgresEntityDatabaseAdapter<
   ): Promise<number> {
     return this.deleteCount;
   }
+
+  protected async fetchCountBySQLFragmentInternalAsync(
+    _queryInterface: any,
+    _tableName: string,
+    _sqlFragment: any,
+  ): Promise<number> {
+    return 0;
+  }
+
+  protected async fetchManyBySQLFragmentWithCountInternalAsync(
+    _queryInterface: any,
+    _tableName: string,
+    _sqlFragment: any,
+    _querySelectionModifiers: any,
+  ): Promise<{ results: object[]; totalCount: number }> {
+    return { results: this.fetchSQLFragmentResults, totalCount: 0 };
+  }
 }
 
 describe(BasePostgresEntityDatabaseAdapter, () => {

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

@@ -299,4 +299,21 @@ 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');
+  }
+
+  protected async fetchManyBySQLFragmentWithCountInternalAsync(
+    _queryInterface: any,
+    _tableName: string,
+    _sqlFragment: SQLFragment,
+    _querySelectionModifiers: TableQuerySelectionModifiersWithOrderByFragment,
+  ): Promise<{ results: object[]; totalCount: number }> {
+    throw new Error('SQL fragments not supported for StubDatabaseAdapter');
+  }
 }