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,32 @@ 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('fetchManyBySQLFragmentWithCountAsync', () => {
+    it('throws because it is unsupported', async () => {
+      const queryContext = instance(mock(EntityQueryContext));
+      const databaseAdapter = new StubPostgresDatabaseAdapter<TestFields, 'customIdField'>(
+        testEntityConfiguration,
+        new Map(),
+      );
+      await expect(
+        databaseAdapter.fetchManyBySQLFragmentWithCountAsync(queryContext, sql``, {}),
+      ).rejects.toThrow();
+    });
+  });
+
   describe('insertAsync', () => {
     it('inserts a record', async () => {
       const queryContext = instance(mock(EntityQueryContext));

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

@@ -92,6 +92,17 @@ interface EntityLoaderBasePaginationArgs<
    * Order the entities by specified columns and orders. If ID field is not included in the orderBy, it will be automatically included as the last orderBy field to ensure stable pagination.
    */
   orderBy?: EntityLoaderOrderByClause<TFields, TSelectedFields>[];
+
+  /**
+   * Whether to calculate and include the Connection totalCount field containing the total count of entities matching the where SQLFragment.
+   *
+   * Note that this may be an expensive operation, especially for large datasets, as it may require an additional SQL query with a COUNT(*) aggregation.
+   * It is recommended to only set this to true when necessary for the client application, such as when implementing pagination UIs that need to know the total number of pages.
+   *
+   * When true and cursor is non-null, the total count is calculated by running a separate COUNT(*) query with the same WHERE clause.
+   * When true and cursor is null, the total count is calculated by running the main query with a window function to count all matching rows, which is slightly faster but still expensive.
+   */
+  includeTotal?: boolean;
 }
 
 /**
@@ -304,6 +315,7 @@ export class AuthorizationResultBasedKnexEntityLoader<
     return {
       edges,
       pageInfo,
+      ...(pageResult.totalCount !== undefined && { totalCount: pageResult.totalCount }),
     };
   }
 }

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

@@ -268,6 +268,68 @@ export abstract class BasePostgresEntityDatabaseAdapter<
     querySelectionModifiers: TableQuerySelectionModifiersWithOrderByFragment,
   ): Promise<object[]>;
 
+  /**
+   * Fetch total count of objects matching the SQL fragment.
+   *
+   * Note that this may be an expensive operation, especially for large datasets, as it does a full query with a COUNT(*) aggregation.
+   *
+   * @param queryContext - query context with which to perform the fetch
+   * @param sqlFragment - SQLFragment for the WHERE clause of the query
+   * @returns total count of objects matching the SQL fragment
+   */
+  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

@@ -212,16 +212,107 @@ export class PostgresEntityDatabaseAdapter<
     sqlFragment: SQLFragment,
     querySelectionModifiers: TableQuerySelectionModifiersWithOrderByFragment,
   ): Promise<object[]> {
-    let query = queryInterface
-      .select()
+    const { results } = await this.fetchManyBySQLFragmentWithOptionalCountInternalAsync(
+      queryInterface,
+      tableName,
+      sqlFragment,
+      querySelectionModifiers,
+      false,
+    );
+    return results;
+  }
+
+  protected async fetchCountBySQLFragmentInternalAsync(
+    queryInterface: Knex,
+    tableName: string,
+    sqlFragment: SQLFragment,
+  ): Promise<number> {
+    const query = queryInterface
+      .count('* as count')
       .from(tableName)
-      .whereRaw(sqlFragment.sql, sqlFragment.getKnexBindings());
+      .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 }> {
+    return await this.fetchManyBySQLFragmentWithOptionalCountInternalAsync(
+      queryInterface,
+      tableName,
+      sqlFragment,
+      querySelectionModifiers,
+      true,
+    );
+  }
+
+  private async fetchManyBySQLFragmentWithOptionalCountInternalAsync(
+    queryInterface: Knex,
+    tableName: string,
+    sqlFragment: SQLFragment,
+    querySelectionModifiers: TableQuerySelectionModifiersWithOrderByFragment,
+    includeTotalCountViaWindowFunction: true,
+  ): Promise<{ results: object[]; totalCount: number }>;
+  private async fetchManyBySQLFragmentWithOptionalCountInternalAsync(
+    queryInterface: Knex,
+    tableName: string,
+    sqlFragment: SQLFragment,
+    querySelectionModifiers: TableQuerySelectionModifiersWithOrderByFragment,
+    includeTotalCountViaWindowFunction: false,
+  ): Promise<{ results: object[]; totalCount: undefined }>;
+  private async fetchManyBySQLFragmentWithOptionalCountInternalAsync(
+    queryInterface: Knex,
+    tableName: string,
+    sqlFragment: SQLFragment,
+    querySelectionModifiers: TableQuerySelectionModifiersWithOrderByFragment,
+    includeTotalCountViaWindowFunction: boolean,
+  ): Promise<{ results: object[]; totalCount: number | undefined }> {
+    // Build the base query with window function for total count
+    let query = queryInterface.select('*');
+    if (includeTotalCountViaWindowFunction) {
+      query = query.select(queryInterface.raw('COUNT(*) OVER() as __total_count'));
+    }
+    query = query.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());
     }
-    return await wrapNativePostgresCallAsync(() => query);
+
+    const rows = await wrapNativePostgresCallAsync(() => query);
+
+    if (!includeTotalCountViaWindowFunction) {
+      return { results: rows, totalCount: undefined };
+    }
+
+    // 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(