private async Task <ResponseMessage> ProcessResourceOperationAsBulkStreamAsync(
OperationType operationType,
RequestOptions requestOptions,
ContainerInternal cosmosContainerCore,
PartitionKey partitionKey,
string itemId,
Stream streamPayload,
CancellationToken cancellationToken)
{
this.ThrowIfDisposed();
ItemRequestOptions itemRequestOptions = requestOptions as ItemRequestOptions;
TransactionalBatchItemRequestOptions batchItemRequestOptions = TransactionalBatchItemRequestOptions.FromItemRequestOptions(itemRequestOptions);
ItemBatchOperation itemBatchOperation = new ItemBatchOperation(
operationType: operationType,
operationIndex: 0,
partitionKey: partitionKey,
id: itemId,
resourceStream: streamPayload,
requestOptions: batchItemRequestOptions,
cosmosClientContext: this);
TransactionalBatchOperationResult batchOperationResult = await cosmosContainerCore.BatchExecutor.AddAsync(
itemBatchOperation,
itemRequestOptions,
cancellationToken);
return(batchOperationResult.ToResponseMessage());
}
private static bool ValidateOperationEPK(
ItemBatchOperation operation,
ItemRequestOptions itemRequestOptions)
{
if (itemRequestOptions.Properties != null &&
(itemRequestOptions.Properties.TryGetValue(WFConstants.BackendHeaders.EffectivePartitionKey, out object epkObj)
| itemRequestOptions.Properties.TryGetValue(WFConstants.BackendHeaders.EffectivePartitionKeyString, out object epkStrObj)
| itemRequestOptions.Properties.TryGetValue(HttpConstants.HttpHeaders.PartitionKey, out object pkStringObj)))
{
byte[] epk = epkObj as byte[];
string epkStr = epkStrObj as string;
string pkString = pkStringObj as string;
if ((epk == null && pkString == null) || epkStr == null)
{
throw new InvalidOperationException(string.Format(
ClientResources.EpkPropertiesPairingExpected,
WFConstants.BackendHeaders.EffectivePartitionKey,
WFConstants.BackendHeaders.EffectivePartitionKeyString));
}
if (operation.PartitionKey != null)
{
throw new InvalidOperationException(ClientResources.PKAndEpkSetTogether);
}
}
return(true);
}
internal virtual async Task ValidateOperationAsync(
ItemBatchOperation operation,
ItemRequestOptions itemRequestOptions = null,
CancellationToken cancellationToken = default(CancellationToken))
{
if (itemRequestOptions != null)
{
if (itemRequestOptions.BaseConsistencyLevel.HasValue ||
itemRequestOptions.PreTriggers != null ||
itemRequestOptions.PostTriggers != null ||
itemRequestOptions.SessionToken != null)
{
throw new InvalidOperationException(ClientResources.UnsupportedBulkRequestOptions);
}
if (itemRequestOptions.DiagnosticContextFactory != null)
{
throw new ArgumentException("DiagnosticContext is not allowed when AllowBulkExecution is set to true");
}
Debug.Assert(BatchAsyncContainerExecutor.ValidateOperationEPK(operation, itemRequestOptions));
}
await operation.EncryptAndMaterializeResourceAsync(this.cosmosClientContext.SerializerCore, cancellationToken);
}
private async Task <ResponseMessage> ProcessResourceOperationAsBulkStreamAsync(
Uri resourceUri,
ResourceType resourceType,
OperationType operationType,
RequestOptions requestOptions,
ContainerCore cosmosContainerCore,
PartitionKey partitionKey,
string itemId,
Stream streamPayload,
CosmosDiagnosticsContext diagnosticsContext,
CancellationToken cancellationToken)
{
ItemRequestOptions itemRequestOptions = requestOptions as ItemRequestOptions;
TransactionalBatchItemRequestOptions batchItemRequestOptions = TransactionalBatchItemRequestOptions.FromItemRequestOptions(itemRequestOptions);
ItemBatchOperation itemBatchOperation = new ItemBatchOperation(
operationType: operationType,
operationIndex: 0,
partitionKey: partitionKey,
id: itemId,
resourceStream: streamPayload,
requestOptions: batchItemRequestOptions,
diagnosticsContext: diagnosticsContext);
TransactionalBatchOperationResult batchOperationResult = await cosmosContainerCore.BatchExecutor.AddAsync(itemBatchOperation, itemRequestOptions, cancellationToken);
return(batchOperationResult.ToResponseMessage());
}
private async Task ReBatchAsync(
ItemBatchOperation operation,
CancellationToken cancellationToken)
{
string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation, cancellationToken).ConfigureAwait(false);
BatchAsyncStreamer streamer = this.GetOrAddStreamerForPartitionKeyRange(resolvedPartitionKeyRangeId);
streamer.Add(operation);
}
private async Task ReBatchAsync(
ItemBatchOperation operation,
CancellationToken cancellationToken)
{
using (ITrace trace = Tracing.Trace.GetRootTrace("Batch Retry Async", TraceComponent.Batch, Tracing.TraceLevel.Verbose))
{
string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation, trace, cancellationToken).ConfigureAwait(false);
operation.Context.ReRouteOperation(resolvedPartitionKeyRangeId, trace);
BatchAsyncStreamer streamer = this.GetOrAddStreamerForPartitionKeyRange(resolvedPartitionKeyRangeId);
streamer.Add(operation);
}
}
private async Task FillOperationPropertiesAsync(ItemBatchOperation operation, CancellationToken cancellationToken)
{
// Same logic from RequestInvokerHandler to manage partition key migration
if (object.ReferenceEquals(operation.PartitionKey, PartitionKey.None))
{
Documents.Routing.PartitionKeyInternal partitionKeyInternal = await this.cosmosContainer.GetNonePartitionKeyValueAsync(cancellationToken).ConfigureAwait(false);
operation.PartitionKeyJson = partitionKeyInternal.ToJsonString();
}
else
{
operation.PartitionKeyJson = operation.PartitionKey.Value.ToString();
}
}
private async Task <string> ResolvePartitionKeyRangeIdAsync(
ItemBatchOperation operation,
CancellationToken cancellationToken)
{
cancellationToken.ThrowIfCancellationRequested();
PartitionKeyDefinition partitionKeyDefinition = await this.cosmosContainer.GetPartitionKeyDefinitionAsync(cancellationToken);
CollectionRoutingMap collectionRoutingMap = await this.cosmosContainer.GetRoutingMapAsync(cancellationToken);
Debug.Assert(operation.RequestOptions?.Properties?.TryGetValue(WFConstants.BackendHeaders.EffectivePartitionKeyString, out object epkObj) == null, "EPK is not supported");
await this.FillOperationPropertiesAsync(operation, cancellationToken);
return(BatchExecUtils.GetPartitionKeyRangeId(operation.PartitionKey.Value, partitionKeyDefinition, collectionRoutingMap));
}
public virtual async Task <TransactionalBatchOperationResult> AddAsync(
ItemBatchOperation operation,
ITrace trace,
ItemRequestOptions itemRequestOptions = null,
CancellationToken cancellationToken = default)
{
if (operation == null)
{
throw new ArgumentNullException(nameof(operation));
}
await this.ValidateOperationAsync(operation, itemRequestOptions, cancellationToken);
string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(
operation,
trace,
cancellationToken).ConfigureAwait(false);
BatchAsyncStreamer streamer = this.GetOrAddStreamerForPartitionKeyRange(resolvedPartitionKeyRangeId);
ItemBatchOperationContext context = new ItemBatchOperationContext(
resolvedPartitionKeyRangeId,
trace,
BatchAsyncContainerExecutor.GetRetryPolicy(this.cosmosContainer, operation.OperationType, this.retryOptions));
if (itemRequestOptions != null && itemRequestOptions.AddRequestHeaders != null)
{
// get the header value if any, passed by the encryption package.
Headers encryptionHeaders = new Headers();
itemRequestOptions.AddRequestHeaders?.Invoke(encryptionHeaders);
// make sure we set the Intended Collection Rid header when we have encrypted payload.
// This primarily would allow CosmosDB Encryption package to detect change in container referenced by a Client
// and prevent creating data with wrong Encryption Policy.
if (encryptionHeaders.TryGetValue(HttpConstants.HttpHeaders.IsClientEncrypted, out string encrypted))
{
context.IsClientEncrypted = bool.Parse(encrypted);
if (context.IsClientEncrypted && encryptionHeaders.TryGetValue(WFConstants.BackendHeaders.IntendedCollectionRid, out string ridValue))
{
context.IntendedCollectionRidValue = ridValue;
}
}
}
operation.AttachContext(context);
streamer.Add(operation);
return(await context.OperationTask);
}
public virtual bool TryAdd(ItemBatchOperation operation)
{
if (this.dispatched)
{
DefaultTrace.TraceCritical($"Add operation attempted on dispatched batch.");
return(false);
}
if (operation == null)
{
throw new ArgumentNullException(nameof(operation));
}
if (operation.Context == null)
{
throw new ArgumentNullException(nameof(operation.Context));
}
if (operation.Context.IsClientEncrypted && !this.isClientEncrypted)
{
this.isClientEncrypted = true;
this.intendedCollectionRidValue = operation.Context.IntendedCollectionRidValue;
}
if (this.batchOperations.Count == this.maxBatchOperationCount)
{
DefaultTrace.TraceInformation($"Batch is full - Max operation count {this.maxBatchOperationCount} reached.");
return(false);
}
int itemByteSize = operation.GetApproximateSerializedLength();
if (this.batchOperations.Count > 0 && itemByteSize + this.currentSize > this.maxBatchByteSize)
{
DefaultTrace.TraceInformation($"Batch is full - Max byte size {this.maxBatchByteSize} reached.");
return(false);
}
this.currentSize += itemByteSize;
// Operation index is in the scope of the current batch
operation.OperationIndex = this.batchOperations.Count;
operation.Context.CurrentBatcher = this;
this.batchOperations.Add(operation);
return(true);
}
private async Task <string> ResolvePartitionKeyRangeIdAsync(
ItemBatchOperation operation,
CancellationToken cancellationToken)
{
cancellationToken.ThrowIfCancellationRequested();
PartitionKeyDefinition partitionKeyDefinition = await this.cosmosContainer.GetPartitionKeyDefinitionAsync(cancellationToken);
CollectionRoutingMap collectionRoutingMap = await this.cosmosContainer.GetRoutingMapAsync(cancellationToken);
Debug.Assert(operation.RequestOptions?.Properties?.TryGetValue(WFConstants.BackendHeaders.EffectivePartitionKeyString, out object epkObj) == null, "EPK is not supported");
Documents.Routing.PartitionKeyInternal partitionKeyInternal = await this.GetPartitionKeyInternalAsync(operation, cancellationToken);
operation.PartitionKeyJson = partitionKeyInternal.ToJsonString();
string effectivePartitionKeyString = partitionKeyInternal.GetEffectivePartitionKeyString(partitionKeyDefinition);
return(collectionRoutingMap.GetRangeByEffectivePartitionKey(effectivePartitionKeyString).Id);
}
private async Task <ResponseMessage> ProcessResourceOperationAsBulkStreamAsync(
Uri resourceUri,
ResourceType resourceType,
OperationType operationType,
RequestOptions requestOptions,
ContainerCore cosmosContainerCore,
PartitionKey partitionKey,
string itemId,
Stream streamPayload,
Action <RequestMessage> requestEnricher,
CancellationToken cancellationToken)
{
ItemRequestOptions itemRequestOptions = requestOptions as ItemRequestOptions;
TransactionalBatchItemRequestOptions batchItemRequestOptions = TransactionalBatchItemRequestOptions.FromItemRequestOptions(itemRequestOptions);
ItemBatchOperation itemBatchOperation = new ItemBatchOperation(operationType, /* index */ 0, partitionKey, itemId, streamPayload, batchItemRequestOptions);
TransactionalBatchOperationResult batchOperationResult = await cosmosContainerCore.BatchExecutor.AddAsync(itemBatchOperation, itemRequestOptions, cancellationToken);
return(batchOperationResult.ToResponseMessage());
}
public void Add(ItemBatchOperation operation)
{
BatchAsyncBatcher toDispatch = null;
lock (this.dispatchLimiter)
{
while (!this.currentBatcher.TryAdd(operation))
{
// Batcher is full
toDispatch = this.GetBatchToDispatchAndCreate();
}
}
if (toDispatch != null)
{
// Discarded for Fire & Forget
_ = toDispatch.DispatchAsync(this.partitionMetric, this.cancellationTokenSource.Token);
}
}
internal virtual async Task ValidateOperationAsync(
ItemBatchOperation operation,
ItemRequestOptions itemRequestOptions = null,
CancellationToken cancellationToken = default)
{
if (itemRequestOptions != null)
{
if (itemRequestOptions.BaseConsistencyLevel.HasValue ||
itemRequestOptions.PreTriggers != null ||
itemRequestOptions.PostTriggers != null ||
itemRequestOptions.SessionToken != null)
{
throw new InvalidOperationException(ClientResources.UnsupportedBulkRequestOptions);
}
Debug.Assert(BatchAsyncContainerExecutor.ValidateOperationEPK(operation, itemRequestOptions));
}
await operation.MaterializeResourceAsync(this.cosmosClientContext.SerializerCore, cancellationToken);
}
public virtual async Task <TransactionalBatchOperationResult> AddAsync(
ItemBatchOperation operation,
ItemRequestOptions itemRequestOptions = null,
CancellationToken cancellationToken = default)
{
if (operation == null)
{
throw new ArgumentNullException(nameof(operation));
}
await this.ValidateOperationAsync(operation, itemRequestOptions, cancellationToken);
string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation, cancellationToken).ConfigureAwait(false);
BatchAsyncStreamer streamer = this.GetOrAddStreamerForPartitionKeyRange(resolvedPartitionKeyRangeId);
ItemBatchOperationContext context = new ItemBatchOperationContext(resolvedPartitionKeyRangeId, BatchAsyncContainerExecutor.GetRetryPolicy(this.cosmosContainer, operation.OperationType, this.retryOptions));
operation.AttachContext(context);
streamer.Add(operation);
return(await context.OperationTask);
}
public async Task <BatchOperationResult> AddAsync(
ItemBatchOperation operation,
ItemRequestOptions itemRequestOptions = null,
CancellationToken cancellationToken = default(CancellationToken))
{
if (operation == null)
{
throw new ArgumentNullException(nameof(operation));
}
await this.ValidateOperationAsync(operation, itemRequestOptions, cancellationToken);
string resolvedPartitionKeyRangeId = await this.ResolvePartitionKeyRangeIdAsync(operation, cancellationToken).ConfigureAwait(false);
BatchAsyncStreamer streamer = this.GetOrAddStreamerForPartitionKeyRange(resolvedPartitionKeyRangeId);
ItemBatchOperationContext context = new ItemBatchOperationContext(resolvedPartitionKeyRangeId);
operation.AttachContext(context);
streamer.Add(operation);
return(await context.Task);
}
private Result WriteOperation(long index, out ReadOnlyMemory <byte> buffer)
{
if (this.bodyStream.Length > this.maxBodyLength)
{
// If there is only one operation within the request, we will keep it even if it
// exceeds the maximum size allowed for the body.
if (index > 1)
{
this.shouldDeleteLastWrittenRecord = true;
}
buffer = default(ReadOnlyMemory <byte>);
return(Result.Success);
}
this.bodyStreamPositionBeforeWritingCurrentRecord = this.bodyStream.Length;
if (index >= this.operations.Count)
{
buffer = default(ReadOnlyMemory <byte>);
return(Result.Success);
}
ItemBatchOperation operation = this.operations.Array[this.operations.Offset + (int)index];
RowBuffer row = new RowBuffer(this.operationResizableWriteBuffer.Memory.Length, this.operationResizableWriteBuffer);
row.InitLayout(HybridRowVersion.V1, BatchSchemaProvider.BatchOperationLayout, BatchSchemaProvider.BatchLayoutResolver);
Result r = RowWriter.WriteBuffer(ref row, operation, ItemBatchOperation.WriteOperation);
if (r != Result.Success)
{
buffer = null;
return(r);
}
this.lastWrittenOperationIndex = (int)index;
buffer = this.operationResizableWriteBuffer.Memory.Slice(0, row.Length);
return(Result.Success);
}
internal static Result WriteOperation(ref RowWriter writer, TypeArgument typeArg, ItemBatchOperation operation)
{
bool pkWritten = false;
Result r = writer.WriteInt32("operationType", (int)operation.OperationType);
if (r != Result.Success)
{
return(r);
}
r = writer.WriteInt32("resourceType", (int)ResourceType.Document);
if (r != Result.Success)
{
return(r);
}
if (operation.PartitionKeyJson != null)
{
r = writer.WriteString("partitionKey", operation.PartitionKeyJson);
if (r != Result.Success)
{
return(r);
}
pkWritten = true;
}
if (operation.Id != null)
{
r = writer.WriteString("id", operation.Id);
if (r != Result.Success)
{
return(r);
}
}
if (!operation.ResourceBody.IsEmpty)
{
r = writer.WriteBinary("resourceBody", operation.ResourceBody.Span);
if (r != Result.Success)
{
return(r);
}
}
if (operation.RequestOptions != null)
{
TransactionalBatchItemRequestOptions options = operation.RequestOptions;
if (options.IndexingDirective.HasValue)
{
string indexingDirectiveString = IndexingDirectiveStrings.FromIndexingDirective(options.IndexingDirective.Value);
r = writer.WriteString("indexingDirective", indexingDirectiveString);
if (r != Result.Success)
{
return(r);
}
}
if (options.IfMatchEtag != null)
{
r = writer.WriteString("ifMatch", options.IfMatchEtag);
if (r != Result.Success)
{
return(r);
}
}
else if (options.IfNoneMatchEtag != null)
{
r = writer.WriteString("ifNoneMatch", options.IfNoneMatchEtag);
if (r != Result.Success)
{
return(r);
}
}
if (options.Properties != null)
{
if (options.Properties.TryGetValue(WFConstants.BackendHeaders.BinaryId, out object binaryIdObj))
{
byte[] binaryId = binaryIdObj as byte[];
if (binaryId != null)
{
r = writer.WriteBinary("binaryId", binaryId);
if (r != Result.Success)
{
return(r);
}
}
}
if (options.Properties.TryGetValue(WFConstants.BackendHeaders.EffectivePartitionKey, out object epkObj))
{
byte[] epk = epkObj as byte[];
if (epk != null)
{
r = writer.WriteBinary("effectivePartitionKey", epk);
if (r != Result.Success)
{
return(r);
}
}
}
if (!pkWritten && options.Properties.TryGetValue(
HttpConstants.HttpHeaders.PartitionKey,
out object pkStrObj))
{
string pkString = pkStrObj as string;
if (pkString != null)
{
r = writer.WriteString("partitionKey", pkString);
if (r != Result.Success)
{
return(r);
}
}
}
if (options.Properties.TryGetValue(WFConstants.BackendHeaders.TimeToLiveInSeconds, out object ttlObj))
{
string ttlStr = ttlObj as string;
if (ttlStr != null && int.TryParse(ttlStr, out int ttl))
{
r = writer.WriteInt32("timeToLiveInSeconds", ttl);
if (r != Result.Success)
{
return(r);
}
}
}
}
}
return(Result.Success);
}
private async Task <Documents.Routing.PartitionKeyInternal> GetPartitionKeyInternalAsync(ItemBatchOperation operation, CancellationToken cancellationToken)
{
Debug.Assert(operation.PartitionKey.HasValue, "PartitionKey should be set on the operation");
if (operation.PartitionKey.Value.IsNone)
{
return(await this.cosmosContainer.GetNonePartitionKeyValueAsync(NoOpTrace.Singleton, cancellationToken).ConfigureAwait(false));
}
return(operation.PartitionKey.Value.InternalKey);
}
