Skip to content

Commit 9458e67

Browse files
wschurmanwschurman
committed
feat: add includeTotalCount to pagination
1 parent 0b4f87d commit 9458e67

10 files changed

Lines changed: 284 additions & 2 deletions

File tree

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -298,4 +298,13 @@ export class StubPostgresDatabaseAdapter<
298298
objectCollection.splice(objectIndex, 1);
299299
return 1;
300300
}
301+
302+
protected async fetchCountBySQLFragmentInternalAsync(
303+
_queryInterface: any,
304+
tableName: string,
305+
_sqlFragment: any,
306+
): Promise<number> {
307+
const objectCollection = this.getObjectCollectionForTable(tableName);
308+
return objectCollection.length;
309+
}
301310
}

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,13 @@ interface EntityLoaderBasePaginationArgs<
9797
* Cursor fields to use for pagination. If not provided, the order by fields will be used as the cursor fields.
9898
*/
9999
cursorFields?: readonly (keyof TFields)[];
100+
101+
/**
102+
* Whether to include the total count of matching entities in the result.
103+
* When true, uses a window function to count all matching rows, which may impact performance on large datasets.
104+
* The total count is available in the connection's totalCount field.
105+
*/
106+
includeTotal?: boolean;
100107
}
101108

102109
/**
@@ -309,6 +316,7 @@ export class AuthorizationResultBasedKnexEntityLoader<
309316
return {
310317
edges,
311318
pageInfo,
319+
...(pageResult.totalCount !== undefined && { totalCount: pageResult.totalCount }),
312320
};
313321
}
314322
}

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -268,6 +268,23 @@ export abstract class BasePostgresEntityDatabaseAdapter<
268268
querySelectionModifiers: TableQuerySelectionModifiersWithOrderByFragment,
269269
): Promise<object[]>;
270270

271+
async fetchCountBySQLFragmentAsync(
272+
queryContext: EntityQueryContext,
273+
sqlFragment: SQLFragment,
274+
): Promise<number> {
275+
return await this.fetchCountBySQLFragmentInternalAsync(
276+
queryContext.getQueryInterface(),
277+
this.entityConfiguration.tableName,
278+
sqlFragment,
279+
);
280+
}
281+
282+
protected abstract fetchCountBySQLFragmentInternalAsync(
283+
queryInterface: Knex,
284+
tableName: string,
285+
sqlFragment: SQLFragment,
286+
): Promise<number>;
287+
271288
private convertToTableQueryModifiersWithOrderByRaw(
272289
querySelectionModifiers: PostgresQuerySelectionModifiersWithOrderByRaw<TFields>,
273290
): TableQuerySelectionModifiersWithOrderByRaw {

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -238,6 +238,7 @@ export class EnforcingKnexEntityLoader<
238238
return {
239239
edges,
240240
pageInfo: pageResult.pageInfo,
241+
...(pageResult.totalCount !== undefined && { totalCount: pageResult.totalCount }),
241242
};
242243
}
243244
}

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -224,6 +224,20 @@ export class PostgresEntityDatabaseAdapter<
224224
return await wrapNativePostgresCallAsync(() => query);
225225
}
226226

227+
protected async fetchCountBySQLFragmentInternalAsync(
228+
queryInterface: Knex,
229+
tableName: string,
230+
sqlFragment: SQLFragment,
231+
): Promise<number> {
232+
const query = queryInterface
233+
.count('* as count')
234+
.from(tableName)
235+
.whereRaw(sqlFragment.sql, sqlFragment.getKnexBindings())
236+
.first();
237+
const result = await wrapNativePostgresCallAsync(() => query);
238+
return parseInt(result?.count ?? '0', 10);
239+
}
240+
227241
protected async insertInternalAsync(
228242
queryInterface: Knex,
229243
tableName: string,

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

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1774,5 +1774,93 @@ describe('postgres entity integration', () => {
17741774
expect(page2.edges).toHaveLength(1);
17751775
expect(page2.pageInfo.hasNextPage).toBe(false);
17761776
});
1777+
1778+
it('includes total count when includeTotal is true', async () => {
1779+
const vc = new ViewerContext(createKnexIntegrationTestEntityCompanionProvider(knexInstance));
1780+
1781+
await PostgresTestEntity.dropPostgresTableAsync(knexInstance);
1782+
await PostgresTestEntity.createOrTruncatePostgresTableAsync(knexInstance);
1783+
1784+
// Create test data
1785+
const names = ['Alice', 'Bob', 'Charlie', 'David', 'Eve', 'Frank', 'Grace', 'Henry'];
1786+
for (let i = 0; i < names.length; i++) {
1787+
await PostgresTestEntity.creator(vc)
1788+
.setField('name', names[i]!)
1789+
.setField('hasACat', i % 2 === 0)
1790+
.setField('hasADog', i % 3 === 0)
1791+
.createAsync();
1792+
}
1793+
1794+
// Test includeTotal with no filter
1795+
const pageWithTotal = await PostgresTestEntity.knexLoader(vc).loadPageBySQLAsync({
1796+
first: 3,
1797+
orderBy: [{ fieldName: 'name', order: OrderByOrdering.ASCENDING }],
1798+
includeTotal: true,
1799+
});
1800+
1801+
expect(pageWithTotal.edges).toHaveLength(3);
1802+
expect(pageWithTotal.totalCount).toBe(8);
1803+
expect(pageWithTotal.edges[0]?.node.getField('name')).toBe('Alice');
1804+
expect(pageWithTotal.edges[1]?.node.getField('name')).toBe('Bob');
1805+
expect(pageWithTotal.edges[2]?.node.getField('name')).toBe('Charlie');
1806+
1807+
// Test includeTotal with where clause
1808+
const pageWithFilter = await PostgresTestEntity.knexLoader(vc).loadPageBySQLAsync({
1809+
first: 2,
1810+
where: sql`has_a_cat = ${true}`,
1811+
orderBy: [{ fieldName: 'name', order: OrderByOrdering.ASCENDING }],
1812+
includeTotal: true,
1813+
});
1814+
1815+
expect(pageWithFilter.edges).toHaveLength(2);
1816+
expect(pageWithFilter.totalCount).toBe(4); // Alice, Charlie, Eve, Grace have cats
1817+
expect(pageWithFilter.edges[0]?.node.getField('name')).toBe('Alice');
1818+
expect(pageWithFilter.edges[1]?.node.getField('name')).toBe('Charlie');
1819+
1820+
// Test that totalCount is not affected by pagination cursor
1821+
const nextPageWithTotal = await PostgresTestEntity.knexLoader(vc).loadPageBySQLAsync({
1822+
first: 2,
1823+
after: pageWithFilter.pageInfo.endCursor!,
1824+
where: sql`has_a_cat = ${true}`,
1825+
orderBy: [{ fieldName: 'name', order: OrderByOrdering.ASCENDING }],
1826+
includeTotal: true,
1827+
});
1828+
1829+
expect(nextPageWithTotal.edges).toHaveLength(2);
1830+
expect(nextPageWithTotal.totalCount).toBe(4); // Total count remains the same
1831+
expect(nextPageWithTotal.edges[0]?.node.getField('name')).toBe('Eve');
1832+
expect(nextPageWithTotal.edges[1]?.node.getField('name')).toBe('Grace');
1833+
1834+
// Test backward pagination with includeTotal
1835+
const backwardPageWithTotal = await PostgresTestEntity.knexLoader(vc).loadPageBySQLAsync({
1836+
last: 2,
1837+
orderBy: [{ fieldName: 'name', order: OrderByOrdering.ASCENDING }],
1838+
includeTotal: true,
1839+
});
1840+
1841+
expect(backwardPageWithTotal.edges).toHaveLength(2);
1842+
expect(backwardPageWithTotal.totalCount).toBe(8);
1843+
expect(backwardPageWithTotal.edges[0]?.node.getField('name')).toBe('Grace');
1844+
expect(backwardPageWithTotal.edges[1]?.node.getField('name')).toBe('Henry');
1845+
1846+
// Test that totalCount is undefined when includeTotal is false
1847+
const pageWithoutTotal = await PostgresTestEntity.knexLoader(vc).loadPageBySQLAsync({
1848+
first: 3,
1849+
orderBy: [{ fieldName: 'name', order: OrderByOrdering.ASCENDING }],
1850+
includeTotal: false,
1851+
});
1852+
1853+
expect(pageWithoutTotal.edges).toHaveLength(3);
1854+
expect(pageWithoutTotal.totalCount).toBeUndefined();
1855+
1856+
// Test that totalCount is undefined when includeTotal is not specified
1857+
const pageDefaultNoTotal = await PostgresTestEntity.knexLoader(vc).loadPageBySQLAsync({
1858+
first: 3,
1859+
orderBy: [{ fieldName: 'name', order: OrderByOrdering.ASCENDING }],
1860+
});
1861+
1862+
expect(pageDefaultNoTotal.edges).toHaveLength(3);
1863+
expect(pageDefaultNoTotal.totalCount).toBeUndefined();
1864+
});
17771865
});
17781866
});

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,14 @@ class TestEntityDatabaseAdapter extends BasePostgresEntityDatabaseAdapter<
126126
): Promise<number> {
127127
return this.deleteCount;
128128
}
129+
130+
protected async fetchCountBySQLFragmentInternalAsync(
131+
_queryInterface: any,
132+
_tableName: string,
133+
_sqlFragment: any,
134+
): Promise<number> {
135+
return 0;
136+
}
129137
}
130138

131139
describe(BasePostgresEntityDatabaseAdapter, () => {

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -299,4 +299,13 @@ export class StubPostgresDatabaseAdapter<
299299
objectCollection.splice(objectIndex, 1);
300300
return 1;
301301
}
302+
303+
protected async fetchCountBySQLFragmentInternalAsync(
304+
_queryInterface: any,
305+
tableName: string,
306+
_sqlFragment: any,
307+
): Promise<number> {
308+
const objectCollection = this.getObjectCollectionForTable(tableName);
309+
return objectCollection.length;
310+
}
302311
}

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

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ interface BasePaginationArgs<TFields extends Record<string, any>> {
2727
where?: SQLFragment;
2828
orderBy?: PostgresOrderByClause<TFields>[];
2929
cursorFields?: readonly (keyof TFields)[];
30+
includeTotal?: boolean;
3031
}
3132

3233
/**
@@ -80,6 +81,11 @@ export interface PageInfo {
8081
export interface Connection<TNode> {
8182
edges: Edge<TNode>[];
8283
pageInfo: PageInfo;
84+
/**
85+
* Total count of all matching entities (not just those in the current page).
86+
* Only present when includeTotal is true in the request.
87+
*/
88+
totalCount?: number;
8389
}
8490

8591
/**
@@ -204,7 +210,7 @@ export class EntityKnexDataManager<
204210
}
205211

206212
const isForward = 'first' in args;
207-
const { where, orderBy, cursorFields } = args;
213+
const { where, orderBy, cursorFields, includeTotal } = args;
208214

209215
let limit: number;
210216
let cursor: string | undefined;
@@ -248,7 +254,8 @@ export class EntityKnexDataManager<
248254
}));
249255

250256
// Fetch data with limit + 1 to check for more pages
251-
const fieldObjects = await timeAndLogLoadEventAsync(
257+
// When includeTotal is true, also run a count query in parallel
258+
const fetchPromise = timeAndLogLoadEventAsync(
252259
this.metricsAdapter,
253260
EntityMetricsLoadType.LOAD_PAGE,
254261
this.entityClassName,
@@ -260,6 +267,13 @@ export class EntityKnexDataManager<
260267
}),
261268
);
262269

270+
const [fieldObjects, totalCount] = await Promise.all([
271+
fetchPromise,
272+
includeTotal
273+
? this.databaseAdapter.fetchCountBySQLFragmentAsync(queryContext, where ?? sql`1 = 1`)
274+
: Promise.resolve(undefined),
275+
]);
276+
263277
// Process results
264278
const hasMore = fieldObjects.length > limit;
265279
const pageFieldObjects = hasMore ? fieldObjects.slice(0, limit) : [...fieldObjects];
@@ -284,6 +298,7 @@ export class EntityKnexDataManager<
284298
return {
285299
edges,
286300
pageInfo,
301+
...(totalCount !== undefined && { totalCount }),
287302
};
288303
}
289304

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

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -253,4 +253,117 @@ describe(EntityKnexDataManager, () => {
253253
});
254254
});
255255
});
256+
257+
describe('loadPageBySQLFragmentAsync', () => {
258+
describe('includeTotal functionality', () => {
259+
it('includes totalCount when includeTotal is true', async () => {
260+
const queryContext = instance(mock<EntityQueryContext>());
261+
const databaseAdapterMock = mock<
262+
PostgresEntityDatabaseAdapter<TestFields, 'customIdField'>
263+
>(PostgresEntityDatabaseAdapter);
264+
265+
when(
266+
databaseAdapterMock.fetchManyBySQLFragmentAsync(anything(), anything(), anything()),
267+
).thenResolve([
268+
{
269+
customIdField: '1',
270+
testIndexedField: 'unique1',
271+
stringField: 'hello',
272+
intField: 1,
273+
dateField: new Date(),
274+
nullableField: null,
275+
},
276+
{
277+
customIdField: '2',
278+
testIndexedField: 'unique2',
279+
stringField: 'world',
280+
intField: 2,
281+
dateField: new Date(),
282+
nullableField: null,
283+
},
284+
]);
285+
286+
when(databaseAdapterMock.fetchCountBySQLFragmentAsync(anything(), anything())).thenResolve(
287+
42,
288+
);
289+
290+
const entityDataManager = new EntityKnexDataManager(
291+
testEntityConfiguration,
292+
instance(databaseAdapterMock),
293+
new NoOpEntityMetricsAdapter(),
294+
TestEntity.name,
295+
);
296+
297+
const result = await entityDataManager.loadPageBySQLFragmentAsync(queryContext, {
298+
first: 10,
299+
includeTotal: true,
300+
});
301+
302+
expect(result.edges).toHaveLength(2);
303+
expect(result.totalCount).toBe(42);
304+
verify(databaseAdapterMock.fetchCountBySQLFragmentAsync(anything(), anything())).once();
305+
});
306+
307+
it('does not include totalCount when includeTotal is false', async () => {
308+
const queryContext = instance(mock<EntityQueryContext>());
309+
const databaseAdapterMock = mock<
310+
PostgresEntityDatabaseAdapter<TestFields, 'customIdField'>
311+
>(PostgresEntityDatabaseAdapter);
312+
313+
when(
314+
databaseAdapterMock.fetchManyBySQLFragmentAsync(anything(), anything(), anything()),
315+
).thenResolve([
316+
{
317+
customIdField: '1',
318+
testIndexedField: 'unique1',
319+
stringField: 'hello',
320+
intField: 1,
321+
dateField: new Date(),
322+
nullableField: null,
323+
},
324+
]);
325+
326+
const entityDataManager = new EntityKnexDataManager(
327+
testEntityConfiguration,
328+
instance(databaseAdapterMock),
329+
new NoOpEntityMetricsAdapter(),
330+
TestEntity.name,
331+
);
332+
333+
const result = await entityDataManager.loadPageBySQLFragmentAsync(queryContext, {
334+
first: 10,
335+
includeTotal: false,
336+
});
337+
338+
expect(result.edges).toHaveLength(1);
339+
expect(result.totalCount).toBeUndefined();
340+
verify(databaseAdapterMock.fetchCountBySQLFragmentAsync(anything(), anything())).never();
341+
});
342+
343+
it('does not include totalCount when includeTotal is not specified', async () => {
344+
const queryContext = instance(mock<EntityQueryContext>());
345+
const databaseAdapterMock = mock<
346+
PostgresEntityDatabaseAdapter<TestFields, 'customIdField'>
347+
>(PostgresEntityDatabaseAdapter);
348+
349+
when(
350+
databaseAdapterMock.fetchManyBySQLFragmentAsync(anything(), anything(), anything()),
351+
).thenResolve([]);
352+
353+
const entityDataManager = new EntityKnexDataManager(
354+
testEntityConfiguration,
355+
instance(databaseAdapterMock),
356+
new NoOpEntityMetricsAdapter(),
357+
TestEntity.name,
358+
);
359+
360+
const result = await entityDataManager.loadPageBySQLFragmentAsync(queryContext, {
361+
first: 10,
362+
});
363+
364+
expect(result.totalCount).toBeUndefined();
365+
verify(databaseAdapterMock.fetchCountBySQLFragmentAsync(anything(), anything())).never();
366+
});
367+
});
368+
});
256369
});

0 commit comments

Comments
 (0)