@azure/cosmos/dist/commonjs/bulk/BulkHelper.js

Microsoft Azure Cosmos DB Service Node.js SDK for NOSQL API

var __defProp = Object.defineProperty;
var __getOwnPropDesc = Object.getOwnPropertyDescriptor;
var __getOwnPropNames = Object.getOwnPropertyNames;
var __hasOwnProp = Object.prototype.hasOwnProperty;
var __export = (target, all) => {
  for (var name in all)
    __defProp(target, name, { get: all[name], enumerable: true });
};
var __copyProps = (to, from, except, desc) => {
  if (from && typeof from === "object" || typeof from === "function") {
    for (let key of __getOwnPropNames(from))
      if (!__hasOwnProp.call(to, key) && key !== except)
        __defProp(to, key, { get: () => from[key], enumerable: !(desc = __getOwnPropDesc(from, key)) || desc.enumerable });
  }
  return to;
};
var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: true }), mod);
var BulkHelper_exports = {};
__export(BulkHelper_exports, {
  BulkHelper: () => BulkHelper
});
module.exports = __toCommonJS(BulkHelper_exports);
var import_ClientUtils = require("../client/ClientUtils.js");
var import_constants = require("../common/constants.js");
var import_helper = require("../common/helper.js");
var import_statusCodes = require("../common/statusCodes.js");
var import_DiagnosticNodeInternal = require("../diagnostics/DiagnosticNodeInternal.js");
var import_PartitionKeyInternal = require("../documents/PartitionKeyInternal.js");
var import__ = require("../index.js");
var import_bulkExecutionRetryPolicy = require("../retry/bulkExecutionRetryPolicy.js");
var import_resourceThrottleRetryPolicy = require("../retry/resourceThrottleRetryPolicy.js");
var import_batch = require("../utils/batch.js");
var import_diagnostics = require("../utils/diagnostics.js");
var import_hash = require("../utils/hashing/hash.js");
var import_HelperPerPartition = require("./HelperPerPartition.js");
var import_index = require("./index.js");
class BulkHelper {
  container;
  clientContext;
  partitionKeyRangeCache;
  helpersByPartitionKeyRangeId;
  options;
  partitionKeyDefinition;
  partitionKeyDefinitionPromise;
  isCancelled;
  processedOperationCountRef = { count: 0 };
  operationPromisesList = [];
  congestionControlTimer;
  congestionControlDelayInMs = 1e3;
  staleRidError;
  operationsPerSleep = 100;
  // Number of operations to add per sleep
  intervalForPartialBatchInMs = 1e3;
  // Sleep interval before adding partial batch to dispatch queue
  /**
   * @internal
   */
  constructor(container, clientContext, partitionKeyRangeCache, options) {
    this.container = container;
    this.clientContext = clientContext;
    this.partitionKeyRangeCache = partitionKeyRangeCache;
    this.helpersByPartitionKeyRangeId = /* @__PURE__ */ new Map();
    this.options = options;
    this.executeRequest = this.executeRequest.bind(this);
    this.reBatchOperation = this.reBatchOperation.bind(this);
    this.refreshPartitionKeyRangeCache = this.refreshPartitionKeyRangeCache.bind(this);
    this.isCancelled = false;
    this.runCongestionControlTimer();
  }
  /**
   * adds operation(s) to the helper
   * @param operationInput - bulk operation or list of bulk operations
   */
  async execute(operationInput) {
    const addOperationPromises = [];
    const minimalPause = 0;
    try {
      for (let i = 0; i < operationInput.length; i++) {
        if (i % this.operationsPerSleep === 0) {
          await (0, import_helper.sleep)(minimalPause);
        }
        addOperationPromises.push(this.addOperation(operationInput[i], i));
      }
      await Promise.allSettled(addOperationPromises);
      while (this.processedOperationCountRef.count < operationInput.length) {
        this.helpersByPartitionKeyRangeId.forEach((helper) => {
          helper.addPartialBatchToQueue();
        });
        await (0, import_helper.sleep)(this.intervalForPartialBatchInMs);
      }
    } finally {
      if (this.congestionControlTimer) {
        clearInterval(this.congestionControlTimer);
      }
    }
    const settledResults = await Promise.allSettled(this.operationPromisesList);
    if (this.isCancelled && this.staleRidError) {
      throw this.staleRidError;
    }
    const bulkOperationResults = settledResults.map(
      (result) => result.status === "fulfilled" ? result.value : result.reason
    );
    const formattedResults = bulkOperationResults.map((result) => {
      if (result && result.error) {
        const { stack, ...otherProps } = result.error;
        const trimmedError = { message: result.error.message, ...otherProps };
        return {
          ...result,
          error: trimmedError
        };
      }
      return result;
    });
    return formattedResults;
  }
  async addOperation(operation, idx) {
    if (this.isCancelled) {
      return;
    }
    if (!operation) {
      this.operationPromisesList[idx] = Promise.resolve({
        operationInput: operation,
        error: Object.assign(new import__.ErrorResponse("Operation cannot be null or undefined."), {
          code: import_statusCodes.StatusCodes.InternalServerError
        })
      });
      return;
    }
    if (operation.operationType === "Create" || operation.operationType === "Upsert" || operation.operationType === "Replace") {
      if (!operation.resourceBody.id) {
        this.operationPromisesList[idx] = Promise.resolve({
          operationInput: operation,
          error: Object.assign(
            new import__.ErrorResponse(
              `Operation resource body must have an 'id' for ${operation.operationType} operations.`
            ),
            { code: import_statusCodes.StatusCodes.InternalServerError }
          )
        });
        this.processedOperationCountRef.count++;
        return;
      }
    }
    if (operation.partitionKey === void 0) {
      this.operationPromisesList[idx] = Promise.resolve({
        operationInput: operation,
        error: Object.assign(
          new import__.ErrorResponse(`PartitionKey is required for ${operation.operationType} operations.`),
          { code: import_statusCodes.StatusCodes.InternalServerError }
        )
      });
      this.processedOperationCountRef.count++;
      return;
    }
    let operationError;
    let diagnosticNode;
    let unencryptedOperation;
    let partitionKeyRangeId;
    try {
      diagnosticNode = new import_DiagnosticNodeInternal.DiagnosticNodeInternal(
        this.clientContext.diagnosticLevel,
        import_DiagnosticNodeInternal.DiagnosticNodeType.CLIENT_REQUEST_NODE,
        null
      );
      if (!this.partitionKeyDefinition) {
        if (!this.partitionKeyDefinitionPromise) {
          this.partitionKeyDefinitionPromise = (async () => {
            try {
              const partitionKeyDefinition = await (0, import_ClientUtils.readPartitionKeyDefinition)(
                diagnosticNode,
                this.container
              );
              this.partitionKeyDefinition = partitionKeyDefinition;
              return partitionKeyDefinition;
            } finally {
              this.partitionKeyDefinitionPromise = null;
            }
          })();
        }
        await this.partitionKeyDefinitionPromise;
      }
      unencryptedOperation = (0, import_helper.copyObject)(operation);
      if (this.clientContext.enableEncryption) {
        operation = (0, import_helper.copyObject)(operation);
        await this.container.checkAndInitializeEncryption();
        diagnosticNode.beginEncryptionDiagnostics(import_constants.Constants.Encryption.DiagnosticsEncryptOperation);
        const { operation: encryptedOp, totalPropertiesEncryptedCount } = await (0, import_batch.encryptOperationInput)(this.container.encryptionProcessor, operation, 0);
        operation = encryptedOp;
        diagnosticNode.endEncryptionDiagnostics(
          import_constants.Constants.Encryption.DiagnosticsEncryptOperation,
          totalPropertiesEncryptedCount
        );
      }
      partitionKeyRangeId = await this.resolvePartitionKeyRangeId(operation, diagnosticNode);
    } catch (error) {
      operationError = error;
    }
    const helperForPartition = this.getHelperForPKRange(partitionKeyRangeId);
    const retryPolicy = this.getRetryPolicy();
    const context = new import_index.ItemOperationContext(partitionKeyRangeId, retryPolicy, diagnosticNode);
    const itemOperation = {
      unencryptedOperationInput: unencryptedOperation,
      operationInput: operation,
      operationContext: context
    };
    this.operationPromisesList[idx] = context.operationPromise;
    if (operationError) {
      const response = {
        operationInput: unencryptedOperation,
        error: Object.assign(new import__.ErrorResponse(operationError.message), {
          code: import_statusCodes.StatusCodes.InternalServerError,
          diagnostics: diagnosticNode?.toDiagnostic(this.clientContext.getClientConfig())
        })
      };
      context.fail(response);
      this.processedOperationCountRef.count++;
      return;
    }
    return helperForPartition.add(itemOperation);
  }
  async resolvePartitionKeyRangeId(operation, diagnosticNode) {
    const partitionKeyRanges = (await this.partitionKeyRangeCache.onCollectionRoutingMap(this.container.url, diagnosticNode)).getOrderedParitionKeyRanges();
    const partitionKey = (0, import_PartitionKeyInternal.convertToInternalPartitionKey)(operation.partitionKey);
    const hashedKey = (0, import_hash.hashPartitionKey)(partitionKey, this.partitionKeyDefinition);
    const matchingRange = partitionKeyRanges.find(
      (range) => (0, import_batch.isKeyInRange)(range.minInclusive, range.maxExclusive, hashedKey)
    );
    if (!matchingRange) {
      throw new Error("No matching partition key range found for the operation.");
    }
    return matchingRange.id;
  }
  getRetryPolicy() {
    const nextRetryPolicy = new import_resourceThrottleRetryPolicy.ResourceThrottleRetryPolicy(this.clientContext.getRetryOptions());
    return new import_bulkExecutionRetryPolicy.BulkExecutionRetryPolicy(nextRetryPolicy);
  }
  async executeRequest(operations, diagnosticNode) {
    if (this.isCancelled) {
      throw new import__.ErrorResponse("Bulk execution cancelled due to a previous error.");
    }
    if (!operations.length) return;
    const pkRangeId = operations[0].operationContext.pkRangeId;
    const path = (0, import_helper.getPathFromLink)(this.container.url, import_constants.ResourceType.item);
    const requestBody = [];
    for (const itemOperation of operations) {
      requestBody.push(this.prepareOperation(itemOperation.operationInput));
    }
    if (!this.options.containerRid) {
      this.options.containerRid = this.container._rid;
    }
    try {
      const response = await (0, import_diagnostics.addDiagnosticChild)(
        async (childNode) => this.clientContext.bulk({
          body: requestBody,
          partitionKeyRangeId: pkRangeId,
          path,
          resourceId: this.container.url,
          options: this.options,
          diagnosticNode: childNode
        }),
        diagnosticNode,
        import_DiagnosticNodeInternal.DiagnosticNodeType.BATCH_REQUEST
      );
      if (!response) {
        throw new import__.ErrorResponse("Failed to fetch bulk response.");
      }
      return import_index.BulkResponse.fromResponseMessage(response, operations);
    } catch (error) {
      if (this.clientContext.enableEncryption) {
        try {
          await this.container.throwIfRequestNeedsARetryPostPolicyRefresh(error);
        } catch (err) {
          await this.cancelExecution(err);
          return import_index.BulkResponse.createEmptyResponse(operations, 0, 0, {});
        }
      }
      return import_index.BulkResponse.fromResponseMessage(error, operations);
    }
  }
  prepareOperation(operationInput) {
    operationInput.partitionKey = (0, import_PartitionKeyInternal.convertToInternalPartitionKey)(operationInput.partitionKey);
    return {
      ...operationInput,
      partitionKey: JSON.stringify(operationInput.partitionKey)
    };
  }
  async reBatchOperation(operation, diagnosticNode) {
    const partitionKeyRangeId = await this.resolvePartitionKeyRangeId(
      operation.operationInput,
      diagnosticNode
    );
    operation.operationContext.updatePKRangeId(partitionKeyRangeId);
    const helper = this.getHelperForPKRange(partitionKeyRangeId);
    await helper.add(operation);
  }
  async cancelExecution(error) {
    this.isCancelled = true;
    this.staleRidError = error;
    for (const helper of this.helpersByPartitionKeyRangeId.values()) {
      await helper.dispose();
    }
    this.helpersByPartitionKeyRangeId.clear();
  }
  getHelperForPKRange(pkRangeId) {
    if (this.helpersByPartitionKeyRangeId.has(pkRangeId)) {
      return this.helpersByPartitionKeyRangeId.get(pkRangeId);
    }
    const newHelper = new import_HelperPerPartition.HelperPerPartition(
      this.executeRequest,
      this.reBatchOperation,
      this.refreshPartitionKeyRangeCache,
      this.clientContext.diagnosticLevel,
      this.clientContext.enableEncryption,
      this.clientContext.getClientConfig(),
      this.container.encryptionProcessor,
      this.processedOperationCountRef
    );
    this.helpersByPartitionKeyRangeId.set(pkRangeId, newHelper);
    return newHelper;
  }
  runCongestionControlTimer() {
    this.congestionControlTimer = setInterval(() => {
      this.helpersByPartitionKeyRangeId.forEach((helper) => {
        helper.runCongestionAlgorithm();
      });
    }, this.congestionControlDelayInMs);
  }
  async refreshPartitionKeyRangeCache(diagnosticNode) {
    await this.partitionKeyRangeCache.onCollectionRoutingMap(
      this.container.url,
      diagnosticNode,
      true
    );
  }
}
// Annotate the CommonJS export names for ESM import in node:
0 && (module.exports = {
  BulkHelper
});