Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
19 changes: 15 additions & 4 deletions src/datasource.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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);
Expand Down
53 changes: 53 additions & 0 deletions tests/integration.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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: "page1@example.com" });
await userDataSource.createOne({ id: "user_page_2", email: "page2@example.com" });
await userDataSource.createOne({ id: "user_page_3", email: "page3@example.com" });
});

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;

Expand Down