Skip to content

Commit

Permalink
Added CosmosRequestContext (#40964)
Browse files Browse the repository at this point in the history
* Added CosmosRequestContext reacting to api review

* removed unnecessary imports

* updated javadoc

* revert change

* initialized in bridgeInternal

* Reacting to api comments

* added missing javadocs

* added missing javadocs

* revert deleting readonly

* reacting to api comments

* reacting to api comments

* reacting to api comments
  • Loading branch information
tvaron3 committed Jul 2, 2024
1 parent 9e8ab07 commit 16a527c
Show file tree
Hide file tree
Showing 13 changed files with 373 additions and 77 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ private static void createReadDeleteBatchEtcOptions(String operationType, String
cosmosRequestOptions.setCosmosEndToEndLatencyPolicyConfig(new CosmosEndToEndOperationLatencyPolicyConfig(true,
Duration.ofSeconds(Long.parseLong(prop.getProperty(E2E_TIMEOUT))),
new ThresholdBasedAvailabilityStrategy()))
.setThresholds(new CosmosDiagnosticsThresholds().setRequestChargeThreshold(Float.parseFloat(prop.getProperty(REQUEST_CHARGE_THRESHOLD))))
.setDiagnosticsThresholds(new CosmosDiagnosticsThresholds().setRequestChargeThreshold(Float.parseFloat(prop.getProperty(REQUEST_CHARGE_THRESHOLD))))
.setConsistencyLevel(ConsistencyLevel.fromServiceSerializedFormat(prop.getProperty(CONSISTENCY_LEVEL)))
.setContentResponseOnWriteEnabled(Boolean.parseBoolean(prop.getProperty(CONTENT_RESPONSE_ON_WRITE)))
.setNonIdempotentWriteRetriesEnabled(Boolean.parseBoolean(prop.getProperty(NON_IDEMPOTENT_WRITE_RETRIES)))
Expand All @@ -112,7 +112,7 @@ private static void createQueryReadAllItemsOptions(String operationType, String
cosmosRequestOptions.setCosmosEndToEndLatencyPolicyConfig(new CosmosEndToEndOperationLatencyPolicyConfig(true,
Duration.ofSeconds(Long.parseLong(prop.getProperty(E2E_TIMEOUT))),
new ThresholdBasedAvailabilityStrategy()))
.setThresholds(new CosmosDiagnosticsThresholds().setRequestChargeThreshold(Float.parseFloat(prop.getProperty(REQUEST_CHARGE_THRESHOLD))))
.setDiagnosticsThresholds(new CosmosDiagnosticsThresholds().setRequestChargeThreshold(Float.parseFloat(prop.getProperty(REQUEST_CHARGE_THRESHOLD))))
.setThroughputControlGroupName(prop.getProperty(THROUGHPUT_CONTROL_GROUP_NAME))
.setDedicatedGatewayRequestOptions(new DedicatedGatewayRequestOptions()
.setIntegratedCacheBypassed(Boolean.parseBoolean(prop.getProperty(BYPASS_CACHE))))
Expand All @@ -135,7 +135,7 @@ private static void createReadManyOptions(String spanName, CosmosRequestOptions
cosmosRequestOptions.setCosmosEndToEndLatencyPolicyConfig(new CosmosEndToEndOperationLatencyPolicyConfig(true,
Duration.ofSeconds(Long.parseLong(prop.getProperty(E2E_TIMEOUT))),
new ThresholdBasedAvailabilityStrategy()))
.setThresholds(new CosmosDiagnosticsThresholds().setRequestChargeThreshold(Float.parseFloat(prop.getProperty(REQUEST_CHARGE_THRESHOLD))))
.setDiagnosticsThresholds(new CosmosDiagnosticsThresholds().setRequestChargeThreshold(Float.parseFloat(prop.getProperty(REQUEST_CHARGE_THRESHOLD))))
.setThroughputControlGroupName(prop.getProperty(THROUGHPUT_CONTROL_GROUP_NAME))
.setDedicatedGatewayRequestOptions(new DedicatedGatewayRequestOptions()
.setIntegratedCacheBypassed(Boolean.parseBoolean(prop.getProperty(BYPASS_CACHE))))
Expand All @@ -158,7 +158,7 @@ private static void createChangeFeedOptions(String spanName, CosmosRequestOption
if (spanName.contains("queryChangeFeed")) {
cosmosRequestOptions.setExcludeRegions((new ArrayList<>(Arrays.asList(prop.getProperty(EXCLUDE_REGIONS).split(",")))))
.setThroughputControlGroupName(prop.getProperty(THROUGHPUT_CONTROL_GROUP_NAME))
.setThresholds(new CosmosDiagnosticsThresholds().setRequestChargeThreshold(Float.parseFloat(prop.getProperty(REQUEST_CHARGE_THRESHOLD))))
.setDiagnosticsThresholds(new CosmosDiagnosticsThresholds().setRequestChargeThreshold(Float.parseFloat(prop.getProperty(REQUEST_CHARGE_THRESHOLD))))
.setMaxPrefetchPageCount(Integer.parseInt(prop.getProperty(MAX_PREFETCH_PAGE_COUNT)))
.setMaxItemCount(Integer.parseInt(prop.getProperty(MAX_ITEM_COUNT)));
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -614,6 +614,7 @@ public static void initializeAllAccessors() {
CosmosClientBuilder.initialize();
CosmosDiagnostics.initialize();
CosmosDiagnosticsContext.initialize();
CosmosRequestContext.initialize();
CosmosException.initialize();
DirectConnectionConfig.initialize();
CosmosAsyncClient.initialize();
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,190 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
package com.azure.cosmos;

import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
import com.azure.cosmos.implementation.OverridableRequestOptions;
import com.azure.cosmos.models.DedicatedGatewayRequestOptions;

import java.util.List;

/**
* Getters for the common request context for operations in CosmosDB.
*/
public final class CosmosRequestContext {

private OverridableRequestOptions requestOptions;

CosmosRequestContext(OverridableRequestOptions requestOptions) {
this.requestOptions = requestOptions;
}

/**
* Gets the CosmosEndToEndLatencyPolicyConfig.
*
* @return the CosmosEndToEndLatencyPolicyConfig. It could be null if not defined or called on an irrelevant operation.
*/
public CosmosEndToEndOperationLatencyPolicyConfig getCosmosEndToEndLatencyPolicyConfig() {
return requestOptions.getCosmosEndToEndLatencyPolicyConfig();
}

/**
* Gets the consistency level.
*
* @return the consistency level. It could be null if not defined or called on an irrelevant operation.
*/
public ConsistencyLevel getConsistencyLevel() {
return requestOptions.getConsistencyLevel();
}

/**
* Gets the content response on write enabled.
*
* @return the content response on write enabled. It could be null if not defined or called on an irrelevant operation.
*/
public Boolean isContentResponseOnWriteEnabled() {
return requestOptions.isContentResponseOnWriteEnabled();
}

/**
* Gets the non idempotent write retries enabled.
*
* @return the non idempotent write retries enabled. It could be null if not defined or called on an irrelevant operation.
*/
public Boolean getNonIdempotentWriteRetriesEnabled() {
return requestOptions.getNonIdempotentWriteRetriesEnabled();
}

/**
* Gets the dedicated gateway request options.
*
* @return the dedicated gateway request options. It could be null if not defined or called on an irrelevant operation.
*/
public DedicatedGatewayRequestOptions getDedicatedGatewayRequestOptions() {
return requestOptions.getDedicatedGatewayRequestOptions();
}

/**
* Gets the excluded regions.
*
* @return the excluded regions.
*/
public List<String> getExcludedRegions() {
return requestOptions.getExcludedRegions();
}

/**
* Gets the resource token.
*
* @return the resource token.
*/
public String getThroughputControlGroupName() {
return requestOptions.getThroughputControlGroupName();
}

/**
* Gets the diagnostics thresholds.
*
* @return the diagnostics thresholds. It could be null if not defined or called on an irrelevant operation.
*/
public CosmosDiagnosticsThresholds getDiagnosticsThresholds() {
return requestOptions.getDiagnosticsThresholds();
}

/**
* Gets the scan in query enabled.
*
* @return the scan in query enabled. It could be null if not defined or called on an irrelevant operation.
*/
public Boolean isScanInQueryEnabled() {
return requestOptions.isScanInQueryEnabled();
}

/**
* Gets the max degree of parallelism.
*
* @return the max degree of parallelism. It could be null if not defined or called on an irrelevant operation.
*/
public Integer getMaxDegreeOfParallelism() {
return requestOptions.getMaxDegreeOfParallelism();
}

/**
* Gets the max buffered item count.
*
* @return the max buffered item count. It could be null if not defined or called on an irrelevant operation.
*/
public Integer getMaxBufferedItemCount() {
return requestOptions.getMaxBufferedItemCount();
}

/**
* Gets the response continuation token limit in KB.
*
* @return the response continuation token limit in KB. It could be null if not defined or called on an irrelevant operation.
*/
public Integer getResponseContinuationTokenLimitInKb() {
return requestOptions.getResponseContinuationTokenLimitInKb();
}

/**
* Gets the max item count.
*
* @return the max item count. It could be null if not defined or called on an irrelevant operation.
*/
public Integer getMaxItemCount() {
return requestOptions.getMaxItemCount();
}

/**
* Gets the query metrics enabled.
*
* @return the query metrics enabled. It could be null if not defined or called on an irrelevant operation.
*/
public Boolean isQueryMetricsEnabled() {
return requestOptions.isQueryMetricsEnabled();
}

/**
* Gets the index metrics enabled.
*
* @return the index metrics enabled. It could be null if not defined or called on an irrelevant operation.
*/
public Boolean isIndexMetricsEnabled() {
return requestOptions.isIndexMetricsEnabled();
}

/**
* Gets the query name.
*
* @return the query name. It could be null if not defined or called on an irrelevant operation.
*/
public Integer getMaxPrefetchPageCount() {
return requestOptions.getMaxPrefetchPageCount();
}

/**
* Gets the query name.
* @param defaultQueryName the default query name.
*
* @return the query name. It could be null if not defined or called on an irrelevant operation.
*/
public String getQueryNameOrDefault(String defaultQueryName) {
return requestOptions.getQueryNameOrDefault(defaultQueryName);
}

///////////////////////////////////////////////////////////////////////////////////////////
// the following helper/accessor only helps to access this class outside of this package.//
///////////////////////////////////////////////////////////////////////////////////////////
static void initialize() {
ImplementationBridgeHelpers.CosmosRequestContextHelper
.setCosmosRequestContextAccessor(
new ImplementationBridgeHelpers.CosmosRequestContextHelper.CosmosRequestContextAccessor() {
@Override
public CosmosRequestContext create(OverridableRequestOptions requestOptions) {
return new CosmosRequestContext(requestOptions);
}
}
);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,9 @@
import com.azure.cosmos.implementation.batch.BatchRequestResponseConstants;
import com.azure.cosmos.implementation.batch.BulkExecutorDiagnosticsTracker;
import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple;
import com.azure.cosmos.models.CosmosBulkExecutionOptions;
import com.azure.cosmos.models.CosmosBulkExecutionThresholdsState;
import com.azure.cosmos.models.CosmosRequestOptions;
import com.azure.cosmos.models.DedicatedGatewayRequestOptions;
import com.azure.cosmos.models.ReadOnlyRequestOptions;
import reactor.core.scheduler.Scheduler;

import java.time.Duration;
Expand Down Expand Up @@ -313,9 +312,9 @@ public BulkExecutorDiagnosticsTracker getDiagnosticsTracker() {
}

@Override
public void override(ReadOnlyRequestOptions readOnlyRequestOptions) {
this.excludeRegions = overrideOption(readOnlyRequestOptions.getExcludedRegions(), this.excludeRegions);
this.throughputControlGroupName = overrideOption(readOnlyRequestOptions.getThroughputControlGroupName(), this.throughputControlGroupName);
public void override(CosmosRequestOptions cosmosRequestOptions) {
this.excludeRegions = overrideOption(cosmosRequestOptions.getExcludedRegions(), this.excludeRegions);
this.throughputControlGroupName = overrideOption(cosmosRequestOptions.getThroughputControlGroupName(), this.throughputControlGroupName);
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,9 @@
import com.azure.cosmos.implementation.changefeed.common.ChangeFeedState;
import com.azure.cosmos.implementation.feedranges.FeedRangeInternal;
import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple;
import com.azure.cosmos.models.CosmosRequestOptions;
import com.azure.cosmos.models.DedicatedGatewayRequestOptions;
import com.azure.cosmos.models.FeedRange;
import com.azure.cosmos.models.ReadOnlyRequestOptions;
import com.azure.cosmos.util.Beta;

import java.util.HashMap;
Expand Down Expand Up @@ -327,12 +327,12 @@ private void addCustomOptionsForFullFidelityMode() {
}

@Override
public void override(ReadOnlyRequestOptions readOnlyRequestOptions) {
this.maxItemCount = overrideOption(readOnlyRequestOptions.getMaxItemCount(), this.maxItemCount);
this.maxPrefetchPageCount = overrideOption(readOnlyRequestOptions.getMaxPrefetchPageCount(), this.maxPrefetchPageCount);
this.excludeRegions = overrideOption(readOnlyRequestOptions.getExcludedRegions(), this.excludeRegions);
this.throughputControlGroupName = overrideOption(readOnlyRequestOptions.getThroughputControlGroupName(), this.throughputControlGroupName);
this.thresholds = overrideOption(readOnlyRequestOptions.getDiagnosticsThresholds(), this.thresholds);
public void override(CosmosRequestOptions cosmosRequestOptions) {
this.maxItemCount = overrideOption(cosmosRequestOptions.getMaxItemCount(), this.maxItemCount);
this.maxPrefetchPageCount = overrideOption(cosmosRequestOptions.getMaxPrefetchPageCount(), this.maxPrefetchPageCount);
this.excludeRegions = overrideOption(cosmosRequestOptions.getExcludedRegions(), this.excludeRegions);
this.throughputControlGroupName = overrideOption(cosmosRequestOptions.getThroughputControlGroupName(), this.throughputControlGroupName);
this.thresholds = overrideOption(cosmosRequestOptions.getDiagnosticsThresholds(), this.thresholds);
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,8 @@
import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList;
import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple;
import com.azure.cosmos.models.CosmosQueryRequestOptions;
import com.azure.cosmos.models.CosmosRequestOptions;
import com.azure.cosmos.models.DedicatedGatewayRequestOptions;
import com.azure.cosmos.models.ReadOnlyRequestOptions;

import java.time.Duration;
import java.util.HashMap;
Expand Down Expand Up @@ -457,16 +457,16 @@ public T setCustomItemSerializer(CosmosItemSerializer customItemSerializer) {
}

@Override
public void override(ReadOnlyRequestOptions readOnlyRequestOptions) {
this.consistencyLevel = overrideOption(readOnlyRequestOptions.getConsistencyLevel(), this.consistencyLevel);
this.throughputControlGroupName = overrideOption(readOnlyRequestOptions.getThroughputControlGroupName(), this.throughputControlGroupName);
this.dedicatedGatewayRequestOptions = overrideOption(readOnlyRequestOptions.getDedicatedGatewayRequestOptions(), this.dedicatedGatewayRequestOptions);
this.cosmosEndToEndOperationLatencyPolicyConfig = overrideOption(readOnlyRequestOptions.getCosmosEndToEndLatencyPolicyConfig(), this.cosmosEndToEndOperationLatencyPolicyConfig);
this.excludeRegions = overrideOption(readOnlyRequestOptions.getExcludedRegions(), this.excludeRegions);
this.thresholds = overrideOption(readOnlyRequestOptions.getDiagnosticsThresholds(), this.thresholds);
this.indexMetricsEnabled = overrideOption(readOnlyRequestOptions.isIndexMetricsEnabled(), this.indexMetricsEnabled);
this.queryMetricsEnabled = overrideOption(readOnlyRequestOptions.isQueryMetricsEnabled(), this.queryMetricsEnabled);
this.responseContinuationTokenLimitInKb = overrideOption(readOnlyRequestOptions.getResponseContinuationTokenLimitInKb(), this.responseContinuationTokenLimitInKb);
public void override(CosmosRequestOptions cosmosRequestOptions) {
this.consistencyLevel = overrideOption(cosmosRequestOptions.getConsistencyLevel(), this.consistencyLevel);
this.throughputControlGroupName = overrideOption(cosmosRequestOptions.getThroughputControlGroupName(), this.throughputControlGroupName);
this.dedicatedGatewayRequestOptions = overrideOption(cosmosRequestOptions.getDedicatedGatewayRequestOptions(), this.dedicatedGatewayRequestOptions);
this.cosmosEndToEndOperationLatencyPolicyConfig = overrideOption(cosmosRequestOptions.getCosmosEndToEndLatencyPolicyConfig(), this.cosmosEndToEndOperationLatencyPolicyConfig);
this.excludeRegions = overrideOption(cosmosRequestOptions.getExcludedRegions(), this.excludeRegions);
this.thresholds = overrideOption(cosmosRequestOptions.getDiagnosticsThresholds(), this.thresholds);
this.indexMetricsEnabled = overrideOption(cosmosRequestOptions.isIndexMetricsEnabled(), this.indexMetricsEnabled);
this.queryMetricsEnabled = overrideOption(cosmosRequestOptions.isQueryMetricsEnabled(), this.queryMetricsEnabled);
this.responseContinuationTokenLimitInKb = overrideOption(cosmosRequestOptions.getResponseContinuationTokenLimitInKb(), this.responseContinuationTokenLimitInKb);
}

public RequestOptions applyToRequestOptions(RequestOptions requestOptions) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,10 @@
package com.azure.cosmos.implementation;

import com.azure.cosmos.CosmosDiagnostics;
import com.azure.cosmos.models.CosmosRequestOptions;
import com.azure.cosmos.models.FeedRange;
import com.azure.cosmos.models.PartitionKey;
import com.azure.cosmos.models.PartitionKeyDefinition;
import com.azure.cosmos.models.ReadOnlyRequestOptions;

import java.util.ArrayList;
import java.util.List;
Expand Down Expand Up @@ -379,12 +379,12 @@ public PartitionKeyDefinition getPartitionKeyDefinition() {
}

@Override
public void override(ReadOnlyRequestOptions readOnlyRequestOptions) {
super.override(readOnlyRequestOptions);
this.scanInQueryEnabled = overrideOption(readOnlyRequestOptions.isScanInQueryEnabled(), this.scanInQueryEnabled);
this.maxDegreeOfParallelism = overrideOption(readOnlyRequestOptions.getMaxDegreeOfParallelism(), this.maxDegreeOfParallelism);
this.maxBufferedItemCount = overrideOption(readOnlyRequestOptions.getMaxBufferedItemCount(), this.maxBufferedItemCount);
this.maxItemCount = overrideOption(readOnlyRequestOptions.getMaxItemCount(), this.maxItemCount);
this.queryName = overrideOption(readOnlyRequestOptions.getQueryNameOrDefault(""), this.queryName);
public void override(CosmosRequestOptions cosmosRequestOptions) {
super.override(cosmosRequestOptions);
this.scanInQueryEnabled = overrideOption(cosmosRequestOptions.isScanInQueryEnabled(), this.scanInQueryEnabled);
this.maxDegreeOfParallelism = overrideOption(cosmosRequestOptions.getMaxDegreeOfParallelism(), this.maxDegreeOfParallelism);
this.maxBufferedItemCount = overrideOption(cosmosRequestOptions.getMaxBufferedItemCount(), this.maxBufferedItemCount);
this.maxItemCount = overrideOption(cosmosRequestOptions.getMaxItemCount(), this.maxItemCount);
this.queryName = overrideOption(cosmosRequestOptions.getQueryName(), this.queryName);
}
}
Loading

0 comments on commit 16a527c

Please sign in to comment.