@azure/cosmos 4.4.1 → 4.5.0-alpha.20250718.3

package/dist/browser/globalPartitionEndpointManager.js

@@ -0,0 +1,278 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+import { OperationType, ResourceType, isReadRequest } from "./common/index.js";
+import { Constants, } from "./index.js";
+import { PartitionKeyRangeFailoverInfo } from "./PartitionKeyRangeFailoverInfo.js";
+import { normalizeEndpoint } from "./utils/checkURL.js";
+import { startBackgroundTask } from "./utils/time.js";
+import { assertNotUndefined } from "./utils/typeChecks.js";
+/**
+ * @hidden
+ * This class is used to failover single partitions to different regions.
+ */
+export class GlobalPartitionEndpointManager {
+    /**
+     * @internal
+     */
+    constructor(options, globalEndpointManager) {
+        this.globalEndpointManager = globalEndpointManager;
+        this.partitionKeyRangeToLocationForWrite = new Map();
+        this.partitionKeyRangeToLocationForReadAndWrite = new Map();
+        this.enablePartitionLevelFailover = options.connectionPolicy.enablePartitionLevelFailover;
+        this.enablePartitionLevelCircuitBreaker =
+            options.connectionPolicy.enablePartitionLevelCircuitBreaker ||
+                options.connectionPolicy.enablePartitionLevelFailover;
+        this.preferredLocations = options.connectionPolicy.preferredLocations;
+        this.preferredLocationsCount = this.preferredLocations ? this.preferredLocations.length : 0;
+        if (this.enablePartitionLevelCircuitBreaker) {
+            this.initiateCircuitBreakerFailbackLoop();
+        }
+    }
+    /**
+     * Checks eligibility of the request for partition failover and
+     * tries to mark the endpoint unavailable for the partition key range. Future
+     * requests will be routed to the next location if available.
+     */
+    async tryPartitionLevelFailover(requestContext, diagnosticNode) {
+        if (!(await this.isRequestEligibleForPartitionFailover(requestContext, true))) {
+            return false;
+        }
+        const isRequestEligibleForPerPartitionAutomaticFailover = this.isRequestEligibleForPerPartitionAutomaticFailover(requestContext);
+        const isRequestEligibleForPartitionLevelCircuitBreaker = this.isRequestEligibleForPartitionLevelCircuitBreaker(requestContext);
+        if (isRequestEligibleForPerPartitionAutomaticFailover ||
+            (isRequestEligibleForPartitionLevelCircuitBreaker &&
+                (await this.incrementFailureCounterAndCheckFailover(requestContext, isRequestEligibleForPerPartitionAutomaticFailover, isRequestEligibleForPartitionLevelCircuitBreaker)))) {
+            return this.tryMarkEndpointUnavailableForPartitionKeyRange(requestContext, diagnosticNode, isRequestEligibleForPerPartitionAutomaticFailover, isRequestEligibleForPartitionLevelCircuitBreaker);
+        }
+        return false;
+    }
+    /**
+     * Updates the DocumentServiceRequest routing location to point
+     * new a location based if a partition level failover occurred.
+     */
+    async tryAddPartitionLevelLocationOverride(requestContext, diagnosticNode) {
+        if (!(await this.isRequestEligibleForPartitionFailover(requestContext, false))) {
+            return requestContext;
+        }
+        const partitionKeyRangeId = requestContext.partitionKeyRangeId;
+        if (this.isRequestEligibleForPerPartitionAutomaticFailover(requestContext)) {
+            if (this.partitionKeyRangeToLocationForWrite.has(partitionKeyRangeId)) {
+                const partitionFailOver = this.partitionKeyRangeToLocationForWrite.get(partitionKeyRangeId);
+                requestContext.endpoint = partitionFailOver.getCurrentEndPoint();
+                diagnosticNode.recordEndpointResolution(requestContext.endpoint);
+                return requestContext;
+            }
+        }
+        else if (this.isRequestEligibleForPartitionLevelCircuitBreaker(requestContext)) {
+            if (this.partitionKeyRangeToLocationForReadAndWrite.has(partitionKeyRangeId)) {
+                const partitionFailOver = this.partitionKeyRangeToLocationForReadAndWrite.get(partitionKeyRangeId);
+                const canCircuitBreakerTriggerPartitionFailOver = await partitionFailOver.canCircuitBreakerTriggerPartitionFailOver(isReadRequest(requestContext.operationType));
+                if (canCircuitBreakerTriggerPartitionFailOver) {
+                    requestContext.endpoint = partitionFailOver.getCurrentEndPoint();
+                    diagnosticNode.recordEndpointResolution(requestContext.endpoint);
+                    return requestContext;
+                }
+            }
+        }
+        return requestContext;
+    }
+    /**
+     * This method clears the background refresher for circuit breaker failback
+     * and stops the periodic checks for unhealthy endpoints.
+     */
+    dispose() {
+        if (this.circuitBreakerFailbackBackgroundRefresher) {
+            clearTimeout(this.circuitBreakerFailbackBackgroundRefresher);
+        }
+    }
+    async tryMarkEndpointUnavailableForPartitionKeyRange(requestContext, diagnosticNode, isRequestEligibleForPerPartitionAutomaticFailover, isRequestEligibleForPartitionLevelCircuitBreaker) {
+        const partitionKeyRangeId = requestContext.partitionKeyRangeId;
+        const failedEndPoint = requestContext.endpoint;
+        const readLocations = await this.globalEndpointManager.getReadLocations();
+        const readEndPoints = [];
+        if (isRequestEligibleForPerPartitionAutomaticFailover) {
+            // For any single master write accounts, the next locations to fail over will be the read regions configured at the account level.
+            for (const location of readLocations) {
+                readEndPoints.push(location.databaseAccountEndpoint);
+            }
+            return this.tryAddOrUpdatePartitionFailoverInfoAndMoveToNextLocation(partitionKeyRangeId, failedEndPoint, readEndPoints, this.partitionKeyRangeToLocationForWrite, diagnosticNode);
+        }
+        else if (isRequestEligibleForPartitionLevelCircuitBreaker) {
+            // For the read requests or multi-master write requests, the next locations to fail over will be the preferred locations
+            // configured at the account level plus any other read locations that are not already in the preferred locations.
+            if (this.preferredLocations && this.preferredLocations.length > 0) {
+                for (const preferredLocation of this.preferredLocations) {
+                    const location = readLocations.find((loc) => normalizeEndpoint(loc.name) === normalizeEndpoint(preferredLocation));
+                    if (location) {
+                        readEndPoints.push(location.databaseAccountEndpoint);
+                    }
+                }
+                // Add the rest of the locations not already added
+                for (const location of readLocations) {
+                    if (!readEndPoints.includes(location.databaseAccountEndpoint)) {
+                        readEndPoints.push(location.databaseAccountEndpoint);
+                    }
+                }
+            }
+            else {
+                for (const location of readLocations) {
+                    readEndPoints.push(location.databaseAccountEndpoint);
+                }
+            }
+            return this.tryAddOrUpdatePartitionFailoverInfoAndMoveToNextLocation(partitionKeyRangeId, failedEndPoint, readEndPoints, this.partitionKeyRangeToLocationForReadAndWrite, diagnosticNode);
+        }
+        return false;
+    }
+    /**
+     * Increments the failure counter for the specified partition and checks if the partition can fail over.
+     * This method is used to determine if a partition should be failed over based on the number of request failures.
+     */
+    async incrementFailureCounterAndCheckFailover(requestContext, isRequestEligibleForPerPartitionAutomaticFailover, isRequestEligibleForPartitionLevelCircuitBreaker) {
+        const partitionKeyRangeId = requestContext.partitionKeyRangeId;
+        const failedEndPoint = requestContext.endpoint;
+        let partitionKeyRangeFailoverInfo;
+        if (isRequestEligibleForPerPartitionAutomaticFailover) {
+            if (!this.partitionKeyRangeToLocationForWrite.has(partitionKeyRangeId)) {
+                // If the partition key range is not already in the map, add it
+                const failoverInfo = new PartitionKeyRangeFailoverInfo(failedEndPoint);
+                this.partitionKeyRangeToLocationForWrite.set(partitionKeyRangeId, failoverInfo);
+            }
+            partitionKeyRangeFailoverInfo =
+                this.partitionKeyRangeToLocationForWrite.get(partitionKeyRangeId);
+        }
+        else if (isRequestEligibleForPartitionLevelCircuitBreaker) {
+            if (!this.partitionKeyRangeToLocationForReadAndWrite.has(partitionKeyRangeId)) {
+                // If the partition key range is not already in the map, add it
+                const failoverInfo = new PartitionKeyRangeFailoverInfo(failedEndPoint);
+                this.partitionKeyRangeToLocationForReadAndWrite.set(partitionKeyRangeId, failoverInfo);
+            }
+            partitionKeyRangeFailoverInfo =
+                this.partitionKeyRangeToLocationForReadAndWrite.get(partitionKeyRangeId);
+        }
+        else {
+            return false;
+        }
+        assertNotUndefined(partitionKeyRangeFailoverInfo, "partitionKeyRangeFailoverInfo should be set if failover flags are true.");
+        const currentTimeInMilliseconds = Date.now();
+        await partitionKeyRangeFailoverInfo.incrementRequestFailureCounts(isReadRequest(requestContext.operationType), currentTimeInMilliseconds);
+        return partitionKeyRangeFailoverInfo.canCircuitBreakerTriggerPartitionFailOver(isReadRequest(requestContext.operationType));
+    }
+    /** Validates if the given request is eligible for partition failover. */
+    async isRequestEligibleForPartitionFailover(requestContext, shouldValidateFailedLocation) {
+        if (!requestContext ||
+            !requestContext.operationType ||
+            !requestContext.resourceType ||
+            !requestContext.partitionKeyRangeId) {
+            return false;
+        }
+        const canUsePartitionLevelFailoverLocations = await this.canUsePartitionLevelFailoverLocations(requestContext.operationType, requestContext.resourceType);
+        if (!canUsePartitionLevelFailoverLocations) {
+            return false;
+        }
+        if (shouldValidateFailedLocation && !requestContext.endpoint) {
+            return false;
+        }
+        return true;
+    }
+    /** Determines if partition level failover locations can be used for the given request. */
+    async canUsePartitionLevelFailoverLocations(operationType, resourceType) {
+        const readEndPoints = await this.globalEndpointManager.getReadEndpoints();
+        if (readEndPoints.length <= 1) {
+            return false;
+        }
+        if (resourceType === ResourceType.item ||
+            (resourceType === ResourceType.sproc && operationType === OperationType.Execute)) {
+            // Right now, for single-master only reads are supported for circuit breaker, and writes are supported for automatic.
+            // failover. For multi master, both reads and writes are supported. Hence return true for both the cases.
+            return true;
+        }
+        return false;
+    }
+    /**
+     * Determines if a request is eligible for per-partition automatic failover.
+     * A request is eligible if it is a write request, partition level failover is enabled,
+     * and the global endpoint manager cannot use multiple write locations for the request.
+     */
+    isRequestEligibleForPerPartitionAutomaticFailover(requestContext) {
+        return (this.enablePartitionLevelFailover &&
+            !isReadRequest(requestContext.operationType) &&
+            !this.globalEndpointManager.canUseMultipleWriteLocations(requestContext.resourceType, requestContext.operationType));
+    }
+    /**
+     * Determines if a request is eligible for partition-level circuit breaker.
+     * This method checks if partition-level circuit breaker is enabled, and if the request is a read-only request or
+     * the global endpoint manager can use multiple write locations for the request.
+     */
+    isRequestEligibleForPartitionLevelCircuitBreaker(requestContext) {
+        if (!this.enablePartitionLevelCircuitBreaker) {
+            return false;
+        }
+        if (isReadRequest(requestContext.operationType)) {
+            return true;
+        }
+        return this.globalEndpointManager.canUseMultipleWriteLocations(requestContext.resourceType, requestContext.operationType);
+    }
+    /**
+     * Attempts to add or update the partition failover information and move to the next available location.
+     * This method checks if the current location for the partition key range has failed and updates the failover
+     * information to route the request to the next available location. If all locations have been tried, it removes
+     * the failover information for the partition key range. Return True if the failover information was successfully
+     * updated and the request was routed to a new location, otherwise false.
+     */
+    async tryAddOrUpdatePartitionFailoverInfoAndMoveToNextLocation(partitionKeyRangeId, failedEndPoint, nextEndPoints, partitionKeyRangeToLocation, diagnosticNode) {
+        if (!partitionKeyRangeToLocation.has(partitionKeyRangeId)) {
+            // If the partition key range is not already in the map, add it
+            const failoverInfo = new PartitionKeyRangeFailoverInfo(failedEndPoint);
+            partitionKeyRangeToLocation.set(partitionKeyRangeId, failoverInfo);
+        }
+        const partitionFailOver = partitionKeyRangeToLocation.get(partitionKeyRangeId);
+        // Will return true if it was able to update to a new region
+        if (await partitionFailOver.tryMoveNextLocation(nextEndPoints, failedEndPoint, diagnosticNode, partitionKeyRangeId)) {
+            return true;
+        }
+        // All the locations have been tried. Remove the override information
+        partitionKeyRangeToLocation.delete(partitionKeyRangeId);
+        return false;
+    }
+    /**
+     * Initiates a background loop that periodically checks for unhealthy endpoints
+     * and attempts to open connections to them. If a connection is successfully
+     * established, it initiates a failback to the original location for the partition key range.
+     * This is useful for scenarios where a partition key range has been marked as unavailable
+     * due to a circuit breaker, and we want to periodically check if the original location
+     * has become healthy again.
+     * The loop runs at a defined interval specified by Constants.StalePartitionUnavailabilityRefreshIntervalInMs.
+     */
+    initiateCircuitBreakerFailbackLoop() {
+        this.circuitBreakerFailbackBackgroundRefresher = startBackgroundTask(async () => {
+            try {
+                await this.openConnectionToUnhealthyEndpointsWithFailback();
+            }
+            catch (err) {
+                console.error("Failed to open connection to unhealthy endpoints: ", err);
+            }
+        }, Constants.StalePartitionUnavailabilityRefreshIntervalInMs);
+    }
+    /**
+     * Attempts to open connections to unhealthy endpoints and initiates failback if the connections are successful.
+     * This method checks the partition key ranges that have failed locations and tries to re-establish connections
+     * to those locations. If a connection is successfully re-established, it initiates a failback to the original
+     * location for the partition key range.
+     */
+    async openConnectionToUnhealthyEndpointsWithFailback() {
+        for (const pkRange of this.partitionKeyRangeToLocationForReadAndWrite.keys()) {
+            const partitionFailover = this.partitionKeyRangeToLocationForReadAndWrite.get(pkRange);
+            if (!partitionFailover)
+                continue;
+            const { firstRequestFailureTime } = await partitionFailover.snapshotPartitionFailoverTimestamps();
+            const now = new Date();
+            if (now.getTime() - firstRequestFailureTime >
+                Constants.AllowedPartitionUnavailabilityDurationInMs) {
+                // Un-deterministically marking the original failed endpoint for the PkRange back to healthy.
+                // Initiate Failback to the original failed location.
+                this.partitionKeyRangeToLocationForReadAndWrite.delete(pkRange);
+            }
+        }
+    }
+}
+//# sourceMappingURL=globalPartitionEndpointManager.js.map

package/dist/browser/globalPartitionEndpointManager.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"globalPartitionEndpointManager.js","sourceRoot":"","sources":["../../src/globalPartitionEndpointManager.ts"],"names":[],"mappings":"AAAA,uCAAuC;AACvC,kCAAkC;AAClC,OAAO,EAAE,aAAa,EAAE,YAAY,EAAE,aAAa,EAAE,MAAM,mBAAmB,CAAC;AAC/E,OAAO,EACL,SAAS,GAKV,MAAM,YAAY,CAAC;AACpB,OAAO,EAAE,6BAA6B,EAAE,MAAM,oCAAoC,CAAC;AACnF,OAAO,EAAE,iBAAiB,EAAE,MAAM,qBAAqB,CAAC;AACxD,OAAO,EAAE,mBAAmB,EAAE,MAAM,iBAAiB,CAAC;AACtD,OAAO,EAAE,kBAAkB,EAAE,MAAM,uBAAuB,CAAC;AAE3D;;;GAGG;AACH,MAAM,OAAO,8BAA8B;IAYzC;;OAEG;IACH,YACE,OAA4B,EACpB,qBAA4C;QAA5C,0BAAqB,GAArB,qBAAqB,CAAuB;QAEpD,IAAI,CAAC,mCAAmC,GAAG,IAAI,GAAG,EAAyC,CAAC;QAC5F,IAAI,CAAC,0CAA0C,GAAG,IAAI,GAAG,EAGtD,CAAC;QAEJ,IAAI,CAAC,4BAA4B,GAAG,OAAO,CAAC,gBAAgB,CAAC,4BAA4B,CAAC;QAC1F,IAAI,CAAC,kCAAkC;YACrC,OAAO,CAAC,gBAAgB,CAAC,kCAAkC;gBAC3D,OAAO,CAAC,gBAAgB,CAAC,4BAA4B,CAAC;QAExD,IAAI,CAAC,kBAAkB,GAAG,OAAO,CAAC,gBAAgB,CAAC,kBAAkB,CAAC;QACtE,IAAI,CAAC,uBAAuB,GAAG,IAAI,CAAC,kBAAkB,CAAC,CAAC,CAAC,IAAI,CAAC,kBAAkB,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC;QAC5F,IAAI,IAAI,CAAC,kCAAkC,EAAE,CAAC;YAC5C,IAAI,CAAC,kCAAkC,EAAE,CAAC;QAC5C,CAAC;IACH,CAAC;IAED;;;;OAIG;IACI,KAAK,CAAC,yBAAyB,CACpC,cAA8B,EAC9B,cAAsC;QAEtC,IAAI,CAAC,CAAC,MAAM,IAAI,CAAC,qCAAqC,CAAC,cAAc,EAAE,IAAI,CAAC,CAAC,EAAE,CAAC;YAC9E,OAAO,KAAK,CAAC;QACf,CAAC;QAED,MAAM,iDAAiD,GACrD,IAAI,CAAC,iDAAiD,CAAC,cAAc,CAAC,CAAC;QACzE,MAAM,gDAAgD,GACpD,IAAI,CAAC,gDAAgD,CAAC,cAAc,CAAC,CAAC;QAExE,IACE,iDAAiD;YACjD,CAAC,gDAAgD;gBAC/C,CAAC,MAAM,IAAI,CAAC,uCAAuC,CACjD,cAAc,EACd,iDAAiD,EACjD,gDAAgD,CACjD,CAAC,CAAC,EACL,CAAC;YACD,OAAO,IAAI,CAAC,8CAA8C,CACxD,cAAc,EACd,cAAc,EACd,iDAAiD,EACjD,gDAAgD,CACjD,CAAC;QACJ,CAAC;QACD,OAAO,KAAK,CAAC;IACf,CAAC;IAED;;;OAGG;IACI,KAAK,CAAC,oCAAoC,CAC/C,cAA8B,EAC9B,cAAsC;QAEtC,IAAI,CAAC,CAAC,MAAM,IAAI,CAAC,qCAAqC,CAAC,cAAc,EAAE,KAAK,CAAC,CAAC,EAAE,CAAC;YAC/E,OAAO,cAAc,CAAC;QACxB,CAAC;QAED,MAAM,mBAAmB,GAAG,cAAc,CAAC,mBAAmB,CAAC;QAE/D,IAAI,IAAI,CAAC,iDAAiD,CAAC,cAAc,CAAC,EAAE,CAAC;YAC3E,IAAI,IAAI,CAAC,mCAAmC,CAAC,GAAG,CAAC,mBAAmB,CAAC,EAAE,CAAC;gBACtE,MAAM,iBAAiB,GAAG,IAAI,CAAC,mCAAmC,CAAC,GAAG,CAAC,mBAAmB,CAAC,CAAC;gBAC5F,cAAc,CAAC,QAAQ,GAAG,iBAAiB,CAAC,kBAAkB,EAAE,CAAC;gBACjE,cAAc,CAAC,wBAAwB,CAAC,cAAc,CAAC,QAAQ,CAAC,CAAC;gBACjE,OAAO,cAAc,CAAC;YACxB,CAAC;QACH,CAAC;aAAM,IAAI,IAAI,CAAC,gDAAgD,CAAC,cAAc,CAAC,EAAE,CAAC;YACjF,IAAI,IAAI,CAAC,0CAA0C,CAAC,GAAG,CAAC,mBAAmB,CAAC,EAAE,CAAC;gBAC7E,MAAM,iBAAiB,GACrB,IAAI,CAAC,0CAA0C,CAAC,GAAG,CAAC,mBAAmB,CAAC,CAAC;gBAE3E,MAAM,yCAAyC,GAC7C,MAAM,iBAAiB,CAAC,yCAAyC,CAC/D,aAAa,CAAC,cAAc,CAAC,aAAa,CAAC,CAC5C,CAAC;gBACJ,IAAI,yCAAyC,EAAE,CAAC;oBAC9C,cAAc,CAAC,QAAQ,GAAG,iBAAiB,CAAC,kBAAkB,EAAE,CAAC;oBACjE,cAAc,CAAC,wBAAwB,CAAC,cAAc,CAAC,QAAQ,CAAC,CAAC;oBACjE,OAAO,cAAc,CAAC;gBACxB,CAAC;YACH,CAAC;QACH,CAAC;QACD,OAAO,cAAc,CAAC;IACxB,CAAC;IAED;;;OAGG;IACI,OAAO;QACZ,IAAI,IAAI,CAAC,yCAAyC,EAAE,CAAC;YACnD,YAAY,CAAC,IAAI,CAAC,yCAAyC,CAAC,CAAC;QAC/D,CAAC;IACH,CAAC;IAEO,KAAK,CAAC,8CAA8C,CAC1D,cAA8B,EAC9B,cAAsC,EACtC,iDAA0D,EAC1D,gDAAyD;QAEzD,MAAM,mBAAmB,GAAG,cAAc,CAAC,mBAAmB,CAAC;QAC/D,MAAM,cAAc,GAAG,cAAc,CAAC,QAAQ,CAAC;QAE/C,MAAM,aAAa,GAAG,MAAM,IAAI,CAAC,qBAAqB,CAAC,gBAAgB,EAAE,CAAC;QAC1E,MAAM,aAAa,GAAa,EAAE,CAAC;QAEnC,IAAI,iDAAiD,EAAE,CAAC;YACtD,kIAAkI;YAClI,KAAK,MAAM,QAAQ,IAAI,aAAa,EAAE,CAAC;gBACrC,aAAa,CAAC,IAAI,CAAC,QAAQ,CAAC,uBAAuB,CAAC,CAAC;YACvD,CAAC;YACD,OAAO,IAAI,CAAC,wDAAwD,CAClE,mBAAmB,EACnB,cAAc,EACd,aAAa,EACb,IAAI,CAAC,mCAAmC,EACxC,cAAc,CACf,CAAC;QACJ,CAAC;aAAM,IAAI,gDAAgD,EAAE,CAAC;YAC5D,wHAAwH;YACxH,iHAAiH;YACjH,IAAI,IAAI,CAAC,kBAAkB,IAAI,IAAI,CAAC,kBAAkB,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC;gBAClE,KAAK,MAAM,iBAAiB,IAAI,IAAI,CAAC,kBAAkB,EAAE,CAAC;oBACxD,MAAM,QAAQ,GAAG,aAAa,CAAC,IAAI,CACjC,CAAC,GAAG,EAAE,EAAE,CAAC,iBAAiB,CAAC,GAAG,CAAC,IAAI,CAAC,KAAK,iBAAiB,CAAC,iBAAiB,CAAC,CAC9E,CAAC;oBACF,IAAI,QAAQ,EAAE,CAAC;wBACb,aAAa,CAAC,IAAI,CAAC,QAAQ,CAAC,uBAAuB,CAAC,CAAC;oBACvD,CAAC;gBACH,CAAC;gBAED,kDAAkD;gBAClD,KAAK,MAAM,QAAQ,IAAI,aAAa,EAAE,CAAC;oBACrC,IAAI,CAAC,aAAa,CAAC,QAAQ,CAAC,QAAQ,CAAC,uBAAuB,CAAC,EAAE,CAAC;wBAC9D,aAAa,CAAC,IAAI,CAAC,QAAQ,CAAC,uBAAuB,CAAC,CAAC;oBACvD,CAAC;gBACH,CAAC;YACH,CAAC;iBAAM,CAAC;gBACN,KAAK,MAAM,QAAQ,IAAI,aAAa,EAAE,CAAC;oBACrC,aAAa,CAAC,IAAI,CAAC,QAAQ,CAAC,uBAAuB,CAAC,CAAC;gBACvD,CAAC;YACH,CAAC;YACD,OAAO,IAAI,CAAC,wDAAwD,CAClE,mBAAmB,EACnB,cAAc,EACd,aAAa,EACb,IAAI,CAAC,0CAA0C,EAC/C,cAAc,CACf,CAAC;QACJ,CAAC;QACD,OAAO,KAAK,CAAC;IACf,CAAC;IAED;;;OAGG;IACK,KAAK,CAAC,uCAAuC,CACnD,cAA8B,EAC9B,iDAA0D,EAC1D,gDAAyD;QAEzD,MAAM,mBAAmB,GAAG,cAAc,CAAC,mBAAmB,CAAC;QAC/D,MAAM,cAAc,GAAG,cAAc,CAAC,QAAQ,CAAC;QAC/C,IAAI,6BAA4D,CAAC;QAEjE,IAAI,iDAAiD,EAAE,CAAC;YACtD,IAAI,CAAC,IAAI,CAAC,mCAAmC,CAAC,GAAG,CAAC,mBAAmB,CAAC,EAAE,CAAC;gBACvE,+DAA+D;gBAC/D,MAAM,YAAY,GAAG,IAAI,6BAA6B,CAAC,cAAc,CAAC,CAAC;gBACvE,IAAI,CAAC,mCAAmC,CAAC,GAAG,CAAC,mBAAmB,EAAE,YAAY,CAAC,CAAC;YAClF,CAAC;YACD,6BAA6B;gBAC3B,IAAI,CAAC,mCAAmC,CAAC,GAAG,CAAC,mBAAmB,CAAC,CAAC;QACtE,CAAC;aAAM,IAAI,gDAAgD,EAAE,CAAC;YAC5D,IAAI,CAAC,IAAI,CAAC,0CAA0C,CAAC,GAAG,CAAC,mBAAmB,CAAC,EAAE,CAAC;gBAC9E,+DAA+D;gBAC/D,MAAM,YAAY,GAAG,IAAI,6BAA6B,CAAC,cAAc,CAAC,CAAC;gBACvE,IAAI,CAAC,0CAA0C,CAAC,GAAG,CAAC,mBAAmB,EAAE,YAAY,CAAC,CAAC;YACzF,CAAC;YACD,6BAA6B;gBAC3B,IAAI,CAAC,0CAA0C,CAAC,GAAG,CAAC,mBAAmB,CAAC,CAAC;QAC7E,CAAC;aAAM,CAAC;YACN,OAAO,KAAK,CAAC;QACf,CAAC;QAED,kBAAkB,CAChB,6BAA6B,EAC7B,yEAAyE,CAC1E,CAAC;QAEF,MAAM,yBAAyB,GAAG,IAAI,CAAC,GAAG,EAAE,CAAC;QAC7C,MAAM,6BAA6B,CAAC,6BAA6B,CAC/D,aAAa,CAAC,cAAc,CAAC,aAAa,CAAC,EAC3C,yBAAyB,CAC1B,CAAC;QAEF,OAAO,6BAA6B,CAAC,yCAAyC,CAC5E,aAAa,CAAC,cAAc,CAAC,aAAa,CAAC,CAC5C,CAAC;IACJ,CAAC;IAED,yEAAyE;IACjE,KAAK,CAAC,qCAAqC,CACjD,cAA8B,EAC9B,4BAAqC;QAErC,IACE,CAAC,cAAc;YACf,CAAC,cAAc,CAAC,aAAa;YAC7B,CAAC,cAAc,CAAC,YAAY;YAC5B,CAAC,cAAc,CAAC,mBAAmB,EACnC,CAAC;YACD,OAAO,KAAK,CAAC;QACf,CAAC;QAED,MAAM,qCAAqC,GAAG,MAAM,IAAI,CAAC,qCAAqC,CAC5F,cAAc,CAAC,aAAa,EAC5B,cAAc,CAAC,YAAY,CAC5B,CAAC;QACF,IAAI,CAAC,qCAAqC,EAAE,CAAC;YAC3C,OAAO,KAAK,CAAC;QACf,CAAC;QAED,IAAI,4BAA4B,IAAI,CAAC,cAAc,CAAC,QAAQ,EAAE,CAAC;YAC7D,OAAO,KAAK,CAAC;QACf,CAAC;QACD,OAAO,IAAI,CAAC;IACd,CAAC;IAED,0FAA0F;IAClF,KAAK,CAAC,qCAAqC,CACjD,aAA6B,EAC7B,YAA2B;QAE3B,MAAM,aAAa,GAAG,MAAM,IAAI,CAAC,qBAAqB,CAAC,gBAAgB,EAAE,CAAC;QAC1E,IAAI,aAAa,CAAC,MAAM,IAAI,CAAC,EAAE,CAAC;YAC9B,OAAO,KAAK,CAAC;QACf,CAAC;QACD,IACE,YAAY,KAAK,YAAY,CAAC,IAAI;YAClC,CAAC,YAAY,KAAK,YAAY,CAAC,KAAK,IAAI,aAAa,KAAK,aAAa,CAAC,OAAO,CAAC,EAChF,CAAC;YACD,qHAAqH;YACrH,yGAAyG;YACzG,OAAO,IAAI,CAAC;QACd,CAAC;QACD,OAAO,KAAK,CAAC;IACf,CAAC;IAED;;;;OAIG;IACK,iDAAiD,CACvD,cAA8B;QAE9B,OAAO,CACL,IAAI,CAAC,4BAA4B;YACjC,CAAC,aAAa,CAAC,cAAc,CAAC,aAAa,CAAC;YAC5C,CAAC,IAAI,CAAC,qBAAqB,CAAC,4BAA4B,CACtD,cAAc,CAAC,YAAY,EAC3B,cAAc,CAAC,aAAa,CAC7B,CACF,CAAC;IACJ,CAAC;IAED;;;;OAIG;IACK,gDAAgD,CACtD,cAA8B;QAE9B,IAAI,CAAC,IAAI,CAAC,kCAAkC,EAAE,CAAC;YAC7C,OAAO,KAAK,CAAC;QACf,CAAC;QACD,IAAI,aAAa,CAAC,cAAc,CAAC,aAAa,CAAC,EAAE,CAAC;YAChD,OAAO,IAAI,CAAC;QACd,CAAC;QACD,OAAO,IAAI,CAAC,qBAAqB,CAAC,4BAA4B,CAC5D,cAAc,CAAC,YAAY,EAC3B,cAAc,CAAC,aAAa,CAC7B,CAAC;IACJ,CAAC;IAED;;;;;;OAMG;IACK,KAAK,CAAC,wDAAwD,CACpE,mBAA2B,EAC3B,cAAsB,EACtB,aAAgC,EAChC,2BAAuE,EACvE,cAAsC;QAEtC,IAAI,CAAC,2BAA2B,CAAC,GAAG,CAAC,mBAAmB,CAAC,EAAE,CAAC;YAC1D,+DAA+D;YAC/D,MAAM,YAAY,GAAG,IAAI,6BAA6B,CAAC,cAAc,CAAC,CAAC;YACvE,2BAA2B,CAAC,GAAG,CAAC,mBAAmB,EAAE,YAAY,CAAC,CAAC;QACrE,CAAC;QACD,MAAM,iBAAiB,GAAG,2BAA2B,CAAC,GAAG,CAAC,mBAAmB,CAAC,CAAC;QAE/E,4DAA4D;QAC5D,IACE,MAAM,iBAAiB,CAAC,mBAAmB,CACzC,aAAa,EACb,cAAc,EACd,cAAc,EACd,mBAAmB,CACpB,EACD,CAAC;YACD,OAAO,IAAI,CAAC;QACd,CAAC;QACD,qEAAqE;QACrE,2BAA2B,CAAC,MAAM,CAAC,mBAAmB,CAAC,CAAC;QACxD,OAAO,KAAK,CAAC;IACf,CAAC;IAED;;;;;;;;OAQG;IACK,kCAAkC;QACxC,IAAI,CAAC,yCAAyC,GAAG,mBAAmB,CAAC,KAAK,IAAI,EAAE;YAC9E,IAAI,CAAC;gBACH,MAAM,IAAI,CAAC,8CAA8C,EAAE,CAAC;YAC9D,CAAC;YAAC,OAAO,GAAG,EAAE,CAAC;gBACb,OAAO,CAAC,KAAK,CAAC,oDAAoD,EAAE,GAAG,CAAC,CAAC;YAC3E,CAAC;QACH,CAAC,EAAE,SAAS,CAAC,+CAA+C,CAAC,CAAC;IAChE,CAAC;IAED;;;;;OAKG;IACK,KAAK,CAAC,8CAA8C;QAC1D,KAAK,MAAM,OAAO,IAAI,IAAI,CAAC,0CAA0C,CAAC,IAAI,EAAE,EAAE,CAAC;YAC7E,MAAM,iBAAiB,GAAG,IAAI,CAAC,0CAA0C,CAAC,GAAG,CAAC,OAAO,CAAC,CAAC;YACvF,IAAI,CAAC,iBAAiB;gBAAE,SAAS;YAEjC,MAAM,EAAE,uBAAuB,EAAE,GAC/B,MAAM,iBAAiB,CAAC,mCAAmC,EAAE,CAAC;YAChE,MAAM,GAAG,GAAG,IAAI,IAAI,EAAE,CAAC;YAEvB,IACE,GAAG,CAAC,OAAO,EAAE,GAAG,uBAAuB;gBACvC,SAAS,CAAC,0CAA0C,EACpD,CAAC;gBACD,6FAA6F;gBAC7F,qDAAqD;gBACrD,IAAI,CAAC,0CAA0C,CAAC,MAAM,CAAC,OAAO,CAAC,CAAC;YAClE,CAAC;QACH,CAAC;IACH,CAAC;CACF","sourcesContent":["// Copyright (c) Microsoft Corporation.\n// Licensed under the MIT License.\nimport { OperationType, ResourceType, isReadRequest } from \"./common/index.js\";\nimport {\n  Constants,\n  DiagnosticNodeInternal,\n  type CosmosClientOptions,\n  type GlobalEndpointManager,\n  type RequestContext,\n} from \"./index.js\";\nimport { PartitionKeyRangeFailoverInfo } from \"./PartitionKeyRangeFailoverInfo.js\";\nimport { normalizeEndpoint } from \"./utils/checkURL.js\";\nimport { startBackgroundTask } from \"./utils/time.js\";\nimport { assertNotUndefined } from \"./utils/typeChecks.js\";\n\n/**\n * @hidden\n * This class is used to failover single partitions to different regions.\n */\nexport class GlobalPartitionEndpointManager {\n  private readonly partitionKeyRangeToLocationForWrite: Map<string, PartitionKeyRangeFailoverInfo>;\n  private readonly partitionKeyRangeToLocationForReadAndWrite: Map<\n    string,\n    PartitionKeyRangeFailoverInfo\n  >;\n  private enablePartitionLevelFailover: boolean;\n  private enablePartitionLevelCircuitBreaker: boolean;\n  private preferredLocations: string[];\n  public preferredLocationsCount: number;\n  private circuitBreakerFailbackBackgroundRefresher: NodeJS.Timeout;\n\n  /**\n   * @internal\n   */\n  constructor(\n    options: CosmosClientOptions,\n    private globalEndpointManager: GlobalEndpointManager,\n  ) {\n    this.partitionKeyRangeToLocationForWrite = new Map<string, PartitionKeyRangeFailoverInfo>();\n    this.partitionKeyRangeToLocationForReadAndWrite = new Map<\n      string,\n      PartitionKeyRangeFailoverInfo\n    >();\n\n    this.enablePartitionLevelFailover = options.connectionPolicy.enablePartitionLevelFailover;\n    this.enablePartitionLevelCircuitBreaker =\n      options.connectionPolicy.enablePartitionLevelCircuitBreaker ||\n      options.connectionPolicy.enablePartitionLevelFailover;\n\n    this.preferredLocations = options.connectionPolicy.preferredLocations;\n    this.preferredLocationsCount = this.preferredLocations ? this.preferredLocations.length : 0;\n    if (this.enablePartitionLevelCircuitBreaker) {\n      this.initiateCircuitBreakerFailbackLoop();\n    }\n  }\n\n  /**\n   * Checks eligibility of the request for partition failover and\n   * tries to mark the endpoint unavailable for the partition key range. Future\n   * requests will be routed to the next location if available.\n   */\n  public async tryPartitionLevelFailover(\n    requestContext: RequestContext,\n    diagnosticNode: DiagnosticNodeInternal,\n  ): Promise<boolean> {\n    if (!(await this.isRequestEligibleForPartitionFailover(requestContext, true))) {\n      return false;\n    }\n\n    const isRequestEligibleForPerPartitionAutomaticFailover =\n      this.isRequestEligibleForPerPartitionAutomaticFailover(requestContext);\n    const isRequestEligibleForPartitionLevelCircuitBreaker =\n      this.isRequestEligibleForPartitionLevelCircuitBreaker(requestContext);\n\n    if (\n      isRequestEligibleForPerPartitionAutomaticFailover ||\n      (isRequestEligibleForPartitionLevelCircuitBreaker &&\n        (await this.incrementFailureCounterAndCheckFailover(\n          requestContext,\n          isRequestEligibleForPerPartitionAutomaticFailover,\n          isRequestEligibleForPartitionLevelCircuitBreaker,\n        )))\n    ) {\n      return this.tryMarkEndpointUnavailableForPartitionKeyRange(\n        requestContext,\n        diagnosticNode,\n        isRequestEligibleForPerPartitionAutomaticFailover,\n        isRequestEligibleForPartitionLevelCircuitBreaker,\n      );\n    }\n    return false;\n  }\n\n  /**\n   * Updates the DocumentServiceRequest routing location to point\n   * new a location based if a partition level failover occurred.\n   */\n  public async tryAddPartitionLevelLocationOverride(\n    requestContext: RequestContext,\n    diagnosticNode: DiagnosticNodeInternal,\n  ): Promise<RequestContext> {\n    if (!(await this.isRequestEligibleForPartitionFailover(requestContext, false))) {\n      return requestContext;\n    }\n\n    const partitionKeyRangeId = requestContext.partitionKeyRangeId;\n\n    if (this.isRequestEligibleForPerPartitionAutomaticFailover(requestContext)) {\n      if (this.partitionKeyRangeToLocationForWrite.has(partitionKeyRangeId)) {\n        const partitionFailOver = this.partitionKeyRangeToLocationForWrite.get(partitionKeyRangeId);\n        requestContext.endpoint = partitionFailOver.getCurrentEndPoint();\n        diagnosticNode.recordEndpointResolution(requestContext.endpoint);\n        return requestContext;\n      }\n    } else if (this.isRequestEligibleForPartitionLevelCircuitBreaker(requestContext)) {\n      if (this.partitionKeyRangeToLocationForReadAndWrite.has(partitionKeyRangeId)) {\n        const partitionFailOver =\n          this.partitionKeyRangeToLocationForReadAndWrite.get(partitionKeyRangeId);\n\n        const canCircuitBreakerTriggerPartitionFailOver =\n          await partitionFailOver.canCircuitBreakerTriggerPartitionFailOver(\n            isReadRequest(requestContext.operationType),\n          );\n        if (canCircuitBreakerTriggerPartitionFailOver) {\n          requestContext.endpoint = partitionFailOver.getCurrentEndPoint();\n          diagnosticNode.recordEndpointResolution(requestContext.endpoint);\n          return requestContext;\n        }\n      }\n    }\n    return requestContext;\n  }\n\n  /**\n   * This method clears the background refresher for circuit breaker failback\n   * and stops the periodic checks for unhealthy endpoints.\n   */\n  public dispose(): void {\n    if (this.circuitBreakerFailbackBackgroundRefresher) {\n      clearTimeout(this.circuitBreakerFailbackBackgroundRefresher);\n    }\n  }\n\n  private async tryMarkEndpointUnavailableForPartitionKeyRange(\n    requestContext: RequestContext,\n    diagnosticNode: DiagnosticNodeInternal,\n    isRequestEligibleForPerPartitionAutomaticFailover: boolean,\n    isRequestEligibleForPartitionLevelCircuitBreaker: boolean,\n  ): Promise<boolean> {\n    const partitionKeyRangeId = requestContext.partitionKeyRangeId;\n    const failedEndPoint = requestContext.endpoint;\n\n    const readLocations = await this.globalEndpointManager.getReadLocations();\n    const readEndPoints: string[] = [];\n\n    if (isRequestEligibleForPerPartitionAutomaticFailover) {\n      // For any single master write accounts, the next locations to fail over will be the read regions configured at the account level.\n      for (const location of readLocations) {\n        readEndPoints.push(location.databaseAccountEndpoint);\n      }\n      return this.tryAddOrUpdatePartitionFailoverInfoAndMoveToNextLocation(\n        partitionKeyRangeId,\n        failedEndPoint,\n        readEndPoints,\n        this.partitionKeyRangeToLocationForWrite,\n        diagnosticNode,\n      );\n    } else if (isRequestEligibleForPartitionLevelCircuitBreaker) {\n      // For the read requests or multi-master write requests, the next locations to fail over will be the preferred locations\n      // configured at the account level plus any other read locations that are not already in the preferred locations.\n      if (this.preferredLocations && this.preferredLocations.length > 0) {\n        for (const preferredLocation of this.preferredLocations) {\n          const location = readLocations.find(\n            (loc) => normalizeEndpoint(loc.name) === normalizeEndpoint(preferredLocation),\n          );\n          if (location) {\n            readEndPoints.push(location.databaseAccountEndpoint);\n          }\n        }\n\n        // Add the rest of the locations not already added\n        for (const location of readLocations) {\n          if (!readEndPoints.includes(location.databaseAccountEndpoint)) {\n            readEndPoints.push(location.databaseAccountEndpoint);\n          }\n        }\n      } else {\n        for (const location of readLocations) {\n          readEndPoints.push(location.databaseAccountEndpoint);\n        }\n      }\n      return this.tryAddOrUpdatePartitionFailoverInfoAndMoveToNextLocation(\n        partitionKeyRangeId,\n        failedEndPoint,\n        readEndPoints,\n        this.partitionKeyRangeToLocationForReadAndWrite,\n        diagnosticNode,\n      );\n    }\n    return false;\n  }\n\n  /**\n   * Increments the failure counter for the specified partition and checks if the partition can fail over.\n   * This method is used to determine if a partition should be failed over based on the number of request failures.\n   */\n  private async incrementFailureCounterAndCheckFailover(\n    requestContext: RequestContext,\n    isRequestEligibleForPerPartitionAutomaticFailover: boolean,\n    isRequestEligibleForPartitionLevelCircuitBreaker: boolean,\n  ): Promise<boolean> {\n    const partitionKeyRangeId = requestContext.partitionKeyRangeId;\n    const failedEndPoint = requestContext.endpoint;\n    let partitionKeyRangeFailoverInfo: PartitionKeyRangeFailoverInfo;\n\n    if (isRequestEligibleForPerPartitionAutomaticFailover) {\n      if (!this.partitionKeyRangeToLocationForWrite.has(partitionKeyRangeId)) {\n        // If the partition key range is not already in the map, add it\n        const failoverInfo = new PartitionKeyRangeFailoverInfo(failedEndPoint);\n        this.partitionKeyRangeToLocationForWrite.set(partitionKeyRangeId, failoverInfo);\n      }\n      partitionKeyRangeFailoverInfo =\n        this.partitionKeyRangeToLocationForWrite.get(partitionKeyRangeId);\n    } else if (isRequestEligibleForPartitionLevelCircuitBreaker) {\n      if (!this.partitionKeyRangeToLocationForReadAndWrite.has(partitionKeyRangeId)) {\n        // If the partition key range is not already in the map, add it\n        const failoverInfo = new PartitionKeyRangeFailoverInfo(failedEndPoint);\n        this.partitionKeyRangeToLocationForReadAndWrite.set(partitionKeyRangeId, failoverInfo);\n      }\n      partitionKeyRangeFailoverInfo =\n        this.partitionKeyRangeToLocationForReadAndWrite.get(partitionKeyRangeId);\n    } else {\n      return false;\n    }\n\n    assertNotUndefined(\n      partitionKeyRangeFailoverInfo,\n      \"partitionKeyRangeFailoverInfo should be set if failover flags are true.\",\n    );\n\n    const currentTimeInMilliseconds = Date.now();\n    await partitionKeyRangeFailoverInfo.incrementRequestFailureCounts(\n      isReadRequest(requestContext.operationType),\n      currentTimeInMilliseconds,\n    );\n\n    return partitionKeyRangeFailoverInfo.canCircuitBreakerTriggerPartitionFailOver(\n      isReadRequest(requestContext.operationType),\n    );\n  }\n\n  /** Validates if the given request is eligible for partition failover. */\n  private async isRequestEligibleForPartitionFailover(\n    requestContext: RequestContext,\n    shouldValidateFailedLocation: boolean,\n  ): Promise<boolean> {\n    if (\n      !requestContext ||\n      !requestContext.operationType ||\n      !requestContext.resourceType ||\n      !requestContext.partitionKeyRangeId\n    ) {\n      return false;\n    }\n\n    const canUsePartitionLevelFailoverLocations = await this.canUsePartitionLevelFailoverLocations(\n      requestContext.operationType,\n      requestContext.resourceType,\n    );\n    if (!canUsePartitionLevelFailoverLocations) {\n      return false;\n    }\n\n    if (shouldValidateFailedLocation && !requestContext.endpoint) {\n      return false;\n    }\n    return true;\n  }\n\n  /** Determines if partition level failover locations can be used for the given request. */\n  private async canUsePartitionLevelFailoverLocations(\n    operationType?: OperationType,\n    resourceType?: ResourceType,\n  ): Promise<boolean> {\n    const readEndPoints = await this.globalEndpointManager.getReadEndpoints();\n    if (readEndPoints.length <= 1) {\n      return false;\n    }\n    if (\n      resourceType === ResourceType.item ||\n      (resourceType === ResourceType.sproc && operationType === OperationType.Execute)\n    ) {\n      // Right now, for single-master only reads are supported for circuit breaker, and writes are supported for automatic.\n      // failover. For multi master, both reads and writes are supported. Hence return true for both the cases.\n      return true;\n    }\n    return false;\n  }\n\n  /**\n   * Determines if a request is eligible for per-partition automatic failover.\n   * A request is eligible if it is a write request, partition level failover is enabled,\n   * and the global endpoint manager cannot use multiple write locations for the request.\n   */\n  private isRequestEligibleForPerPartitionAutomaticFailover(\n    requestContext: RequestContext,\n  ): boolean {\n    return (\n      this.enablePartitionLevelFailover &&\n      !isReadRequest(requestContext.operationType) &&\n      !this.globalEndpointManager.canUseMultipleWriteLocations(\n        requestContext.resourceType,\n        requestContext.operationType,\n      )\n    );\n  }\n\n  /**\n   * Determines if a request is eligible for partition-level circuit breaker.\n   * This method checks if partition-level circuit breaker is enabled, and if the request is a read-only request or\n   * the global endpoint manager can use multiple write locations for the request.\n   */\n  private isRequestEligibleForPartitionLevelCircuitBreaker(\n    requestContext: RequestContext,\n  ): boolean {\n    if (!this.enablePartitionLevelCircuitBreaker) {\n      return false;\n    }\n    if (isReadRequest(requestContext.operationType)) {\n      return true;\n    }\n    return this.globalEndpointManager.canUseMultipleWriteLocations(\n      requestContext.resourceType,\n      requestContext.operationType,\n    );\n  }\n\n  /**\n   * Attempts to add or update the partition failover information and move to the next available location.\n   * This method checks if the current location for the partition key range has failed and updates the failover\n   * information to route the request to the next available location. If all locations have been tried, it removes\n   * the failover information for the partition key range. Return True if the failover information was successfully\n   * updated and the request was routed to a new location, otherwise false.\n   */\n  private async tryAddOrUpdatePartitionFailoverInfoAndMoveToNextLocation(\n    partitionKeyRangeId: string,\n    failedEndPoint: string,\n    nextEndPoints: readonly string[],\n    partitionKeyRangeToLocation: Map<string, PartitionKeyRangeFailoverInfo>,\n    diagnosticNode: DiagnosticNodeInternal,\n  ): Promise<boolean> {\n    if (!partitionKeyRangeToLocation.has(partitionKeyRangeId)) {\n      // If the partition key range is not already in the map, add it\n      const failoverInfo = new PartitionKeyRangeFailoverInfo(failedEndPoint);\n      partitionKeyRangeToLocation.set(partitionKeyRangeId, failoverInfo);\n    }\n    const partitionFailOver = partitionKeyRangeToLocation.get(partitionKeyRangeId);\n\n    // Will return true if it was able to update to a new region\n    if (\n      await partitionFailOver.tryMoveNextLocation(\n        nextEndPoints,\n        failedEndPoint,\n        diagnosticNode,\n        partitionKeyRangeId,\n      )\n    ) {\n      return true;\n    }\n    // All the locations have been tried. Remove the override information\n    partitionKeyRangeToLocation.delete(partitionKeyRangeId);\n    return false;\n  }\n\n  /**\n   * Initiates a background loop that periodically checks for unhealthy endpoints\n   * and attempts to open connections to them. If a connection is successfully\n   * established, it initiates a failback to the original location for the partition key range.\n   * This is useful for scenarios where a partition key range has been marked as unavailable\n   * due to a circuit breaker, and we want to periodically check if the original location\n   * has become healthy again.\n   * The loop runs at a defined interval specified by Constants.StalePartitionUnavailabilityRefreshIntervalInMs.\n   */\n  private initiateCircuitBreakerFailbackLoop(): void {\n    this.circuitBreakerFailbackBackgroundRefresher = startBackgroundTask(async () => {\n      try {\n        await this.openConnectionToUnhealthyEndpointsWithFailback();\n      } catch (err) {\n        console.error(\"Failed to open connection to unhealthy endpoints: \", err);\n      }\n    }, Constants.StalePartitionUnavailabilityRefreshIntervalInMs);\n  }\n\n  /**\n   * Attempts to open connections to unhealthy endpoints and initiates failback if the connections are successful.\n   * This method checks the partition key ranges that have failed locations and tries to re-establish connections\n   * to those locations. If a connection is successfully re-established, it initiates a failback to the original\n   * location for the partition key range.\n   */\n  private async openConnectionToUnhealthyEndpointsWithFailback(): Promise<void> {\n    for (const pkRange of this.partitionKeyRangeToLocationForReadAndWrite.keys()) {\n      const partitionFailover = this.partitionKeyRangeToLocationForReadAndWrite.get(pkRange);\n      if (!partitionFailover) continue;\n\n      const { firstRequestFailureTime } =\n        await partitionFailover.snapshotPartitionFailoverTimestamps();\n      const now = new Date();\n\n      if (\n        now.getTime() - firstRequestFailureTime >\n        Constants.AllowedPartitionUnavailabilityDurationInMs\n      ) {\n        // Un-deterministically marking the original failed endpoint for the PkRange back to healthy.\n        // Initiate Failback to the original failed location.\n        this.partitionKeyRangeToLocationForReadAndWrite.delete(pkRange);\n      }\n    }\n  }\n}\n"]}

package/dist/browser/queryExecutionContext/Aggregators/index.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../../../../src/queryExecutionContext/Aggregators/index.ts"],"names":[],"mappings":"AAEA,OAAO,EAAE,iBAAiB,EAAE,MAAM,wBAAwB,CAAC;AAC3D,OAAO,EAAE,eAAe,EAAE,MAAM,sBAAsB,CAAC;AACvD,OAAO,EAAE,aAAa,EAAE,MAAM,oBAAoB,CAAC;AACnD,OAAO,EAAE,aAAa,EAAE,MAAM,oBAAoB,CAAC;AACnD,OAAO,EAAE,aAAa,EAAE,MAAM,oBAAoB,CAAC;AACnD,OAAO,EAAE,qBAAqB,EAAE,MAAM,4BAA4B,CAAC;AACnE,OAAO,KAAK,EAAE,aAAa,EAAE,MAAM,gCAAgC,CAAC;AAIpE,wBAAgB,gBAAgB,CAC9B,aAAa,EAAE,aAAa,GAE1B,iBAAiB,GACjB,eAAe,GACf,aAAa,GACb,aAAa,GACb,aAAa,GACb,qBAAqB,CAmBxB;AAED,OAAO,EAAE,iBAAiB,EAAE,eAAe,EAAE,aAAa,EAAE,aAAa,EAAE,aAAa,EAAE,CAAC;AAC3F,OAAO,EAAE,UAAU,EAAE,MAAM,iBAAiB,CAAC"}
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../../../../src/queryExecutionContext/Aggregators/index.ts"],"names":[],"mappings":"AAEA,OAAO,EAAE,iBAAiB,EAAE,MAAM,wBAAwB,CAAC;AAC3D,OAAO,EAAE,eAAe,EAAE,MAAM,sBAAsB,CAAC;AACvD,OAAO,EAAE,aAAa,EAAE,MAAM,oBAAoB,CAAC;AACnD,OAAO,EAAE,aAAa,EAAE,MAAM,oBAAoB,CAAC;AACnD,OAAO,EAAE,aAAa,EAAE,MAAM,oBAAoB,CAAC;AACnD,OAAO,EAAE,qBAAqB,EAAE,MAAM,4BAA4B,CAAC;AACnE,OAAO,KAAK,EAAE,aAAa,EAAE,MAAM,gCAAgC,CAAC;AAIpE,wBAAgB,gBAAgB,CAC9B,aAAa,EAAE,aAAa,GAE1B,iBAAiB,GACjB,eAAe,GACf,aAAa,GACb,aAAa,GACb,aAAa,GACb,qBAAqB,CAoBxB;AAED,OAAO,EAAE,iBAAiB,EAAE,eAAe,EAAE,aAAa,EAAE,aAAa,EAAE,aAAa,EAAE,CAAC;AAC3F,OAAO,EAAE,UAAU,EAAE,MAAM,iBAAiB,CAAC"}

package/dist/browser/queryExecutionContext/Aggregators/index.js

@@ -13,6 +13,7 @@ export function createAggregator(aggregateType) {
         case "Average":
             return new AverageAggregator();
         case "Count":
+        case "CountIf":
             return new CountAggregator();
         case "Max":
             return new MaxAggregator();

package/dist/browser/queryExecutionContext/Aggregators/index.js.map

@@ -1 +1 @@
-{"version":3,"file":"index.js","sourceRoot":"","sources":["../../../../src/queryExecutionContext/Aggregators/index.ts"],"names":[],"mappings":"AAAA,uCAAuC;AACvC,kCAAkC;AAClC,OAAO,EAAE,iBAAiB,EAAE,MAAM,wBAAwB,CAAC;AAC3D,OAAO,EAAE,eAAe,EAAE,MAAM,sBAAsB,CAAC;AACvD,OAAO,EAAE,aAAa,EAAE,MAAM,oBAAoB,CAAC;AACnD,OAAO,EAAE,aAAa,EAAE,MAAM,oBAAoB,CAAC;AACnD,OAAO,EAAE,aAAa,EAAE,MAAM,oBAAoB,CAAC;AACnD,OAAO,EAAE,qBAAqB,EAAE,MAAM,4BAA4B,CAAC;AAEnE,OAAO,EAAE,kBAAkB,EAAE,MAAM,yBAAyB,CAAC;AAC7D,OAAO,EAAE,iBAAiB,EAAE,MAAM,wBAAwB,CAAC;AAE3D,MAAM,UAAU,gBAAgB,CAC9B,aAA4B;IAQ5B,QAAQ,aAAa,EAAE,CAAC;QACtB,KAAK,SAAS;YACZ,OAAO,IAAI,iBAAiB,EAAE,CAAC;QACjC,KAAK,OAAO;YACV,OAAO,IAAI,eAAe,EAAE,CAAC;QAC/B,KAAK,KAAK;YACR,OAAO,IAAI,aAAa,EAAE,CAAC;QAC7B,KAAK,KAAK;YACR,OAAO,IAAI,aAAa,EAAE,CAAC;QAC7B,KAAK,KAAK;YACR,OAAO,IAAI,aAAa,EAAE,CAAC;QAC7B,KAAK,UAAU;YACb,OAAO,IAAI,kBAAkB,EAAE,CAAC;QAClC,KAAK,SAAS;YACZ,OAAO,IAAI,iBAAiB,EAAE,CAAC;QACjC;YACE,OAAO,IAAI,qBAAqB,EAAE,CAAC;IACvC,CAAC;AACH,CAAC;AAED,OAAO,EAAE,iBAAiB,EAAE,eAAe,EAAE,aAAa,EAAE,aAAa,EAAE,aAAa,EAAE,CAAC","sourcesContent":["// Copyright (c) Microsoft Corporation.\n// Licensed under the MIT License.\nimport { AverageAggregator } from \"./AverageAggregator.js\";\nimport { CountAggregator } from \"./CountAggregator.js\";\nimport { MaxAggregator } from \"./MaxAggregator.js\";\nimport { MinAggregator } from \"./MinAggregator.js\";\nimport { SumAggregator } from \"./SumAggregator.js\";\nimport { StaticValueAggregator } from \"./StaticValueAggregator.js\";\nimport type { AggregateType } from \"../../request/ErrorResponse.js\";\nimport { MakeListAggregator } from \"./MakeListAggregator.js\";\nimport { MakeSetAggregator } from \"./MakeSetAggregator.js\";\n\nexport function createAggregator(\n  aggregateType: AggregateType,\n):\n  | AverageAggregator\n  | CountAggregator\n  | MaxAggregator\n  | MinAggregator\n  | SumAggregator\n  | StaticValueAggregator {\n  switch (aggregateType) {\n    case \"Average\":\n      return new AverageAggregator();\n    case \"Count\":\n      return new CountAggregator();\n    case \"Max\":\n      return new MaxAggregator();\n    case \"Min\":\n      return new MinAggregator();\n    case \"Sum\":\n      return new SumAggregator();\n    case \"MakeList\":\n      return new MakeListAggregator();\n    case \"MakeSet\":\n      return new MakeSetAggregator();\n    default:\n      return new StaticValueAggregator();\n  }\n}\n\nexport { AverageAggregator, CountAggregator, MaxAggregator, MinAggregator, SumAggregator };\nexport { Aggregator } from \"./Aggregator.js\";\n"]}
+{"version":3,"file":"index.js","sourceRoot":"","sources":["../../../../src/queryExecutionContext/Aggregators/index.ts"],"names":[],"mappings":"AAAA,uCAAuC;AACvC,kCAAkC;AAClC,OAAO,EAAE,iBAAiB,EAAE,MAAM,wBAAwB,CAAC;AAC3D,OAAO,EAAE,eAAe,EAAE,MAAM,sBAAsB,CAAC;AACvD,OAAO,EAAE,aAAa,EAAE,MAAM,oBAAoB,CAAC;AACnD,OAAO,EAAE,aAAa,EAAE,MAAM,oBAAoB,CAAC;AACnD,OAAO,EAAE,aAAa,EAAE,MAAM,oBAAoB,CAAC;AACnD,OAAO,EAAE,qBAAqB,EAAE,MAAM,4BAA4B,CAAC;AAEnE,OAAO,EAAE,kBAAkB,EAAE,MAAM,yBAAyB,CAAC;AAC7D,OAAO,EAAE,iBAAiB,EAAE,MAAM,wBAAwB,CAAC;AAE3D,MAAM,UAAU,gBAAgB,CAC9B,aAA4B;IAQ5B,QAAQ,aAAa,EAAE,CAAC;QACtB,KAAK,SAAS;YACZ,OAAO,IAAI,iBAAiB,EAAE,CAAC;QACjC,KAAK,OAAO,CAAC;QACb,KAAK,SAAS;YACZ,OAAO,IAAI,eAAe,EAAE,CAAC;QAC/B,KAAK,KAAK;YACR,OAAO,IAAI,aAAa,EAAE,CAAC;QAC7B,KAAK,KAAK;YACR,OAAO,IAAI,aAAa,EAAE,CAAC;QAC7B,KAAK,KAAK;YACR,OAAO,IAAI,aAAa,EAAE,CAAC;QAC7B,KAAK,UAAU;YACb,OAAO,IAAI,kBAAkB,EAAE,CAAC;QAClC,KAAK,SAAS;YACZ,OAAO,IAAI,iBAAiB,EAAE,CAAC;QACjC;YACE,OAAO,IAAI,qBAAqB,EAAE,CAAC;IACvC,CAAC;AACH,CAAC;AAED,OAAO,EAAE,iBAAiB,EAAE,eAAe,EAAE,aAAa,EAAE,aAAa,EAAE,aAAa,EAAE,CAAC","sourcesContent":["// Copyright (c) Microsoft Corporation.\n// Licensed under the MIT License.\nimport { AverageAggregator } from \"./AverageAggregator.js\";\nimport { CountAggregator } from \"./CountAggregator.js\";\nimport { MaxAggregator } from \"./MaxAggregator.js\";\nimport { MinAggregator } from \"./MinAggregator.js\";\nimport { SumAggregator } from \"./SumAggregator.js\";\nimport { StaticValueAggregator } from \"./StaticValueAggregator.js\";\nimport type { AggregateType } from \"../../request/ErrorResponse.js\";\nimport { MakeListAggregator } from \"./MakeListAggregator.js\";\nimport { MakeSetAggregator } from \"./MakeSetAggregator.js\";\n\nexport function createAggregator(\n  aggregateType: AggregateType,\n):\n  | AverageAggregator\n  | CountAggregator\n  | MaxAggregator\n  | MinAggregator\n  | SumAggregator\n  | StaticValueAggregator {\n  switch (aggregateType) {\n    case \"Average\":\n      return new AverageAggregator();\n    case \"Count\":\n    case \"CountIf\":\n      return new CountAggregator();\n    case \"Max\":\n      return new MaxAggregator();\n    case \"Min\":\n      return new MinAggregator();\n    case \"Sum\":\n      return new SumAggregator();\n    case \"MakeList\":\n      return new MakeListAggregator();\n    case \"MakeSet\":\n      return new MakeSetAggregator();\n    default:\n      return new StaticValueAggregator();\n  }\n}\n\nexport { AverageAggregator, CountAggregator, MaxAggregator, MinAggregator, SumAggregator };\nexport { Aggregator } from \"./Aggregator.js\";\n"]}

package/dist/browser/queryExecutionContext/documentProducer.d.ts

@@ -5,6 +5,7 @@ import type { FeedOptions } from "../request/index.js";
 import type { Response } from "../request/index.js";
 import type { FetchFunctionCallback } from "./defaultQueryExecutionContext.js";
 import { FetchResult } from "./FetchResult.js";
+import type { CosmosHeaders } from "./headerUtils.js";
 import type { SqlQuerySpec } from "./index.js";
 /** @hidden */
 export declare class DocumentProducer {
@@ -18,7 +19,6 @@ export declare class DocumentProducer {
     previousContinuationToken: string;
     continuationToken: string;
     generation: number;
-    private respHeaders;
     private internalExecutionContext;
     startEpk: string;
     endEpk: string;
@@ -36,13 +36,12 @@ export declare class DocumentProducer {
     fetchFunction: FetchFunctionCallback;
     hasMoreResults(): boolean;
     gotSplit(): boolean;
-    private _getAndResetActiveResponseHeaders;
     private _updateStates;
     private static _needPartitionKeyRangeCacheRefresh;
     /**
      * Fetches and bufferes the next page of results in internal buffer
      */
-    bufferMore(diagnosticNode: DiagnosticNodeInternal): Promise<void>;
+    bufferMore(diagnosticNode: DiagnosticNodeInternal): Promise<CosmosHeaders>;
     getTargetParitionKeyRange(): PartitionKeyRange;
     /**
      * Peak the next item in the buffer

package/dist/browser/queryExecutionContext/documentProducer.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"documentProducer.d.ts","sourceRoot":"","sources":["../../../src/queryExecutionContext/documentProducer.ts"],"names":[],"mappings":"AAEA,OAAO,KAAK,EAAE,iBAAiB,EAAY,MAAM,oBAAoB,CAAC;AACtE,OAAO,KAAK,EAAE,aAAa,EAAE,MAAM,qBAAqB,CAAC;AASzD,OAAO,KAAK,EAAE,sBAAsB,EAAE,MAAM,0CAA0C,CAAC;AACvF,OAAO,KAAK,EAAE,WAAW,EAAE,MAAM,qBAAqB,CAAC;AACvD,OAAO,KAAK,EAAE,QAAQ,EAAE,MAAM,qBAAqB,CAAC;AAEpD,OAAO,KAAK,EAAE,qBAAqB,EAAE,MAAM,mCAAmC,CAAC;AAC/E,OAAO,EAAE,WAAW,EAAmB,MAAM,kBAAkB,CAAC;AAGhE,OAAO,KAAK,EAAE,YAAY,EAAE,MAAM,YAAY,CAAC;AAE/C,cAAc;AACd,qBAAa,gBAAgB;IAyBzB,OAAO,CAAC,aAAa;IAxBvB,OAAO,CAAC,cAAc,CAAS;IAC/B,OAAO,CAAC,KAAK,CAAwB;IAC9B,uBAAuB,EAAE,iBAAiB,CAAC;IAC3C,YAAY,EAAE,WAAW,EAAE,CAAC;IAC5B,UAAU,EAAE,OAAO,CAAC;IAC3B,OAAO,CAAC,GAAG,CAAQ;IACZ,yBAAyB,EAAE,MAAM,CAAC;IAClC,iBAAiB,EAAE,MAAM,CAAC;IAC1B,UAAU,EAAE,MAAM,CAAK;IAC9B,OAAO,CAAC,WAAW,CAAgB;IACnC,OAAO,CAAC,wBAAwB,CAA+B;IACxD,QAAQ,EAAE,MAAM,CAAC;IACjB,MAAM,EAAE,MAAM,CAAC;IACf,uBAAuB,EAAE,OAAO,CAAC;IAExC;;;;;;;OAOG;gBAEO,aAAa,EAAE,aAAa,EACpC,cAAc,EAAE,MAAM,EACtB,KAAK,EAAE,YAAY,EACnB,uBAAuB,EAAE,iBAAiB,EAC1C,OAAO,EAAE,WAAW,EACpB,oBAAoB,EAAE,MAAM,EAC5B,QAAQ,CAAC,EAAE,MAAM,EACjB,MAAM,CAAC,EAAE,MAAM,EACf,uBAAuB,GAAE,OAAe;IAwBnC,iBAAiB,IAAI,GAAG,EAAE;IAmB1B,aAAa,EAAE,qBAAqB,CAwBzC;IAEK,cAAc,IAAI,OAAO;IAIzB,QAAQ,IAAI,OAAO;IAY1B,OAAO,CAAC,iCAAiC;IAMzC,OAAO,CAAC,aAAa;IAgBrB,OAAO,CAAC,MAAM,CAAC,kCAAkC;IASjD;;OAEG;IACU,UAAU,CAAC,cAAc,EAAE,sBAAsB,GAAG,OAAO,CAAC,IAAI,CAAC;IAkDvE,yBAAyB,IAAI,iBAAiB;IAGrD;;OAEG;IACI,YAAY,IAAI,GAAG;IAqB1B;;OAEG;IACU,aAAa,IAAI,OAAO,CAAC,QAAQ,CAAC,GAAG,CAAC,CAAC;IAoBpD;;OAEG;IACU,kBAAkB,IAAI,OAAO,CAAC,QAAQ,CAAC,GAAG,EAAE,CAAC,CAAC;IA4B3D;;OAEG;IACH,OAAO,CAAC,OAAO;CAiChB"}
+{"version":3,"file":"documentProducer.d.ts","sourceRoot":"","sources":["../../../src/queryExecutionContext/documentProducer.ts"],"names":[],"mappings":"AAEA,OAAO,KAAK,EAAE,iBAAiB,EAAY,MAAM,oBAAoB,CAAC;AACtE,OAAO,KAAK,EAAE,aAAa,EAAE,MAAM,qBAAqB,CAAC;AASzD,OAAO,KAAK,EAAE,sBAAsB,EAAE,MAAM,0CAA0C,CAAC;AACvF,OAAO,KAAK,EAAE,WAAW,EAAE,MAAM,qBAAqB,CAAC;AACvD,OAAO,KAAK,EAAE,QAAQ,EAAE,MAAM,qBAAqB,CAAC;AAEpD,OAAO,KAAK,EAAE,qBAAqB,EAAE,MAAM,mCAAmC,CAAC;AAC/E,OAAO,EAAE,WAAW,EAAmB,MAAM,kBAAkB,CAAC;AAEhE,OAAO,KAAK,EAAE,aAAa,EAAE,MAAM,kBAAkB,CAAC;AACtD,OAAO,KAAK,EAAE,YAAY,EAAE,MAAM,YAAY,CAAC;AAE/C,cAAc;AACd,qBAAa,gBAAgB;IAwBzB,OAAO,CAAC,aAAa;IAvBvB,OAAO,CAAC,cAAc,CAAS;IAC/B,OAAO,CAAC,KAAK,CAAwB;IAC9B,uBAAuB,EAAE,iBAAiB,CAAC;IAC3C,YAAY,EAAE,WAAW,EAAE,CAAC;IAC5B,UAAU,EAAE,OAAO,CAAC;IAC3B,OAAO,CAAC,GAAG,CAAQ;IACZ,yBAAyB,EAAE,MAAM,CAAC;IAClC,iBAAiB,EAAE,MAAM,CAAC;IAC1B,UAAU,EAAE,MAAM,CAAK;IAC9B,OAAO,CAAC,wBAAwB,CAA+B;IACxD,QAAQ,EAAE,MAAM,CAAC;IACjB,MAAM,EAAE,MAAM,CAAC;IACf,uBAAuB,EAAE,OAAO,CAAC;IAExC;;;;;;;OAOG;gBAEO,aAAa,EAAE,aAAa,EACpC,cAAc,EAAE,MAAM,EACtB,KAAK,EAAE,YAAY,EACnB,uBAAuB,EAAE,iBAAiB,EAC1C,OAAO,EAAE,WAAW,EACpB,oBAAoB,EAAE,MAAM,EAC5B,QAAQ,CAAC,EAAE,MAAM,EACjB,MAAM,CAAC,EAAE,MAAM,EACf,uBAAuB,GAAE,OAAe;IAuBnC,iBAAiB,IAAI,GAAG,EAAE;IAmB1B,aAAa,EAAE,qBAAqB,CAwBzC;IAEK,cAAc,IAAI,OAAO;IAIzB,QAAQ,IAAI,OAAO;IAY1B,OAAO,CAAC,aAAa;IAgBrB,OAAO,CAAC,MAAM,CAAC,kCAAkC;IASjD;;OAEG;IACU,UAAU,CAAC,cAAc,EAAE,sBAAsB,GAAG,OAAO,CAAC,aAAa,CAAC;IAkDhF,yBAAyB,IAAI,iBAAiB;IAGrD;;OAEG;IACI,YAAY,IAAI,GAAG;IAqB1B;;OAEG;IACU,aAAa,IAAI,OAAO,CAAC,QAAQ,CAAC,GAAG,CAAC,CAAC;IAoBpD;;OAEG;IACU,kBAAkB,IAAI,OAAO,CAAC,QAAQ,CAAC,GAAG,EAAE,CAAC,CAAC;IA6B3D;;OAEG;IACH,OAAO,CAAC,OAAO;CAiChB"}

package/dist/browser/queryExecutionContext/documentProducer.js

@@ -1,7 +1,7 @@
 import { Constants, getIdFromLink, getPathFromLink, ResourceType, StatusCodes, SubStatusCodes, } from "../common/index.js";
 import { DefaultQueryExecutionContext } from "./defaultQueryExecutionContext.js";
 import { FetchResult, FetchResultType } from "./FetchResult.js";
-import { getInitialHeader, mergeHeaders } from "./headerUtils.js";
+import { getInitialHeader } from "./headerUtils.js";
 /** @hidden */
 export class DocumentProducer {
     /**
@@ -44,7 +44,6 @@ export class DocumentProducer {
         this.err = undefined;
         this.previousContinuationToken = undefined;
         this.continuationToken = undefined;
-        this.respHeaders = getInitialHeader();
         this.internalExecutionContext = new DefaultQueryExecutionContext(options, this.fetchFunction, correlatedActivityId);
         this.startEpk = startEpk;
         this.endEpk = endEpk;
@@ -82,11 +81,6 @@ export class DocumentProducer {
         }
         return false;
     }
-    _getAndResetActiveResponseHeaders() {
-        const ret = this.respHeaders;
-        this.respHeaders = getInitialHeader();
-        return ret;
-    }
     _updateStates(err, allFetched) {
         if (err) {
             this.err = err;
@@ -136,7 +130,7 @@ export class DocumentProducer {
                 headerResponse[Constants.HttpHeaders.QueryMetrics][this.targetPartitionKeyRange.id] =
                     queryMetrics;
             }
-            mergeHeaders(this.respHeaders, headerResponse);
+            return headerResponse;
         }
         catch (err) {
             if (DocumentProducer._needPartitionKeyRangeCacheRefresh(err)) {
@@ -144,7 +138,7 @@ export class DocumentProducer {
                 // Buffer the error so the execution context can still get the feedResponses in the itemBuffer
                 const bufferedError = new FetchResult(undefined, err);
                 this.fetchResults.push(bufferedError);
-                mergeHeaders(this.respHeaders, err.headers);
+                return err.headers;
             }
             else {
                 this._updateStates(err, err.resources === undefined);
@@ -184,7 +178,7 @@ export class DocumentProducer {
             throw this.err;
         }
         if (this.allFetched) {
-            return { result: undefined, headers: this._getAndResetActiveResponseHeaders() };
+            return { result: undefined, headers: getInitialHeader() };
         }
         try {
             const { result, headers } = this.current();
@@ -208,23 +202,24 @@ export class DocumentProducer {
             throw this.err;
         }
         if (this.allFetched) {
-            return { result: undefined, headers: this._getAndResetActiveResponseHeaders() };
+            return { result: undefined, headers: getInitialHeader() };
         }
         const resources = [];
-        const resHeaders = getInitialHeader();
         try {
             while (this.fetchResults.length > 0) {
-                const { result, headers } = this.current();
+                const { result } = this.current();
                 this._updateStates(undefined, result === undefined);
-                mergeHeaders(resHeaders, headers);
                 if (result === undefined) {
-                    return { result: resources.length > 0 ? resources : undefined, headers: resHeaders };
+                    return {
+                        result: resources.length > 0 ? resources : undefined,
+                        headers: getInitialHeader(),
+                    };
                 }
                 else {
                     resources.push(result);
                 }
             }
-            return { result: resources, headers: resHeaders };
+            return { result: resources, headers: getInitialHeader() };
         }
         catch (err) {
             this._updateStates(err, err.item === undefined);
@@ -243,15 +238,15 @@ export class DocumentProducer {
                 case FetchResultType.Done:
                     return {
                         result: undefined,
-                        headers: this._getAndResetActiveResponseHeaders(),
+                        headers: getInitialHeader(),
                     };
                 case FetchResultType.Exception:
-                    fetchResult.error.headers = this._getAndResetActiveResponseHeaders();
+                    fetchResult.error.headers = getInitialHeader();
                     throw fetchResult.error;
                 case FetchResultType.Result:
                     return {
                         result: fetchResult.feedResponse,
-                        headers: this._getAndResetActiveResponseHeaders(),
+                        headers: getInitialHeader(),
                     };
             }
         }
@@ -259,11 +254,11 @@ export class DocumentProducer {
         if (this.allFetched) {
             return {
                 result: undefined,
-                headers: this._getAndResetActiveResponseHeaders(),
+                headers: getInitialHeader(),
             };
         }
         // If the internal buffer is empty, return empty result
-        return { result: [], headers: this._getAndResetActiveResponseHeaders() };
+        return { result: [], headers: getInitialHeader() };
     }
 }
 //# sourceMappingURL=documentProducer.js.map

package/dist/browser/queryExecutionContext/documentProducer.js.map

@@ -1 +1 @@
-{"version":3,"file":"documentProducer.js","sourceRoot":"","sources":["../../../src/queryExecutionContext/documentProducer.ts"],"names":[],"mappings":"AAIA,OAAO,EACL,SAAS,EACT,aAAa,EACb,eAAe,EACf,YAAY,EACZ,WAAW,EACX,cAAc,GACf,MAAM,oBAAoB,CAAC;AAI5B,OAAO,EAAE,4BAA4B,EAAE,MAAM,mCAAmC,CAAC;AAEjF,OAAO,EAAE,WAAW,EAAE,eAAe,EAAE,MAAM,kBAAkB,CAAC;AAChE,OAAO,EAAE,gBAAgB,EAAE,YAAY,EAAE,MAAM,kBAAkB,CAAC;AAIlE,cAAc;AACd,MAAM,OAAO,gBAAgB;IAgB3B;;;;;;;OAOG;IACH,YACU,aAA4B,EACpC,cAAsB,EACtB,KAAmB,EACnB,uBAA0C,EAC1C,OAAoB,EACpB,oBAA4B,EAC5B,QAAiB,EACjB,MAAe,EACf,0BAAmC,KAAK;QARhC,kBAAa,GAAb,aAAa,CAAe;QAhB/B,eAAU,GAAW,CAAC,CAAC;QAmEvB,kBAAa,GAA0B,KAAK,EACjD,cAAsC,EACtC,OAAoB,EACpB,oBAA4B,EACC,EAAE;YAC/B,MAAM,IAAI,GAAG,eAAe,CAAC,IAAI,CAAC,cAAc,EAAE,YAAY,CAAC,IAAI,CAAC,CAAC;YACrE,cAAc,CAAC,OAAO,CAAC,EAAE,mBAAmB,EAAE,IAAI,CAAC,uBAAuB,CAAC,EAAE,EAAE,CAAC,CAAC;YACjF,MAAM,EAAE,GAAG,aAAa,CAAC,IAAI,CAAC,cAAc,CAAC,CAAC;YAC9C,MAAM,QAAQ,GAAG,IAAI,CAAC,uBAAuB,CAAC,CAAC,CAAC,IAAI,CAAC,QAAQ,CAAC,CAAC,CAAC,SAAS,CAAC;YAC1E,MAAM,MAAM,GAAG,IAAI,CAAC,uBAAuB,CAAC,CAAC,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC,CAAC,SAAS,CAAC;YAEtE,OAAO,IAAI,CAAC,aAAa,CAAC,SAAS,CAAC;gBAClC,IAAI;gBACJ,YAAY,EAAE,YAAY,CAAC,IAAI;gBAC/B,UAAU,EAAE,EAAE;gBACd,QAAQ,EAAE,CAAC,MAAW,EAAE,EAAE,CAAC,MAAM,CAAC,SAAS;gBAC3C,KAAK,EAAE,IAAI,CAAC,KAAK;gBACjB,OAAO;gBACP,cAAc;gBACd,mBAAmB,EAAE,IAAI,CAAC,uBAAuB,CAAC,IAAI,CAAC;gBACvD,oBAAoB,EAAE,oBAAoB;gBAC1C,QAAQ,EAAE,QAAQ;gBAClB,MAAM,EAAE,MAAM;aACf,CAAC,CAAC;QACL,CAAC,CAAC;QAjEA,oBAAoB;QACpB,IAAI,CAAC,cAAc,GAAG,cAAc,CAAC;QACrC,IAAI,CAAC,KAAK,GAAG,KAAK,CAAC;QACnB,IAAI,CAAC,uBAAuB,GAAG,uBAAuB,CAAC;QACvD,IAAI,CAAC,YAAY,GAAG,EAAE,CAAC;QAEvB,IAAI,CAAC,UAAU,GAAG,KAAK,CAAC;QACxB,IAAI,CAAC,GAAG,GAAG,SAAS,CAAC;QAErB,IAAI,CAAC,yBAAyB,GAAG,SAAS,CAAC;QAC3C,IAAI,CAAC,iBAAiB,GAAG,SAAS,CAAC;QACnC,IAAI,CAAC,WAAW,GAAG,gBAAgB,EAAE,CAAC;QAEtC,IAAI,CAAC,wBAAwB,GAAG,IAAI,4BAA4B,CAC9D,OAAO,EACP,IAAI,CAAC,aAAa,EAClB,oBAAoB,CACrB,CAAC;QACF,IAAI,CAAC,QAAQ,GAAG,QAAQ,CAAC;QACzB,IAAI,CAAC,MAAM,GAAG,MAAM,CAAC;QACrB,IAAI,CAAC,uBAAuB,GAAG,uBAAuB,CAAC;IACzD,CAAC;IACM,iBAAiB;QACtB,MAAM,eAAe,GAAG,EAAE,CAAC;QAC3B,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,IAAI,GAAG,KAAK,EAAE,CAAC,GAAG,IAAI,CAAC,YAAY,CAAC,MAAM,IAAI,CAAC,IAAI,EAAE,CAAC,EAAE,EAAE,CAAC;YACzE,MAAM,WAAW,GAAG,IAAI,CAAC,YAAY,CAAC,CAAC,CAAC,CAAC;YACzC,QAAQ,WAAW,CAAC,eAAe,EAAE,CAAC;gBACpC,KAAK,eAAe,CAAC,IAAI;oBACvB,IAAI,GAAG,IAAI,CAAC;oBACZ,MAAM;gBACR,KAAK,eAAe,CAAC,SAAS;oBAC5B,IAAI,GAAG,IAAI,CAAC;oBACZ,MAAM;gBACR,KAAK,eAAe,CAAC,MAAM;oBACzB,eAAe,CAAC,IAAI,CAAC,WAAW,CAAC,YAAY,CAAC,CAAC;oBAC/C,MAAM;YACV,CAAC;QACH,CAAC;QACD,OAAO,eAAe,CAAC;IACzB,CAAC;IA4BM,cAAc;QACnB,OAAO,IAAI,CAAC,wBAAwB,CAAC,cAAc,EAAE,IAAI,IAAI,CAAC,YAAY,CAAC,MAAM,KAAK,CAAC,CAAC;IAC1F,CAAC;IAEM,QAAQ;QACb,IAAI,IAAI,CAAC,YAAY,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;YACnC,MAAM,WAAW,GAAG,IAAI,CAAC,YAAY,CAAC,CAAC,CAAC,CAAC;YACzC,IAAI,WAAW,CAAC,eAAe,KAAK,eAAe,CAAC,SAAS,EAAE,CAAC;gBAC9D,IAAI,gBAAgB,CAAC,kCAAkC,CAAC,WAAW,CAAC,KAAK,CAAC,EAAE,CAAC;oBAC3E,OAAO,IAAI,CAAC;gBACd,CAAC;YACH,CAAC;QACH,CAAC;QACD,OAAO,KAAK,CAAC;IACf,CAAC;IAEO,iCAAiC;QACvC,MAAM,GAAG,GAAG,IAAI,CAAC,WAAW,CAAC;QAC7B,IAAI,CAAC,WAAW,GAAG,gBAAgB,EAAE,CAAC;QACtC,OAAO,GAAG,CAAC;IACb,CAAC;IAEO,aAAa,CAAC,GAAQ,EAAE,UAAmB;QACjD,IAAI,GAAG,EAAE,CAAC;YACR,IAAI,CAAC,GAAG,GAAG,GAAG,CAAC;YACf,OAAO;QACT,CAAC;QACD,IAAI,UAAU,EAAE,CAAC;YACf,IAAI,CAAC,UAAU,GAAG,IAAI,CAAC;QACzB,CAAC;QACD,IAAI,IAAI,CAAC,wBAAwB,CAAC,iBAAiB,KAAK,IAAI,CAAC,iBAAiB,EAAE,CAAC;YAC/E,kBAAkB;YAClB,OAAO;QACT,CAAC;QACD,IAAI,CAAC,yBAAyB,GAAG,IAAI,CAAC,iBAAiB,CAAC;QACxD,IAAI,CAAC,iBAAiB,GAAG,IAAI,CAAC,wBAAwB,CAAC,iBAAiB,CAAC;IAC3E,CAAC;IAEO,MAAM,CAAC,kCAAkC,CAAC,KAAU;QAC1D,cAAc;QACd,OAAO,CACL,KAAK,CAAC,IAAI,KAAK,WAAW,CAAC,IAAI;YAC/B,WAAW,IAAI,KAAK;YACpB,KAAK,CAAC,WAAW,CAAC,KAAK,cAAc,CAAC,qBAAqB,CAC5D,CAAC;IACJ,CAAC;IAED;;OAEG;IACI,KAAK,CAAC,UAAU,CAAC,cAAsC;QAC5D,IAAI,IAAI,CAAC,GAAG,EAAE,CAAC;YACb,MAAM,IAAI,CAAC,GAAG,CAAC;QACjB,CAAC;QAED,IAAI,CAAC;YACH,MAAM,EAAE,MAAM,EAAE,SAAS,EAAE,OAAO,EAAE,cAAc,EAAE,GAClD,MAAM,IAAI,CAAC,wBAAwB,CAAC,SAAS,CAAC,cAAc,CAAC,CAAC;YAChE,EAAE,IAAI,CAAC,UAAU,CAAC;YAClB,IAAI,CAAC,aAAa,CAAC,SAAS,EAAE,SAAS,KAAK,SAAS,CAAC,CAAC;YACvD,IAAI,SAAS,KAAK,SAAS,EAAE,CAAC;gBAC5B,sDAAsD;gBACtD,IAAI,sBAAsB,GAAG,IAAI,CAAC;gBAClC,SAAS,CAAC,OAAO,CAAC,CAAC,OAAY,EAAE,EAAE;oBACjC,IAAI,CAAC,YAAY,CAAC,IAAI,CACpB,IAAI,WAAW,CACb,OAAO,EACP,SAAS,EACT,sBAAsB,CAAC,CAAC,CAAC,cAAc,CAAC,CAAC,CAAC,gBAAgB,EAAE,CAC7D,CACF,CAAC;oBACF,sBAAsB,GAAG,KAAK,CAAC;gBACjC,CAAC,CAAC,CAAC;YACL,CAAC;YAED,iFAAiF;YACjF,IAAI,cAAc,IAAI,IAAI,IAAI,SAAS,CAAC,WAAW,CAAC,YAAY,IAAI,cAAc,EAAE,CAAC;gBACnF,gEAAgE;gBAChE,MAAM,YAAY,GAAG,cAAc,CAAC,SAAS,CAAC,WAAW,CAAC,YAAY,CAAC,CAAC,GAAG,CAAC,CAAC;gBAE7E,gFAAgF;gBAChF,cAAc,CAAC,SAAS,CAAC,WAAW,CAAC,YAAY,CAAC,GAAG,EAAE,CAAC;gBACxD,cAAc,CAAC,SAAS,CAAC,WAAW,CAAC,YAAY,CAAC,CAAC,IAAI,CAAC,uBAAuB,CAAC,EAAE,CAAC;oBACjF,YAAY,CAAC;YACjB,CAAC;YACD,YAAY,CAAC,IAAI,CAAC,WAAW,EAAE,cAAc,CAAC,CAAC;QACjD,CAAC;QAAC,OAAO,GAAQ,EAAE,CAAC;YAClB,IAAI,gBAAgB,CAAC,kCAAkC,CAAC,GAAG,CAAC,EAAE,CAAC;gBAC7D,qBAAqB;gBACrB,8FAA8F;gBAC9F,MAAM,aAAa,GAAG,IAAI,WAAW,CAAC,SAAS,EAAE,GAAG,CAAC,CAAC;gBACtD,IAAI,CAAC,YAAY,CAAC,IAAI,CAAC,aAAa,CAAC,CAAC;gBACtC,YAAY,CAAC,IAAI,CAAC,WAAW,EAAE,GAAG,CAAC,OAAO,CAAC,CAAC;YAC9C,CAAC;iBAAM,CAAC;gBACN,IAAI,CAAC,aAAa,CAAC,GAAG,EAAE,GAAG,CAAC,SAAS,KAAK,SAAS,CAAC,CAAC;gBACrD,MAAM,GAAG,CAAC;YACZ,CAAC;QACH,CAAC;IACH,CAAC;IAEM,yBAAyB;QAC9B,OAAO,IAAI,CAAC,uBAAuB,CAAC;IACtC,CAAC;IACD;;OAEG;IACI,YAAY;QACjB,IAAI,IAAI,CAAC,GAAG,EAAE,CAAC;YACb,MAAM,IAAI,CAAC,GAAG,CAAC;QACjB,CAAC;QACD,IAAI,IAAI,CAAC,UAAU,IAAI,IAAI,CAAC,YAAY,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;YACtD,OAAO,SAAS,CAAC;QACnB,CAAC;QACD,MAAM,WAAW,GAAG,IAAI,CAAC,YAAY,CAAC,CAAC,CAAC,CAAC;QAEzC,QAAQ,WAAW,CAAC,eAAe,EAAE,CAAC;YACpC,KAAK,eAAe,CAAC,IAAI;gBACvB,OAAO,SAAS,CAAC;YAEnB,KAAK,eAAe,CAAC,SAAS,EAAE,0BAA0B;gBACxD,OAAO,SAAS,CAAC;YAEnB,KAAK,eAAe,CAAC,MAAM;gBACzB,OAAO,WAAW,CAAC,YAAY,CAAC;QACpC,CAAC;IACH,CAAC;IAED;;OAEG;IACI,KAAK,CAAC,aAAa;QACxB,IAAI,IAAI,CAAC,GAAG,EAAE,CAAC;YACb,IAAI,CAAC,aAAa,CAAC,IAAI,CAAC,GAAG,EAAE,SAAS,CAAC,CAAC;YACxC,MAAM,IAAI,CAAC,GAAG,CAAC;QACjB,CAAC;QACD,IAAI,IAAI,CAAC,UAAU,EAAE,CAAC;YACpB,OAAO,EAAE,MAAM,EAAE,SAAS,EAAE,OAAO,EAAE,IAAI,CAAC,iCAAiC,EAAE,EAAE,CAAC;QAClF,CAAC;QACD,IAAI,CAAC;YACH,MAAM,EAAE,MAAM,EAAE,OAAO,EAAE,GAAG,IAAI,CAAC,OAAO,EAAE,CAAC;YAC3C,IAAI,CAAC,aAAa,CAAC,SAAS,EAAE,MAAM,KAAK,SAAS,CAAC,CAAC;YACpD,IAAI,MAAM,KAAK,SAAS,IAAI,MAAM,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;gBAChD,OAAO,EAAE,MAAM,EAAE,SAAS,EAAE,OAAO,EAAE,CAAC;YACxC,CAAC;YACD,OAAO,EAAE,MAAM,EAAE,OAAO,EAAE,CAAC,CAAC,EAAE;QAChC,CAAC;QAAC,OAAO,GAAQ,EAAE,CAAC;YAClB,IAAI,CAAC,aAAa,CAAC,GAAG,EAAE,GAAG,CAAC,IAAI,KAAK,SAAS,CAAC,CAAC;YAChD,MAAM,GAAG,CAAC;QACZ,CAAC;IACH,CAAC;IACD;;OAEG;IACI,KAAK,CAAC,kBAAkB;QAC7B,IAAI,IAAI,CAAC,GAAG,EAAE,CAAC;YACb,IAAI,CAAC,aAAa,CAAC,IAAI,CAAC,GAAG,EAAE,SAAS,CAAC,CAAC;YACxC,MAAM,IAAI,CAAC,GAAG,CAAC;QACjB,CAAC;QACD,IAAI,IAAI,CAAC,UAAU,EAAE,CAAC;YACpB,OAAO,EAAE,MAAM,EAAE,SAAS,EAAE,OAAO,EAAE,IAAI,CAAC,iCAAiC,EAAE,EAAE,CAAC;QAClF,CAAC;QACD,MAAM,SAAS,GAAU,EAAE,CAAC;QAC5B,MAAM,UAAU,GAAkB,gBAAgB,EAAE,CAAC;QACrD,IAAI,CAAC;YACH,OAAO,IAAI,CAAC,YAAY,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC;gBACpC,MAAM,EAAE,MAAM,EAAE,OAAO,EAAE,GAAG,IAAI,CAAC,OAAO,EAAE,CAAC;gBAC3C,IAAI,CAAC,aAAa,CAAC,SAAS,EAAE,MAAM,KAAK,SAAS,CAAC,CAAC;gBACpD,YAAY,CAAC,UAAU,EAAE,OAAO,CAAC,CAAC;gBAClC,IAAI,MAAM,KAAK,SAAS,EAAE,CAAC;oBACzB,OAAO,EAAE,MAAM,EAAE,SAAS,CAAC,MAAM,GAAG,CAAC,CAAC,CAAC,CAAC,SAAS,CAAC,CAAC,CAAC,SAAS,EAAE,OAAO,EAAE,UAAU,EAAE,CAAC;gBACvF,CAAC;qBAAM,CAAC;oBACN,SAAS,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;gBACzB,CAAC;YACH,CAAC;YACD,OAAO,EAAE,MAAM,EAAE,SAAS,EAAE,OAAO,EAAE,UAAU,EAAE,CAAC;QACpD,CAAC;QAAC,OAAO,GAAQ,EAAE,CAAC;YAClB,IAAI,CAAC,aAAa,CAAC,GAAG,EAAE,GAAG,CAAC,IAAI,KAAK,SAAS,CAAC,CAAC;YAChD,MAAM,GAAG,CAAC;QACZ,CAAC;IACH,CAAC;IAED;;OAEG;IACK,OAAO;QACb,0CAA0C;QAC1C,IAAI,IAAI,CAAC,YAAY,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC;YACjC,MAAM,WAAW,GAAG,IAAI,CAAC,YAAY,CAAC,KAAK,EAAE,CAAC;YAC9C,+BAA+B;YAC/B,QAAQ,WAAW,CAAC,eAAe,EAAE,CAAC;gBACpC,KAAK,eAAe,CAAC,IAAI;oBACvB,OAAO;wBACL,MAAM,EAAE,SAAS;wBACjB,OAAO,EAAE,IAAI,CAAC,iCAAiC,EAAE;qBAClD,CAAC;gBACJ,KAAK,eAAe,CAAC,SAAS;oBAC5B,WAAW,CAAC,KAAK,CAAC,OAAO,GAAG,IAAI,CAAC,iCAAiC,EAAE,CAAC;oBACrE,MAAM,WAAW,CAAC,KAAK,CAAC;gBAC1B,KAAK,eAAe,CAAC,MAAM;oBACzB,OAAO;wBACL,MAAM,EAAE,WAAW,CAAC,YAAY;wBAChC,OAAO,EAAE,IAAI,CAAC,iCAAiC,EAAE;qBAClD,CAAC;YACN,CAAC;QACH,CAAC;QAED,qEAAqE;QACrE,IAAI,IAAI,CAAC,UAAU,EAAE,CAAC;YACpB,OAAO;gBACL,MAAM,EAAE,SAAS;gBACjB,OAAO,EAAE,IAAI,CAAC,iCAAiC,EAAE;aAClD,CAAC;QACJ,CAAC;QAED,uDAAuD;QACvD,OAAO,EAAE,MAAM,EAAE,EAAE,EAAE,OAAO,EAAE,IAAI,CAAC,iCAAiC,EAAE,EAAE,CAAC;IAC3E,CAAC;CACF","sourcesContent":["// Copyright (c) Microsoft Corporation.\n// Licensed under the MIT License.\nimport type { PartitionKeyRange, Resource } from \"../client/index.js\";\nimport type { ClientContext } from \"../ClientContext.js\";\nimport {\n  Constants,\n  getIdFromLink,\n  getPathFromLink,\n  ResourceType,\n  StatusCodes,\n  SubStatusCodes,\n} from \"../common/index.js\";\nimport type { DiagnosticNodeInternal } from \"../diagnostics/DiagnosticNodeInternal.js\";\nimport type { FeedOptions } from \"../request/index.js\";\nimport type { Response } from \"../request/index.js\";\nimport { DefaultQueryExecutionContext } from \"./defaultQueryExecutionContext.js\";\nimport type { FetchFunctionCallback } from \"./defaultQueryExecutionContext.js\";\nimport { FetchResult, FetchResultType } from \"./FetchResult.js\";\nimport { getInitialHeader, mergeHeaders } from \"./headerUtils.js\";\nimport type { CosmosHeaders } from \"./headerUtils.js\";\nimport type { SqlQuerySpec } from \"./index.js\";\n\n/** @hidden */\nexport class DocumentProducer {\n  private collectionLink: string;\n  private query: string | SqlQuerySpec;\n  public targetPartitionKeyRange: PartitionKeyRange;\n  public fetchResults: FetchResult[];\n  public allFetched: boolean;\n  private err: Error;\n  public previousContinuationToken: string;\n  public continuationToken: string;\n  public generation: number = 0;\n  private respHeaders: CosmosHeaders;\n  private internalExecutionContext: DefaultQueryExecutionContext;\n  public startEpk: string;\n  public endEpk: string;\n  public populateEpkRangeHeaders: boolean;\n\n  /**\n   * Provides the Target Partition Range Query Execution Context.\n   * @param clientContext  - The service endpoint to use to create the client.\n   * @param collectionLink - Represents collection link\n   * @param query          - A SQL query.\n   * @param targetPartitionKeyRange - Query Target Partition key Range\n   * @hidden\n   */\n  constructor(\n    private clientContext: ClientContext,\n    collectionLink: string,\n    query: SqlQuerySpec,\n    targetPartitionKeyRange: PartitionKeyRange,\n    options: FeedOptions,\n    correlatedActivityId: string,\n    startEpk?: string,\n    endEpk?: string,\n    populateEpkRangeHeaders: boolean = false,\n  ) {\n    // TODO: any options\n    this.collectionLink = collectionLink;\n    this.query = query;\n    this.targetPartitionKeyRange = targetPartitionKeyRange;\n    this.fetchResults = [];\n\n    this.allFetched = false;\n    this.err = undefined;\n\n    this.previousContinuationToken = undefined;\n    this.continuationToken = undefined;\n    this.respHeaders = getInitialHeader();\n\n    this.internalExecutionContext = new DefaultQueryExecutionContext(\n      options,\n      this.fetchFunction,\n      correlatedActivityId,\n    );\n    this.startEpk = startEpk;\n    this.endEpk = endEpk;\n    this.populateEpkRangeHeaders = populateEpkRangeHeaders;\n  }\n  public peekBufferedItems(): any[] {\n    const bufferedResults = [];\n    for (let i = 0, done = false; i < this.fetchResults.length && !done; i++) {\n      const fetchResult = this.fetchResults[i];\n      switch (fetchResult.fetchResultType) {\n        case FetchResultType.Done:\n          done = true;\n          break;\n        case FetchResultType.Exception:\n          done = true;\n          break;\n        case FetchResultType.Result:\n          bufferedResults.push(fetchResult.feedResponse);\n          break;\n      }\n    }\n    return bufferedResults;\n  }\n\n  public fetchFunction: FetchFunctionCallback = async (\n    diagnosticNode: DiagnosticNodeInternal,\n    options: FeedOptions,\n    correlatedActivityId: string,\n  ): Promise<Response<Resource>> => {\n    const path = getPathFromLink(this.collectionLink, ResourceType.item);\n    diagnosticNode.addData({ partitionKeyRangeId: this.targetPartitionKeyRange.id });\n    const id = getIdFromLink(this.collectionLink);\n    const startEpk = this.populateEpkRangeHeaders ? this.startEpk : undefined;\n    const endEpk = this.populateEpkRangeHeaders ? this.endEpk : undefined;\n\n    return this.clientContext.queryFeed({\n      path,\n      resourceType: ResourceType.item,\n      resourceId: id,\n      resultFn: (result: any) => result.Documents,\n      query: this.query,\n      options,\n      diagnosticNode,\n      partitionKeyRangeId: this.targetPartitionKeyRange[\"id\"],\n      correlatedActivityId: correlatedActivityId,\n      startEpk: startEpk,\n      endEpk: endEpk,\n    });\n  };\n\n  public hasMoreResults(): boolean {\n    return this.internalExecutionContext.hasMoreResults() || this.fetchResults.length !== 0;\n  }\n\n  public gotSplit(): boolean {\n    if (this.fetchResults.length !== 0) {\n      const fetchResult = this.fetchResults[0];\n      if (fetchResult.fetchResultType === FetchResultType.Exception) {\n        if (DocumentProducer._needPartitionKeyRangeCacheRefresh(fetchResult.error)) {\n          return true;\n        }\n      }\n    }\n    return false;\n  }\n\n  private _getAndResetActiveResponseHeaders(): CosmosHeaders {\n    const ret = this.respHeaders;\n    this.respHeaders = getInitialHeader();\n    return ret;\n  }\n\n  private _updateStates(err: any, allFetched: boolean): void {\n    if (err) {\n      this.err = err;\n      return;\n    }\n    if (allFetched) {\n      this.allFetched = true;\n    }\n    if (this.internalExecutionContext.continuationToken === this.continuationToken) {\n      // nothing changed\n      return;\n    }\n    this.previousContinuationToken = this.continuationToken;\n    this.continuationToken = this.internalExecutionContext.continuationToken;\n  }\n\n  private static _needPartitionKeyRangeCacheRefresh(error: any): boolean {\n    // TODO: error\n    return (\n      error.code === StatusCodes.Gone &&\n      \"substatus\" in error &&\n      error[\"substatus\"] === SubStatusCodes.PartitionKeyRangeGone\n    );\n  }\n\n  /**\n   * Fetches and bufferes the next page of results in internal buffer\n   */\n  public async bufferMore(diagnosticNode: DiagnosticNodeInternal): Promise<void> {\n    if (this.err) {\n      throw this.err;\n    }\n\n    try {\n      const { result: resources, headers: headerResponse } =\n        await this.internalExecutionContext.fetchMore(diagnosticNode);\n      ++this.generation;\n      this._updateStates(undefined, resources === undefined);\n      if (resources !== undefined) {\n        // add fetched header to the 1st element in the buffer\n        let addHeaderToFetchResult = true;\n        resources.forEach((element: any) => {\n          this.fetchResults.push(\n            new FetchResult(\n              element,\n              undefined,\n              addHeaderToFetchResult ? headerResponse : getInitialHeader(),\n            ),\n          );\n          addHeaderToFetchResult = false;\n        });\n      }\n\n      // need to modify the header response so that the query metrics are per partition\n      if (headerResponse != null && Constants.HttpHeaders.QueryMetrics in headerResponse) {\n        // \"0\" is the default partition before one is actually assigned.\n        const queryMetrics = headerResponse[Constants.HttpHeaders.QueryMetrics][\"0\"];\n\n        // Wraping query metrics in a object where the keys are the partition key range.\n        headerResponse[Constants.HttpHeaders.QueryMetrics] = {};\n        headerResponse[Constants.HttpHeaders.QueryMetrics][this.targetPartitionKeyRange.id] =\n          queryMetrics;\n      }\n      mergeHeaders(this.respHeaders, headerResponse);\n    } catch (err: any) {\n      if (DocumentProducer._needPartitionKeyRangeCacheRefresh(err)) {\n        // Split just happend\n        // Buffer the error so the execution context can still get the feedResponses in the itemBuffer\n        const bufferedError = new FetchResult(undefined, err);\n        this.fetchResults.push(bufferedError);\n        mergeHeaders(this.respHeaders, err.headers);\n      } else {\n        this._updateStates(err, err.resources === undefined);\n        throw err;\n      }\n    }\n  }\n\n  public getTargetParitionKeyRange(): PartitionKeyRange {\n    return this.targetPartitionKeyRange;\n  }\n  /**\n   * Peak the next item in the buffer\n   */\n  public peakNextItem(): any {\n    if (this.err) {\n      throw this.err;\n    }\n    if (this.allFetched || this.fetchResults.length === 0) {\n      return undefined;\n    }\n    const fetchResult = this.fetchResults[0];\n\n    switch (fetchResult.fetchResultType) {\n      case FetchResultType.Done:\n        return undefined;\n\n      case FetchResultType.Exception: // do not throw this error\n        return undefined;\n\n      case FetchResultType.Result:\n        return fetchResult.feedResponse;\n    }\n  }\n\n  /**\n   * Returns the first item in the buffered results if any, or [] otherwise.\n   */\n  public async fetchNextItem(): Promise<Response<any>> {\n    if (this.err) {\n      this._updateStates(this.err, undefined);\n      throw this.err;\n    }\n    if (this.allFetched) {\n      return { result: undefined, headers: this._getAndResetActiveResponseHeaders() };\n    }\n    try {\n      const { result, headers } = this.current();\n      this._updateStates(undefined, result === undefined);\n      if (result === undefined || result.length === 0) {\n        return { result: undefined, headers };\n      }\n      return { result, headers }; //\n    } catch (err: any) {\n      this._updateStates(err, err.item === undefined);\n      throw err;\n    }\n  }\n  /**\n   * Fetches all the buffered results\n   */\n  public async fetchBufferedItems(): Promise<Response<any[]>> {\n    if (this.err) {\n      this._updateStates(this.err, undefined);\n      throw this.err;\n    }\n    if (this.allFetched) {\n      return { result: undefined, headers: this._getAndResetActiveResponseHeaders() };\n    }\n    const resources: any[] = [];\n    const resHeaders: CosmosHeaders = getInitialHeader();\n    try {\n      while (this.fetchResults.length > 0) {\n        const { result, headers } = this.current();\n        this._updateStates(undefined, result === undefined);\n        mergeHeaders(resHeaders, headers);\n        if (result === undefined) {\n          return { result: resources.length > 0 ? resources : undefined, headers: resHeaders };\n        } else {\n          resources.push(result);\n        }\n      }\n      return { result: resources, headers: resHeaders };\n    } catch (err: any) {\n      this._updateStates(err, err.item === undefined);\n      throw err;\n    }\n  }\n\n  /**\n   * Retrieve the current element on the DocumentProducer.\n   */\n  private current(): Response<any> {\n    // If something is buffered just give that\n    if (this.fetchResults.length > 0) {\n      const fetchResult = this.fetchResults.shift();\n      // Need to unwrap fetch results\n      switch (fetchResult.fetchResultType) {\n        case FetchResultType.Done:\n          return {\n            result: undefined,\n            headers: this._getAndResetActiveResponseHeaders(),\n          };\n        case FetchResultType.Exception:\n          fetchResult.error.headers = this._getAndResetActiveResponseHeaders();\n          throw fetchResult.error;\n        case FetchResultType.Result:\n          return {\n            result: fetchResult.feedResponse,\n            headers: this._getAndResetActiveResponseHeaders(),\n          };\n      }\n    }\n\n    // If there isn't anymore items left to fetch then let the user know.\n    if (this.allFetched) {\n      return {\n        result: undefined,\n        headers: this._getAndResetActiveResponseHeaders(),\n      };\n    }\n\n    // If the internal buffer is empty, return empty result\n    return { result: [], headers: this._getAndResetActiveResponseHeaders() };\n  }\n}\n"]}
+{"version":3,"file":"documentProducer.js","sourceRoot":"","sources":["../../../src/queryExecutionContext/documentProducer.ts"],"names":[],"mappings":"AAIA,OAAO,EACL,SAAS,EACT,aAAa,EACb,eAAe,EACf,YAAY,EACZ,WAAW,EACX,cAAc,GACf,MAAM,oBAAoB,CAAC;AAI5B,OAAO,EAAE,4BAA4B,EAAE,MAAM,mCAAmC,CAAC;AAEjF,OAAO,EAAE,WAAW,EAAE,eAAe,EAAE,MAAM,kBAAkB,CAAC;AAChE,OAAO,EAAE,gBAAgB,EAAE,MAAM,kBAAkB,CAAC;AAIpD,cAAc;AACd,MAAM,OAAO,gBAAgB;IAe3B;;;;;;;OAOG;IACH,YACU,aAA4B,EACpC,cAAsB,EACtB,KAAmB,EACnB,uBAA0C,EAC1C,OAAoB,EACpB,oBAA4B,EAC5B,QAAiB,EACjB,MAAe,EACf,0BAAmC,KAAK;QARhC,kBAAa,GAAb,aAAa,CAAe;QAf/B,eAAU,GAAW,CAAC,CAAC;QAiEvB,kBAAa,GAA0B,KAAK,EACjD,cAAsC,EACtC,OAAoB,EACpB,oBAA4B,EACC,EAAE;YAC/B,MAAM,IAAI,GAAG,eAAe,CAAC,IAAI,CAAC,cAAc,EAAE,YAAY,CAAC,IAAI,CAAC,CAAC;YACrE,cAAc,CAAC,OAAO,CAAC,EAAE,mBAAmB,EAAE,IAAI,CAAC,uBAAuB,CAAC,EAAE,EAAE,CAAC,CAAC;YACjF,MAAM,EAAE,GAAG,aAAa,CAAC,IAAI,CAAC,cAAc,CAAC,CAAC;YAC9C,MAAM,QAAQ,GAAG,IAAI,CAAC,uBAAuB,CAAC,CAAC,CAAC,IAAI,CAAC,QAAQ,CAAC,CAAC,CAAC,SAAS,CAAC;YAC1E,MAAM,MAAM,GAAG,IAAI,CAAC,uBAAuB,CAAC,CAAC,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC,CAAC,SAAS,CAAC;YAEtE,OAAO,IAAI,CAAC,aAAa,CAAC,SAAS,CAAC;gBAClC,IAAI;gBACJ,YAAY,EAAE,YAAY,CAAC,IAAI;gBAC/B,UAAU,EAAE,EAAE;gBACd,QAAQ,EAAE,CAAC,MAAW,EAAE,EAAE,CAAC,MAAM,CAAC,SAAS;gBAC3C,KAAK,EAAE,IAAI,CAAC,KAAK;gBACjB,OAAO;gBACP,cAAc;gBACd,mBAAmB,EAAE,IAAI,CAAC,uBAAuB,CAAC,IAAI,CAAC;gBACvD,oBAAoB,EAAE,oBAAoB;gBAC1C,QAAQ,EAAE,QAAQ;gBAClB,MAAM,EAAE,MAAM;aACf,CAAC,CAAC;QACL,CAAC,CAAC;QAhEA,oBAAoB;QACpB,IAAI,CAAC,cAAc,GAAG,cAAc,CAAC;QACrC,IAAI,CAAC,KAAK,GAAG,KAAK,CAAC;QACnB,IAAI,CAAC,uBAAuB,GAAG,uBAAuB,CAAC;QACvD,IAAI,CAAC,YAAY,GAAG,EAAE,CAAC;QAEvB,IAAI,CAAC,UAAU,GAAG,KAAK,CAAC;QACxB,IAAI,CAAC,GAAG,GAAG,SAAS,CAAC;QAErB,IAAI,CAAC,yBAAyB,GAAG,SAAS,CAAC;QAC3C,IAAI,CAAC,iBAAiB,GAAG,SAAS,CAAC;QAEnC,IAAI,CAAC,wBAAwB,GAAG,IAAI,4BAA4B,CAC9D,OAAO,EACP,IAAI,CAAC,aAAa,EAClB,oBAAoB,CACrB,CAAC;QACF,IAAI,CAAC,QAAQ,GAAG,QAAQ,CAAC;QACzB,IAAI,CAAC,MAAM,GAAG,MAAM,CAAC;QACrB,IAAI,CAAC,uBAAuB,GAAG,uBAAuB,CAAC;IACzD,CAAC;IACM,iBAAiB;QACtB,MAAM,eAAe,GAAG,EAAE,CAAC;QAC3B,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,IAAI,GAAG,KAAK,EAAE,CAAC,GAAG,IAAI,CAAC,YAAY,CAAC,MAAM,IAAI,CAAC,IAAI,EAAE,CAAC,EAAE,EAAE,CAAC;YACzE,MAAM,WAAW,GAAG,IAAI,CAAC,YAAY,CAAC,CAAC,CAAC,CAAC;YACzC,QAAQ,WAAW,CAAC,eAAe,EAAE,CAAC;gBACpC,KAAK,eAAe,CAAC,IAAI;oBACvB,IAAI,GAAG,IAAI,CAAC;oBACZ,MAAM;gBACR,KAAK,eAAe,CAAC,SAAS;oBAC5B,IAAI,GAAG,IAAI,CAAC;oBACZ,MAAM;gBACR,KAAK,eAAe,CAAC,MAAM;oBACzB,eAAe,CAAC,IAAI,CAAC,WAAW,CAAC,YAAY,CAAC,CAAC;oBAC/C,MAAM;YACV,CAAC;QACH,CAAC;QACD,OAAO,eAAe,CAAC;IACzB,CAAC;IA4BM,cAAc;QACnB,OAAO,IAAI,CAAC,wBAAwB,CAAC,cAAc,EAAE,IAAI,IAAI,CAAC,YAAY,CAAC,MAAM,KAAK,CAAC,CAAC;IAC1F,CAAC;IAEM,QAAQ;QACb,IAAI,IAAI,CAAC,YAAY,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;YACnC,MAAM,WAAW,GAAG,IAAI,CAAC,YAAY,CAAC,CAAC,CAAC,CAAC;YACzC,IAAI,WAAW,CAAC,eAAe,KAAK,eAAe,CAAC,SAAS,EAAE,CAAC;gBAC9D,IAAI,gBAAgB,CAAC,kCAAkC,CAAC,WAAW,CAAC,KAAK,CAAC,EAAE,CAAC;oBAC3E,OAAO,IAAI,CAAC;gBACd,CAAC;YACH,CAAC;QACH,CAAC;QACD,OAAO,KAAK,CAAC;IACf,CAAC;IAEO,aAAa,CAAC,GAAQ,EAAE,UAAmB;QACjD,IAAI,GAAG,EAAE,CAAC;YACR,IAAI,CAAC,GAAG,GAAG,GAAG,CAAC;YACf,OAAO;QACT,CAAC;QACD,IAAI,UAAU,EAAE,CAAC;YACf,IAAI,CAAC,UAAU,GAAG,IAAI,CAAC;QACzB,CAAC;QACD,IAAI,IAAI,CAAC,wBAAwB,CAAC,iBAAiB,KAAK,IAAI,CAAC,iBAAiB,EAAE,CAAC;YAC/E,kBAAkB;YAClB,OAAO;QACT,CAAC;QACD,IAAI,CAAC,yBAAyB,GAAG,IAAI,CAAC,iBAAiB,CAAC;QACxD,IAAI,CAAC,iBAAiB,GAAG,IAAI,CAAC,wBAAwB,CAAC,iBAAiB,CAAC;IAC3E,CAAC;IAEO,MAAM,CAAC,kCAAkC,CAAC,KAAU;QAC1D,cAAc;QACd,OAAO,CACL,KAAK,CAAC,IAAI,KAAK,WAAW,CAAC,IAAI;YAC/B,WAAW,IAAI,KAAK;YACpB,KAAK,CAAC,WAAW,CAAC,KAAK,cAAc,CAAC,qBAAqB,CAC5D,CAAC;IACJ,CAAC;IAED;;OAEG;IACI,KAAK,CAAC,UAAU,CAAC,cAAsC;QAC5D,IAAI,IAAI,CAAC,GAAG,EAAE,CAAC;YACb,MAAM,IAAI,CAAC,GAAG,CAAC;QACjB,CAAC;QAED,IAAI,CAAC;YACH,MAAM,EAAE,MAAM,EAAE,SAAS,EAAE,OAAO,EAAE,cAAc,EAAE,GAClD,MAAM,IAAI,CAAC,wBAAwB,CAAC,SAAS,CAAC,cAAc,CAAC,CAAC;YAChE,EAAE,IAAI,CAAC,UAAU,CAAC;YAClB,IAAI,CAAC,aAAa,CAAC,SAAS,EAAE,SAAS,KAAK,SAAS,CAAC,CAAC;YACvD,IAAI,SAAS,KAAK,SAAS,EAAE,CAAC;gBAC5B,sDAAsD;gBACtD,IAAI,sBAAsB,GAAG,IAAI,CAAC;gBAClC,SAAS,CAAC,OAAO,CAAC,CAAC,OAAY,EAAE,EAAE;oBACjC,IAAI,CAAC,YAAY,CAAC,IAAI,CACpB,IAAI,WAAW,CACb,OAAO,EACP,SAAS,EACT,sBAAsB,CAAC,CAAC,CAAC,cAAc,CAAC,CAAC,CAAC,gBAAgB,EAAE,CAC7D,CACF,CAAC;oBACF,sBAAsB,GAAG,KAAK,CAAC;gBACjC,CAAC,CAAC,CAAC;YACL,CAAC;YAED,iFAAiF;YACjF,IAAI,cAAc,IAAI,IAAI,IAAI,SAAS,CAAC,WAAW,CAAC,YAAY,IAAI,cAAc,EAAE,CAAC;gBACnF,gEAAgE;gBAChE,MAAM,YAAY,GAAG,cAAc,CAAC,SAAS,CAAC,WAAW,CAAC,YAAY,CAAC,CAAC,GAAG,CAAC,CAAC;gBAE7E,gFAAgF;gBAChF,cAAc,CAAC,SAAS,CAAC,WAAW,CAAC,YAAY,CAAC,GAAG,EAAE,CAAC;gBACxD,cAAc,CAAC,SAAS,CAAC,WAAW,CAAC,YAAY,CAAC,CAAC,IAAI,CAAC,uBAAuB,CAAC,EAAE,CAAC;oBACjF,YAAY,CAAC;YACjB,CAAC;YACD,OAAO,cAAc,CAAC;QACxB,CAAC;QAAC,OAAO,GAAQ,EAAE,CAAC;YAClB,IAAI,gBAAgB,CAAC,kCAAkC,CAAC,GAAG,CAAC,EAAE,CAAC;gBAC7D,qBAAqB;gBACrB,8FAA8F;gBAC9F,MAAM,aAAa,GAAG,IAAI,WAAW,CAAC,SAAS,EAAE,GAAG,CAAC,CAAC;gBACtD,IAAI,CAAC,YAAY,CAAC,IAAI,CAAC,aAAa,CAAC,CAAC;gBACtC,OAAO,GAAG,CAAC,OAAO,CAAC;YACrB,CAAC;iBAAM,CAAC;gBACN,IAAI,CAAC,aAAa,CAAC,GAAG,EAAE,GAAG,CAAC,SAAS,KAAK,SAAS,CAAC,CAAC;gBACrD,MAAM,GAAG,CAAC;YACZ,CAAC;QACH,CAAC;IACH,CAAC;IAEM,yBAAyB;QAC9B,OAAO,IAAI,CAAC,uBAAuB,CAAC;IACtC,CAAC;IACD;;OAEG;IACI,YAAY;QACjB,IAAI,IAAI,CAAC,GAAG,EAAE,CAAC;YACb,MAAM,IAAI,CAAC,GAAG,CAAC;QACjB,CAAC;QACD,IAAI,IAAI,CAAC,UAAU,IAAI,IAAI,CAAC,YAAY,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;YACtD,OAAO,SAAS,CAAC;QACnB,CAAC;QACD,MAAM,WAAW,GAAG,IAAI,CAAC,YAAY,CAAC,CAAC,CAAC,CAAC;QAEzC,QAAQ,WAAW,CAAC,eAAe,EAAE,CAAC;YACpC,KAAK,eAAe,CAAC,IAAI;gBACvB,OAAO,SAAS,CAAC;YAEnB,KAAK,eAAe,CAAC,SAAS,EAAE,0BAA0B;gBACxD,OAAO,SAAS,CAAC;YAEnB,KAAK,eAAe,CAAC,MAAM;gBACzB,OAAO,WAAW,CAAC,YAAY,CAAC;QACpC,CAAC;IACH,CAAC;IAED;;OAEG;IACI,KAAK,CAAC,aAAa;QACxB,IAAI,IAAI,CAAC,GAAG,EAAE,CAAC;YACb,IAAI,CAAC,aAAa,CAAC,IAAI,CAAC,GAAG,EAAE,SAAS,CAAC,CAAC;YACxC,MAAM,IAAI,CAAC,GAAG,CAAC;QACjB,CAAC;QACD,IAAI,IAAI,CAAC,UAAU,EAAE,CAAC;YACpB,OAAO,EAAE,MAAM,EAAE,SAAS,EAAE,OAAO,EAAE,gBAAgB,EAAE,EAAE,CAAC;QAC5D,CAAC;QACD,IAAI,CAAC;YACH,MAAM,EAAE,MAAM,EAAE,OAAO,EAAE,GAAG,IAAI,CAAC,OAAO,EAAE,CAAC;YAC3C,IAAI,CAAC,aAAa,CAAC,SAAS,EAAE,MAAM,KAAK,SAAS,CAAC,CAAC;YACpD,IAAI,MAAM,KAAK,SAAS,IAAI,MAAM,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;gBAChD,OAAO,EAAE,MAAM,EAAE,SAAS,EAAE,OAAO,EAAE,CAAC;YACxC,CAAC;YACD,OAAO,EAAE,MAAM,EAAE,OAAO,EAAE,CAAC,CAAC,EAAE;QAChC,CAAC;QAAC,OAAO,GAAQ,EAAE,CAAC;YAClB,IAAI,CAAC,aAAa,CAAC,GAAG,EAAE,GAAG,CAAC,IAAI,KAAK,SAAS,CAAC,CAAC;YAChD,MAAM,GAAG,CAAC;QACZ,CAAC;IACH,CAAC;IACD;;OAEG;IACI,KAAK,CAAC,kBAAkB;QAC7B,IAAI,IAAI,CAAC,GAAG,EAAE,CAAC;YACb,IAAI,CAAC,aAAa,CAAC,IAAI,CAAC,GAAG,EAAE,SAAS,CAAC,CAAC;YACxC,MAAM,IAAI,CAAC,GAAG,CAAC;QACjB,CAAC;QACD,IAAI,IAAI,CAAC,UAAU,EAAE,CAAC;YACpB,OAAO,EAAE,MAAM,EAAE,SAAS,EAAE,OAAO,EAAE,gBAAgB,EAAE,EAAE,CAAC;QAC5D,CAAC;QACD,MAAM,SAAS,GAAU,EAAE,CAAC;QAC5B,IAAI,CAAC;YACH,OAAO,IAAI,CAAC,YAAY,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC;gBACpC,MAAM,EAAE,MAAM,EAAE,GAAG,IAAI,CAAC,OAAO,EAAE,CAAC;gBAClC,IAAI,CAAC,aAAa,CAAC,SAAS,EAAE,MAAM,KAAK,SAAS,CAAC,CAAC;gBACpD,IAAI,MAAM,KAAK,SAAS,EAAE,CAAC;oBACzB,OAAO;wBACL,MAAM,EAAE,SAAS,CAAC,MAAM,GAAG,CAAC,CAAC,CAAC,CAAC,SAAS,CAAC,CAAC,CAAC,SAAS;wBACpD,OAAO,EAAE,gBAAgB,EAAE;qBAC5B,CAAC;gBACJ,CAAC;qBAAM,CAAC;oBACN,SAAS,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;gBACzB,CAAC;YACH,CAAC;YACD,OAAO,EAAE,MAAM,EAAE,SAAS,EAAE,OAAO,EAAE,gBAAgB,EAAE,EAAE,CAAC;QAC5D,CAAC;QAAC,OAAO,GAAQ,EAAE,CAAC;YAClB,IAAI,CAAC,aAAa,CAAC,GAAG,EAAE,GAAG,CAAC,IAAI,KAAK,SAAS,CAAC,CAAC;YAChD,MAAM,GAAG,CAAC;QACZ,CAAC;IACH,CAAC;IAED;;OAEG;IACK,OAAO;QACb,0CAA0C;QAC1C,IAAI,IAAI,CAAC,YAAY,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC;YACjC,MAAM,WAAW,GAAG,IAAI,CAAC,YAAY,CAAC,KAAK,EAAE,CAAC;YAC9C,+BAA+B;YAC/B,QAAQ,WAAW,CAAC,eAAe,EAAE,CAAC;gBACpC,KAAK,eAAe,CAAC,IAAI;oBACvB,OAAO;wBACL,MAAM,EAAE,SAAS;wBACjB,OAAO,EAAE,gBAAgB,EAAE;qBAC5B,CAAC;gBACJ,KAAK,eAAe,CAAC,SAAS;oBAC5B,WAAW,CAAC,KAAK,CAAC,OAAO,GAAG,gBAAgB,EAAE,CAAC;oBAC/C,MAAM,WAAW,CAAC,KAAK,CAAC;gBAC1B,KAAK,eAAe,CAAC,MAAM;oBACzB,OAAO;wBACL,MAAM,EAAE,WAAW,CAAC,YAAY;wBAChC,OAAO,EAAE,gBAAgB,EAAE;qBAC5B,CAAC;YACN,CAAC;QACH,CAAC;QAED,qEAAqE;QACrE,IAAI,IAAI,CAAC,UAAU,EAAE,CAAC;YACpB,OAAO;gBACL,MAAM,EAAE,SAAS;gBACjB,OAAO,EAAE,gBAAgB,EAAE;aAC5B,CAAC;QACJ,CAAC;QAED,uDAAuD;QACvD,OAAO,EAAE,MAAM,EAAE,EAAE,EAAE,OAAO,EAAE,gBAAgB,EAAE,EAAE,CAAC;IACrD,CAAC;CACF","sourcesContent":["// Copyright (c) Microsoft Corporation.\n// Licensed under the MIT License.\nimport type { PartitionKeyRange, Resource } from \"../client/index.js\";\nimport type { ClientContext } from \"../ClientContext.js\";\nimport {\n  Constants,\n  getIdFromLink,\n  getPathFromLink,\n  ResourceType,\n  StatusCodes,\n  SubStatusCodes,\n} from \"../common/index.js\";\nimport type { DiagnosticNodeInternal } from \"../diagnostics/DiagnosticNodeInternal.js\";\nimport type { FeedOptions } from \"../request/index.js\";\nimport type { Response } from \"../request/index.js\";\nimport { DefaultQueryExecutionContext } from \"./defaultQueryExecutionContext.js\";\nimport type { FetchFunctionCallback } from \"./defaultQueryExecutionContext.js\";\nimport { FetchResult, FetchResultType } from \"./FetchResult.js\";\nimport { getInitialHeader } from \"./headerUtils.js\";\nimport type { CosmosHeaders } from \"./headerUtils.js\";\nimport type { SqlQuerySpec } from \"./index.js\";\n\n/** @hidden */\nexport class DocumentProducer {\n  private collectionLink: string;\n  private query: string | SqlQuerySpec;\n  public targetPartitionKeyRange: PartitionKeyRange;\n  public fetchResults: FetchResult[];\n  public allFetched: boolean;\n  private err: Error;\n  public previousContinuationToken: string;\n  public continuationToken: string;\n  public generation: number = 0;\n  private internalExecutionContext: DefaultQueryExecutionContext;\n  public startEpk: string;\n  public endEpk: string;\n  public populateEpkRangeHeaders: boolean;\n\n  /**\n   * Provides the Target Partition Range Query Execution Context.\n   * @param clientContext  - The service endpoint to use to create the client.\n   * @param collectionLink - Represents collection link\n   * @param query          - A SQL query.\n   * @param targetPartitionKeyRange - Query Target Partition key Range\n   * @hidden\n   */\n  constructor(\n    private clientContext: ClientContext,\n    collectionLink: string,\n    query: SqlQuerySpec,\n    targetPartitionKeyRange: PartitionKeyRange,\n    options: FeedOptions,\n    correlatedActivityId: string,\n    startEpk?: string,\n    endEpk?: string,\n    populateEpkRangeHeaders: boolean = false,\n  ) {\n    // TODO: any options\n    this.collectionLink = collectionLink;\n    this.query = query;\n    this.targetPartitionKeyRange = targetPartitionKeyRange;\n    this.fetchResults = [];\n\n    this.allFetched = false;\n    this.err = undefined;\n\n    this.previousContinuationToken = undefined;\n    this.continuationToken = undefined;\n\n    this.internalExecutionContext = new DefaultQueryExecutionContext(\n      options,\n      this.fetchFunction,\n      correlatedActivityId,\n    );\n    this.startEpk = startEpk;\n    this.endEpk = endEpk;\n    this.populateEpkRangeHeaders = populateEpkRangeHeaders;\n  }\n  public peekBufferedItems(): any[] {\n    const bufferedResults = [];\n    for (let i = 0, done = false; i < this.fetchResults.length && !done; i++) {\n      const fetchResult = this.fetchResults[i];\n      switch (fetchResult.fetchResultType) {\n        case FetchResultType.Done:\n          done = true;\n          break;\n        case FetchResultType.Exception:\n          done = true;\n          break;\n        case FetchResultType.Result:\n          bufferedResults.push(fetchResult.feedResponse);\n          break;\n      }\n    }\n    return bufferedResults;\n  }\n\n  public fetchFunction: FetchFunctionCallback = async (\n    diagnosticNode: DiagnosticNodeInternal,\n    options: FeedOptions,\n    correlatedActivityId: string,\n  ): Promise<Response<Resource>> => {\n    const path = getPathFromLink(this.collectionLink, ResourceType.item);\n    diagnosticNode.addData({ partitionKeyRangeId: this.targetPartitionKeyRange.id });\n    const id = getIdFromLink(this.collectionLink);\n    const startEpk = this.populateEpkRangeHeaders ? this.startEpk : undefined;\n    const endEpk = this.populateEpkRangeHeaders ? this.endEpk : undefined;\n\n    return this.clientContext.queryFeed({\n      path,\n      resourceType: ResourceType.item,\n      resourceId: id,\n      resultFn: (result: any) => result.Documents,\n      query: this.query,\n      options,\n      diagnosticNode,\n      partitionKeyRangeId: this.targetPartitionKeyRange[\"id\"],\n      correlatedActivityId: correlatedActivityId,\n      startEpk: startEpk,\n      endEpk: endEpk,\n    });\n  };\n\n  public hasMoreResults(): boolean {\n    return this.internalExecutionContext.hasMoreResults() || this.fetchResults.length !== 0;\n  }\n\n  public gotSplit(): boolean {\n    if (this.fetchResults.length !== 0) {\n      const fetchResult = this.fetchResults[0];\n      if (fetchResult.fetchResultType === FetchResultType.Exception) {\n        if (DocumentProducer._needPartitionKeyRangeCacheRefresh(fetchResult.error)) {\n          return true;\n        }\n      }\n    }\n    return false;\n  }\n\n  private _updateStates(err: any, allFetched: boolean): void {\n    if (err) {\n      this.err = err;\n      return;\n    }\n    if (allFetched) {\n      this.allFetched = true;\n    }\n    if (this.internalExecutionContext.continuationToken === this.continuationToken) {\n      // nothing changed\n      return;\n    }\n    this.previousContinuationToken = this.continuationToken;\n    this.continuationToken = this.internalExecutionContext.continuationToken;\n  }\n\n  private static _needPartitionKeyRangeCacheRefresh(error: any): boolean {\n    // TODO: error\n    return (\n      error.code === StatusCodes.Gone &&\n      \"substatus\" in error &&\n      error[\"substatus\"] === SubStatusCodes.PartitionKeyRangeGone\n    );\n  }\n\n  /**\n   * Fetches and bufferes the next page of results in internal buffer\n   */\n  public async bufferMore(diagnosticNode: DiagnosticNodeInternal): Promise<CosmosHeaders> {\n    if (this.err) {\n      throw this.err;\n    }\n\n    try {\n      const { result: resources, headers: headerResponse } =\n        await this.internalExecutionContext.fetchMore(diagnosticNode);\n      ++this.generation;\n      this._updateStates(undefined, resources === undefined);\n      if (resources !== undefined) {\n        // add fetched header to the 1st element in the buffer\n        let addHeaderToFetchResult = true;\n        resources.forEach((element: any) => {\n          this.fetchResults.push(\n            new FetchResult(\n              element,\n              undefined,\n              addHeaderToFetchResult ? headerResponse : getInitialHeader(),\n            ),\n          );\n          addHeaderToFetchResult = false;\n        });\n      }\n\n      // need to modify the header response so that the query metrics are per partition\n      if (headerResponse != null && Constants.HttpHeaders.QueryMetrics in headerResponse) {\n        // \"0\" is the default partition before one is actually assigned.\n        const queryMetrics = headerResponse[Constants.HttpHeaders.QueryMetrics][\"0\"];\n\n        // Wraping query metrics in a object where the keys are the partition key range.\n        headerResponse[Constants.HttpHeaders.QueryMetrics] = {};\n        headerResponse[Constants.HttpHeaders.QueryMetrics][this.targetPartitionKeyRange.id] =\n          queryMetrics;\n      }\n      return headerResponse;\n    } catch (err: any) {\n      if (DocumentProducer._needPartitionKeyRangeCacheRefresh(err)) {\n        // Split just happend\n        // Buffer the error so the execution context can still get the feedResponses in the itemBuffer\n        const bufferedError = new FetchResult(undefined, err);\n        this.fetchResults.push(bufferedError);\n        return err.headers;\n      } else {\n        this._updateStates(err, err.resources === undefined);\n        throw err;\n      }\n    }\n  }\n\n  public getTargetParitionKeyRange(): PartitionKeyRange {\n    return this.targetPartitionKeyRange;\n  }\n  /**\n   * Peak the next item in the buffer\n   */\n  public peakNextItem(): any {\n    if (this.err) {\n      throw this.err;\n    }\n    if (this.allFetched || this.fetchResults.length === 0) {\n      return undefined;\n    }\n    const fetchResult = this.fetchResults[0];\n\n    switch (fetchResult.fetchResultType) {\n      case FetchResultType.Done:\n        return undefined;\n\n      case FetchResultType.Exception: // do not throw this error\n        return undefined;\n\n      case FetchResultType.Result:\n        return fetchResult.feedResponse;\n    }\n  }\n\n  /**\n   * Returns the first item in the buffered results if any, or [] otherwise.\n   */\n  public async fetchNextItem(): Promise<Response<any>> {\n    if (this.err) {\n      this._updateStates(this.err, undefined);\n      throw this.err;\n    }\n    if (this.allFetched) {\n      return { result: undefined, headers: getInitialHeader() };\n    }\n    try {\n      const { result, headers } = this.current();\n      this._updateStates(undefined, result === undefined);\n      if (result === undefined || result.length === 0) {\n        return { result: undefined, headers };\n      }\n      return { result, headers }; //\n    } catch (err: any) {\n      this._updateStates(err, err.item === undefined);\n      throw err;\n    }\n  }\n  /**\n   * Fetches all the buffered results\n   */\n  public async fetchBufferedItems(): Promise<Response<any[]>> {\n    if (this.err) {\n      this._updateStates(this.err, undefined);\n      throw this.err;\n    }\n    if (this.allFetched) {\n      return { result: undefined, headers: getInitialHeader() };\n    }\n    const resources: any[] = [];\n    try {\n      while (this.fetchResults.length > 0) {\n        const { result } = this.current();\n        this._updateStates(undefined, result === undefined);\n        if (result === undefined) {\n          return {\n            result: resources.length > 0 ? resources : undefined,\n            headers: getInitialHeader(),\n          };\n        } else {\n          resources.push(result);\n        }\n      }\n      return { result: resources, headers: getInitialHeader() };\n    } catch (err: any) {\n      this._updateStates(err, err.item === undefined);\n      throw err;\n    }\n  }\n\n  /**\n   * Retrieve the current element on the DocumentProducer.\n   */\n  private current(): Response<any> {\n    // If something is buffered just give that\n    if (this.fetchResults.length > 0) {\n      const fetchResult = this.fetchResults.shift();\n      // Need to unwrap fetch results\n      switch (fetchResult.fetchResultType) {\n        case FetchResultType.Done:\n          return {\n            result: undefined,\n            headers: getInitialHeader(),\n          };\n        case FetchResultType.Exception:\n          fetchResult.error.headers = getInitialHeader();\n          throw fetchResult.error;\n        case FetchResultType.Result:\n          return {\n            result: fetchResult.feedResponse,\n            headers: getInitialHeader(),\n          };\n      }\n    }\n\n    // If there isn't anymore items left to fetch then let the user know.\n    if (this.allFetched) {\n      return {\n        result: undefined,\n        headers: getInitialHeader(),\n      };\n    }\n\n    // If the internal buffer is empty, return empty result\n    return { result: [], headers: getInitialHeader() };\n  }\n}\n"]}