Skip to content

Commit 637f5f7

Browse files
authored
feat: Add maxItemCount option for pagination (#112)
1 parent c118abe commit 637f5f7

3 files changed

Lines changed: 69 additions & 5 deletions

File tree

README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,7 @@ const userPair = await users.findManyByIds([id1, id2], {ttl}) // => Promise<(T |
133133
await users.deleteFromCacheById(id}) // => Promise<void>
134134

135135
// query method; note that this returns the CosmosDB FeedResponse object because sometimes this extra information is useful
136-
const userListResponse = await users.findManyByQuery('SELECT * from c where c.type="User"', {ttl, requestOptions}) // => Promise<FeedResponse<T>>
136+
const userListResponse = await users.findManyByQuery('SELECT * from c where c.type="User"', {ttl, requestOptions, maxItemCount}) // => Promise<FeedResponse<T>>
137137
console.log(userListResponse.resources) // user array from query
138138

139139
// create method returns the CosmosDB ItemResponse object

src/datasource.ts

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,11 @@ export interface CosmosQueryDbArgs {
2121
* See https://docs.microsoft.com/en-us/javascript/api/%40azure/cosmos/feedoptions?view=azure-node-latest
2222
*/
2323
requestOptions?: FeedOptions;
24+
/** Specifies the maximum number of items to be returned per page by Cosmos DB.
25+
* This is a convenience shorthand for `requestOptions: { maxItemCount: value }`.
26+
* If both are provided, this top-level `maxItemCount` takes precedence.
27+
*/
28+
maxItemCount?: number;
2429
}
2530

2631
export type QueryFindArgs = FindArgs & CosmosQueryDbArgs;
@@ -47,15 +52,21 @@ export class CosmosDataSource<TData extends { id: string }, TContext = any>
4752
*/
4853
async findManyByQuery(
4954
query: string | SqlQuerySpec,
50-
{ ttl, requestOptions }: QueryFindArgs = {}
55+
{ ttl, requestOptions, maxItemCount }: QueryFindArgs = {}
5156
) {
5257
this.options?.logger?.debug(
5358
// eslint-disable-next-line @typescript-eslint/no-explicit-any
5459
`findManyByQuery: CosmosQuery: ${(query as any).query || query}`
5560
);
56-
const results = await this.container.items
57-
.query<TData>(query, requestOptions)
58-
.fetchAll();
61+
62+
const finalRequestOptions: FeedOptions = { ...requestOptions };
63+
if (maxItemCount !== undefined) {
64+
finalRequestOptions.maxItemCount = maxItemCount;
65+
}
66+
67+
const iterator = this.container.items.query<TData>(query, finalRequestOptions);
68+
const results = await iterator.fetchNext();
69+
5970
// prime these into the dataloader and maybe the cache
6071
if (this.dataLoader && results.resources) {
6172
this.primeLoader(results.resources, ttl);

tests/integration.test.ts

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,59 @@ describe("basic crud", () => {
6969
);
7070
});
7171

72+
describe("pagination tests", function () {
73+
let cosmosDB: CosmosDB;
74+
let userDataSource: UserDataSource;
75+
76+
before(async function () {
77+
cosmosDB = new CosmosDB();
78+
await cosmosDB.createDatabase();
79+
const { container } = await cosmosDB.createContainer({
80+
id: "pagination-test-collection",
81+
});
82+
userDataSource = new UserDataSource(container);
83+
userDataSource.initialize({});
84+
85+
// Create some initial data
86+
await userDataSource.createOne({ id: "user_page_1", email: "[email protected]" });
87+
await userDataSource.createOne({ id: "user_page_2", email: "[email protected]" });
88+
await userDataSource.createOne({ id: "user_page_3", email: "[email protected]" });
89+
});
90+
91+
after(async function () {
92+
await userDataSource.container.delete();
93+
cosmosDB.close();
94+
});
95+
96+
it("should respect maxItemCount when querying", async function () {
97+
const query = "SELECT * FROM c";
98+
const results = await userDataSource.findManyByQuery(query, { maxItemCount: 2 });
99+
100+
expect(results.resources.length).to.equal(2);
101+
// Note: @vercel/cosmosdb-server might not fully implement hasMoreResults
102+
// If this assertion fails, it might be a limitation of the simulator
103+
expect(results.hasMoreResults).to.equal(true);
104+
});
105+
106+
it("should respect maxItemCount via requestOptions when querying", async function () {
107+
const query = "SELECT * FROM c";
108+
const results = await userDataSource.findManyByQuery(query, { requestOptions: { maxItemCount: 1 } });
109+
110+
expect(results.resources.length).to.equal(1);
111+
// Note: @vercel/cosmosdb-server might not fully implement hasMoreResults
112+
expect(results.hasMoreResults).to.equal(true);
113+
});
114+
115+
it("should return all items if maxItemCount is larger than total items", async function () {
116+
const query = "SELECT * FROM c";
117+
const results = await userDataSource.findManyByQuery(query, { maxItemCount: 5 });
118+
119+
expect(results.resources.length).to.equal(3);
120+
// In this case, hasMoreResults should be false as all items are returned
121+
expect(results.hasMoreResults).to.equal(false);
122+
});
123+
});
124+
72125
describe("partitionKey crud tests", function () {
73126
let cosmosDB: CosmosDB;
74127

0 commit comments

Comments
 (0)