@azure/cosmos
Version:
Microsoft Azure Cosmos DB Service Node.js SDK for NOSQL API
295 lines (294 loc) • 13.4 kB
JavaScript
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 pipelinedQueryExecutionContext_exports = {};
__export(pipelinedQueryExecutionContext_exports, {
PipelinedQueryExecutionContext: () => PipelinedQueryExecutionContext
});
module.exports = __toCommonJS(pipelinedQueryExecutionContext_exports);
var import_ErrorResponse = require("../request/ErrorResponse.js");
var import_OffsetLimitEndpointComponent = require("./EndpointComponent/OffsetLimitEndpointComponent.js");
var import_OrderByEndpointComponent = require("./EndpointComponent/OrderByEndpointComponent.js");
var import_OrderedDistinctEndpointComponent = require("./EndpointComponent/OrderedDistinctEndpointComponent.js");
var import_UnorderedDistinctEndpointComponent = require("./EndpointComponent/UnorderedDistinctEndpointComponent.js");
var import_GroupByEndpointComponent = require("./EndpointComponent/GroupByEndpointComponent.js");
var import_orderByQueryExecutionContext = require("./orderByQueryExecutionContext.js");
var import_parallelQueryExecutionContext = require("./parallelQueryExecutionContext.js");
var import_GroupByValueEndpointComponent = require("./EndpointComponent/GroupByValueEndpointComponent.js");
var import_NonStreamingOrderByDistinctEndpointComponent = require("./EndpointComponent/NonStreamingOrderByDistinctEndpointComponent.js");
var import_NonStreamingOrderByEndpointComponent = require("./EndpointComponent/NonStreamingOrderByEndpointComponent.js");
var import_QueryValidationHelper = require("./QueryValidationHelper.js");
var import_ContinuationTokenParser = require("./ContinuationTokenParser.js");
var import_LegacyFetchImplementation = require("./LegacyFetchImplementation.js");
var import_QueryControlFetchImplementation = require("./QueryControlFetchImplementation.js");
var import_constants = require("../common/constants.js");
class PipelinedQueryExecutionContext {
constructor(clientContext, collectionLink, query, options, partitionedQueryExecutionInfo, correlatedActivityId, emitRawOrderByPayload = false, supportsContinuationTokens = true) {
this.clientContext = clientContext;
this.collectionLink = collectionLink;
this.query = query;
this.options = options;
this.partitionedQueryExecutionInfo = partitionedQueryExecutionInfo;
this.emitRawOrderByPayload = emitRawOrderByPayload;
this.supportsContinuationTokens = supportsContinuationTokens;
if (!partitionedQueryExecutionInfo.queryInfo) {
throw new import_ErrorResponse.ErrorResponse(
"Query execution requires valid query plan information. The partitioned query execution info is missing queryInfo. This may indicate an invalid query or a problem with query planning."
);
}
if (!this.options.maxItemCount) {
this.options.maxItemCount = import_constants.QueryExecution.DEFAULT_PAGE_SIZE;
}
const pageSize = this.options.maxItemCount;
const analyzedQueryInfo = this.analyzeQueryInfo(partitionedQueryExecutionInfo.queryInfo);
const {
sortOrders,
nonStreamingOrderBy,
isOrderByQuery,
isGroupByQuery,
isUnorderedDistinctQuery,
querySupportsTokens
} = analyzedQueryInfo;
if (!querySupportsTokens) {
(0, import_QueryValidationHelper.rejectContinuationTokenForUnsupportedQueries)(this.options.continuationToken, [
import_QueryValidationHelper.QueryTypes.nonStreamingOrderBy(nonStreamingOrderBy),
import_QueryValidationHelper.QueryTypes.groupBy(isGroupByQuery),
import_QueryValidationHelper.QueryTypes.unorderedDistinct(isUnorderedDistinctQuery)
]);
}
const queryContinuationFields = this.options.continuationToken ? (0, import_ContinuationTokenParser.parseContinuationTokenFields)(this.options.continuationToken) : void 0;
this.endpoint = nonStreamingOrderBy ? this.createNonStreamingEndpoint(
partitionedQueryExecutionInfo,
sortOrders,
correlatedActivityId,
options
) : this.createStreamingEndpoint(
partitionedQueryExecutionInfo,
sortOrders,
correlatedActivityId,
isGroupByQuery,
queryContinuationFields
);
this.fetchBuffer = [];
if (this.options.enableQueryControl) {
const querySupportsContinuationTokens = this.supportsContinuationTokens && querySupportsTokens;
this.fetchImplementation = new import_QueryControlFetchImplementation.QueryControlFetchImplementation(
this.endpoint,
pageSize,
this.collectionLink,
this.options.continuationToken,
isOrderByQuery,
querySupportsContinuationTokens
);
} else {
this.fetchImplementation = new import_LegacyFetchImplementation.LegacyFetchImplementation(this.endpoint, pageSize);
}
}
fetchBuffer;
endpoint;
fetchImplementation;
hasMoreResults() {
return this.fetchBuffer.length !== 0 || this.endpoint.hasMoreResults();
}
async fetchMore(diagnosticNode) {
return this.fetchImplementation.fetchMore(diagnosticNode, this.fetchBuffer);
}
/**
* Creates a non-streaming endpoint for vector search and similar queries that require buffering
*/
createNonStreamingEndpoint(partitionedQueryExecutionInfo, sortOrders, correlatedActivityId, options) {
const queryInfo = partitionedQueryExecutionInfo.queryInfo;
if (!options.allowUnboundedNonStreamingQueries) {
this.checkQueryConstraints(queryInfo);
}
const vectorSearchBufferSize = this.calculateVectorSearchBufferSize(queryInfo, options);
this.validateVectorSearchBufferSize(vectorSearchBufferSize, options);
const baseContext = new import_parallelQueryExecutionContext.ParallelQueryExecutionContext(
this.clientContext,
this.collectionLink,
this.query,
this.options,
this.partitionedQueryExecutionInfo,
correlatedActivityId
);
return this.wrapWithNonStreamingComponent(
baseContext,
queryInfo,
sortOrders,
vectorSearchBufferSize
);
}
/**
* Creates a streaming endpoint with proper pipeline components
*/
createStreamingEndpoint(partitionedQueryExecutionInfo, sortOrders, correlatedActivityId, isGroupByQuery, queryContinuationFields) {
const queryInfo = partitionedQueryExecutionInfo.queryInfo;
let endpoint = this.createBaseExecutionContext(
partitionedQueryExecutionInfo,
sortOrders,
correlatedActivityId
);
endpoint = this.applyGroupByComponents(endpoint, queryInfo, isGroupByQuery);
endpoint = this.applyDistinctComponents(endpoint, queryInfo, queryContinuationFields);
endpoint = this.applyLimitComponents(endpoint, queryInfo, queryContinuationFields);
return endpoint;
}
/**
* Creates the base execution context (OrderBy or Parallel)
*/
createBaseExecutionContext(partitionedQueryExecutionInfo, sortOrders, correlatedActivityId) {
if (Array.isArray(sortOrders) && sortOrders.length > 0) {
return new import_OrderByEndpointComponent.OrderByEndpointComponent(
new import_orderByQueryExecutionContext.OrderByQueryExecutionContext(
this.clientContext,
this.collectionLink,
this.query,
this.options,
partitionedQueryExecutionInfo,
correlatedActivityId
),
this.emitRawOrderByPayload
);
}
return new import_parallelQueryExecutionContext.ParallelQueryExecutionContext(
this.clientContext,
this.collectionLink,
this.query,
this.options,
partitionedQueryExecutionInfo,
correlatedActivityId
);
}
/**
* Wraps base context with appropriate non-streaming component
*/
wrapWithNonStreamingComponent(baseContext, queryInfo, sortOrders, vectorSearchBufferSize) {
const distinctType = queryInfo.distinctType;
if (distinctType === "None") {
return new import_NonStreamingOrderByEndpointComponent.NonStreamingOrderByEndpointComponent(
baseContext,
sortOrders,
vectorSearchBufferSize,
queryInfo.offset,
this.emitRawOrderByPayload
);
}
return new import_NonStreamingOrderByDistinctEndpointComponent.NonStreamingOrderByDistinctEndpointComponent(
baseContext,
queryInfo,
vectorSearchBufferSize,
this.emitRawOrderByPayload
);
}
/**
* Applies GROUP BY components to the pipeline if needed
*/
applyGroupByComponents(endpoint, queryInfo, isGroupByQuery) {
if (!isGroupByQuery) {
return endpoint;
}
return queryInfo.hasSelectValue ? new import_GroupByValueEndpointComponent.GroupByValueEndpointComponent(endpoint, queryInfo) : new import_GroupByEndpointComponent.GroupByEndpointComponent(endpoint, queryInfo);
}
/**
* Applies DISTINCT components to the pipeline if needed
*/
applyDistinctComponents(endpoint, queryInfo, queryContinuationFields) {
const distinctType = queryInfo.distinctType;
if (distinctType === "Ordered") {
const lastHash = queryContinuationFields?.hashedLastResult;
return new import_OrderedDistinctEndpointComponent.OrderedDistinctEndpointComponent(endpoint, lastHash);
}
if (distinctType === "Unordered") {
return new import_UnorderedDistinctEndpointComponent.UnorderedDistinctEndpointComponent(endpoint);
}
return endpoint;
}
/**
* Applies TOP and OFFSET+LIMIT components to the pipeline if needed
*/
applyLimitComponents(endpoint, queryInfo, queryContinuationFields) {
let top = queryInfo.top;
if (typeof top === "number") {
if (queryContinuationFields?.limit !== void 0) {
top = queryContinuationFields.limit;
}
endpoint = new import_OffsetLimitEndpointComponent.OffsetLimitEndpointComponent(endpoint, 0, top);
}
let limit = queryInfo.limit;
let offset = queryInfo.offset;
if (queryContinuationFields) {
if (queryContinuationFields.limit !== void 0) {
limit = queryContinuationFields.limit;
}
if (queryContinuationFields.offset !== void 0) {
offset = queryContinuationFields.offset;
}
}
if (typeof limit === "number" && typeof offset === "number") {
endpoint = new import_OffsetLimitEndpointComponent.OffsetLimitEndpointComponent(endpoint, offset, limit);
}
return endpoint;
}
/**
* Validates vector search buffer size constraints
*/
validateVectorSearchBufferSize(vectorSearchBufferSize, options) {
const maxBufferSize = options["vectorSearchBufferSize"] ? options["vectorSearchBufferSize"] : import_constants.QueryExecution.DEFAULT_MAX_VECTOR_SEARCH_BUFFER_SIZE;
if (vectorSearchBufferSize > maxBufferSize) {
throw new import_ErrorResponse.ErrorResponse(
`Executing a vector search query with TOP or OFFSET + LIMIT value ${vectorSearchBufferSize} larger than the vectorSearchBufferSize ${maxBufferSize} is not allowed`
);
}
}
calculateVectorSearchBufferSize(queryInfo, options) {
if (queryInfo.top === 0 || queryInfo.limit === 0) return 0;
return queryInfo.top ? queryInfo.top : queryInfo.limit ? queryInfo.offset + queryInfo.limit : options["vectorSearchBufferSize"] && options["vectorSearchBufferSize"] > 0 ? options["vectorSearchBufferSize"] : import_constants.QueryExecution.DEFAULT_MAX_VECTOR_SEARCH_BUFFER_SIZE;
}
checkQueryConstraints(queryInfo) {
const hasTop = queryInfo.top || queryInfo.top === 0;
const hasLimit = queryInfo.limit || queryInfo.limit === 0;
if (!hasTop && !hasLimit) {
throw new import_ErrorResponse.ErrorResponse(
"Executing a non-streaming search query without TOP or LIMIT can consume a large number of RUs very fast and have long runtimes. Please ensure you are using one of the above two filters with your vector search query."
);
}
return;
}
/**
* Analyzes query information and extracts key characteristics for query execution planning
*/
analyzeQueryInfo(queryInfo) {
const sortOrders = queryInfo.orderBy;
const nonStreamingOrderBy = queryInfo.hasNonStreamingOrderBy;
const isOrderByQuery = Array.isArray(sortOrders) && sortOrders.length > 0;
const isGroupByQuery = Object.keys(queryInfo.groupByAliasToAggregateType || {}).length > 0 || (queryInfo.aggregates?.length || 0) > 0 || (queryInfo.groupByExpressions?.length || 0) > 0;
const isUnorderedDistinctQuery = queryInfo.distinctType === "Unordered";
const querySupportsTokens = !isUnorderedDistinctQuery && !isGroupByQuery && !nonStreamingOrderBy;
return {
sortOrders,
nonStreamingOrderBy,
isOrderByQuery,
isGroupByQuery,
isUnorderedDistinctQuery,
querySupportsTokens
};
}
}
// Annotate the CommonJS export names for ESM import in node:
0 && (module.exports = {
PipelinedQueryExecutionContext
});