Fix bug where offset limit iterator was being called for each item under the offset count (#398)

src/queryExecutionContext/EndpointComponent/OffsetLimitEndpointComponent.ts

@@ -1,22 +1,24 @@
 import { Response } from "../../request";
 import { ExecutionContext } from "../ExecutionContext";
-import { getInitialHeader } from "../headerUtils";
+import { getInitialHeader, mergeHeaders } from "../headerUtils";
 
 /** @hidden */
 export class OffsetLimitEndpointComponent implements ExecutionContext {
   constructor(private executionContext: ExecutionContext, private offset: number, private limit: number) {}
 
   public async nextItem(): Promise<Response<any>> {
-    if (this.offset > 0) {
+    const aggregateHeaders = getInitialHeader();
+    while (this.offset > 0) {
       // Grab next item but ignore the result. We only need the headers
       const { headers } = await this.executionContext.nextItem();
       this.offset--;
-      return { result: undefined, headers };
+      mergeHeaders(aggregateHeaders, headers);
     }
     if (this.limit > 0) {
-      const response = await this.executionContext.nextItem();
+      const { result, headers } = await this.executionContext.nextItem();
       this.limit--;
-      return response;
+      mergeHeaders(aggregateHeaders, headers);
+      return { result, headers: aggregateHeaders };
     }
     // If both limit and offset are 0, return nothing
     return { result: undefined, headers: getInitialHeader() };

test/integration/crossPartition.spec.ts

@@ -114,16 +114,19 @@ describe("Cross Partition", function() {
       queryIterator: QueryIterator<any>,
       expectedOrderIds: string[],
       fetchAllResponse: FeedResponse<any>,
-      expectedCount: number
+      expectedCount: number,
+      expectedIteratorCalls: number
     ) {
       const pageSize = options["maxItemCount"];
       let totalExecuteNextRequestCharge = 0;
+      let totalIteratorCalls = 0;
       let totalFetchedResults: any[] = [];
       const expectedLength =
         expectedCount || (expectedOrderIds && expectedOrderIds.length) || documentDefinitions.length;
 
       while (queryIterator.hasMoreResults()) {
         const { resources: results, queryMetrics, requestCharge } = await queryIterator.fetchNext();
+        totalIteratorCalls++;
         assert(queryMetrics, "expected response have query metrics");
 
         if (totalFetchedResults.length > expectedLength) {
@@ -146,6 +149,10 @@ describe("Cross Partition", function() {
         }
       }
 
+      if (expectedIteratorCalls) {
+        assert.equal(totalIteratorCalls, expectedIteratorCalls);
+      }
+
       // no more results
       validateResults(totalFetchedResults, expectedOrderIds, expectedCount);
       assert.equal(queryIterator.hasMoreResults(), false, "hasMoreResults: no more results is left");
@@ -183,13 +190,15 @@ describe("Cross Partition", function() {
       options,
       expectedOrderIds,
       expectedCount,
-      expectedRus
+      expectedRus,
+      expectedIteratorCalls
     }: {
       query: string | SqlQuerySpec;
       options: any;
       expectedOrderIds?: any[];
       expectedCount?: number;
       expectedRus?: number;
+      expectedIteratorCalls?: number;
     }) {
       options.populateQueryMetrics = true;
       const queryIterator = container.items.query(query, options);
@@ -207,7 +216,8 @@ describe("Cross Partition", function() {
         queryIterator,
         expectedOrderIds,
         fetchAllResponse,
-        expectedCount
+        expectedCount,
+        expectedIteratorCalls
       );
       queryIterator.reset();
       await validateAsyncIterator(queryIterator, expectedOrderIds, expectedCount);
@@ -772,7 +782,7 @@ describe("Cross Partition", function() {
 
     it("Validate simple LIMIT OFFSET", async function() {
       const limit = 1;
-      const offset = 2;
+      const offset = 7;
 
       const querySpec = {
         query: `SELECT * FROM root r OFFSET ${offset} LIMIT ${limit}`
@@ -785,7 +795,8 @@ describe("Cross Partition", function() {
       await executeQueryAndValidateResults({
         query: querySpec,
         options,
-        expectedCount: 1
+        expectedCount: 1,
+        expectedIteratorCalls: 1
       });
     });