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

@@ -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

@@ -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(