feat: Add maxItemCount option for pagination (#112)

Author: andrejpk

README.md

@@ -133,7 +133,7 @@ const userPair = await users.findManyByIds([id1, id2], {ttl}) // => Promise<(T |
 await users.deleteFromCacheById(id}) // => Promise<void>
 
 // query method; note that this returns the CosmosDB FeedResponse object because sometimes this extra information is useful
-const userListResponse = await users.findManyByQuery('SELECT * from c where c.type="User"', {ttl, requestOptions}) // => Promise<FeedResponse<T>>
+const userListResponse = await users.findManyByQuery('SELECT * from c where c.type="User"', {ttl, requestOptions, maxItemCount}) // => Promise<FeedResponse<T>>
 console.log(userListResponse.resources) // user array from query
 
 // create method returns the CosmosDB ItemResponse object

src/datasource.ts

@@ -21,6 +21,11 @@ export interface CosmosQueryDbArgs {
    * See https://docs.microsoft.com/en-us/javascript/api/%40azure/cosmos/feedoptions?view=azure-node-latest
    */
   requestOptions?: FeedOptions;
+  /** Specifies the maximum number of items to be returned per page by Cosmos DB.
+   * This is a convenience shorthand for `requestOptions: { maxItemCount: value }`.
+   * If both are provided, this top-level `maxItemCount` takes precedence.
+   */
+  maxItemCount?: number;
 }
 
 export type QueryFindArgs = FindArgs & CosmosQueryDbArgs;
@@ -47,15 +52,21 @@ export class CosmosDataSource<TData extends { id: string }, TContext = any>
    */
   async findManyByQuery(
     query: string | SqlQuerySpec,
-    { ttl, requestOptions }: QueryFindArgs = {}
+    { ttl, requestOptions, maxItemCount }: QueryFindArgs = {}
   ) {
     this.options?.logger?.debug(
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       `findManyByQuery: CosmosQuery: ${(query as any).query || query}`
     );
-    const results = await this.container.items
-      .query<TData>(query, requestOptions)
-      .fetchAll();
+
+    const finalRequestOptions: FeedOptions = { ...requestOptions };
+    if (maxItemCount !== undefined) {
+      finalRequestOptions.maxItemCount = maxItemCount;
+    }
+
+    const iterator = this.container.items.query<TData>(query, finalRequestOptions);
+    const results = await iterator.fetchNext();
+
     // prime these into the dataloader and maybe the cache
     if (this.dataLoader && results.resources) {
       this.primeLoader(results.resources, ttl);

tests/integration.test.ts

@@ -69,6 +69,59 @@ describe("basic crud", () => {
   );
 });
 
+describe("pagination tests", function () {
+  let cosmosDB: CosmosDB;
+  let userDataSource: UserDataSource;
+
+  before(async function () {
+    cosmosDB = new CosmosDB();
+    await cosmosDB.createDatabase();
+    const { container } = await cosmosDB.createContainer({
+      id: "pagination-test-collection",
+    });
+    userDataSource = new UserDataSource(container);
+    userDataSource.initialize({});
+
+    // Create some initial data
+    await userDataSource.createOne({ id: "user_page_1", email: "[email protected]" });
+    await userDataSource.createOne({ id: "user_page_2", email: "[email protected]" });
+    await userDataSource.createOne({ id: "user_page_3", email: "[email protected]" });
+  });
+
+  after(async function () {
+    await userDataSource.container.delete();
+    cosmosDB.close();
+  });
+
+  it("should respect maxItemCount when querying", async function () {
+    const query = "SELECT * FROM c";
+    const results = await userDataSource.findManyByQuery(query, { maxItemCount: 2 });
+
+    expect(results.resources.length).to.equal(2);
+    // Note: @vercel/cosmosdb-server might not fully implement hasMoreResults
+    // If this assertion fails, it might be a limitation of the simulator
+    expect(results.hasMoreResults).to.equal(true); 
+  });
+
+  it("should respect maxItemCount via requestOptions when querying", async function () {
+    const query = "SELECT * FROM c";
+    const results = await userDataSource.findManyByQuery(query, { requestOptions: { maxItemCount: 1 } });
+
+    expect(results.resources.length).to.equal(1);
+    // Note: @vercel/cosmosdb-server might not fully implement hasMoreResults
+    expect(results.hasMoreResults).to.equal(true);
+  });
+
+  it("should return all items if maxItemCount is larger than total items", async function () {
+    const query = "SELECT * FROM c";
+    const results = await userDataSource.findManyByQuery(query, { maxItemCount: 5 });
+
+    expect(results.resources.length).to.equal(3);
+    // In this case, hasMoreResults should be false as all items are returned
+    expect(results.hasMoreResults).to.equal(false);
+  });
+});
+
 describe("partitionKey crud tests", function () {
   let cosmosDB: CosmosDB;