Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance.
Project price only 1 $
You can buy this project and download/modify it how often you want.
com.azure.cosmos.implementation.RxDocumentClientImpl Maven / Gradle / Ivy
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
package com.azure.cosmos.implementation;
import com.azure.core.credential.AzureKeyCredential;
import com.azure.core.credential.SimpleTokenCache;
import com.azure.core.credential.TokenCredential;
import com.azure.core.credential.TokenRequestContext;
import com.azure.cosmos.BridgeInternal;
import com.azure.cosmos.ConnectionMode;
import com.azure.cosmos.ConsistencyLevel;
import com.azure.cosmos.CosmosAsyncClient;
import com.azure.cosmos.CosmosContainerProactiveInitConfig;
import com.azure.cosmos.CosmosDiagnostics;
import com.azure.cosmos.CosmosDiagnosticsContext;
import com.azure.cosmos.CosmosEndToEndOperationLatencyPolicyConfig;
import com.azure.cosmos.CosmosException;
import com.azure.cosmos.CosmosItemSerializer;
import com.azure.cosmos.CosmosOperationPolicy;
import com.azure.cosmos.DirectConnectionConfig;
import com.azure.cosmos.SessionRetryOptions;
import com.azure.cosmos.ThresholdBasedAvailabilityStrategy;
import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
import com.azure.cosmos.implementation.apachecommons.lang.tuple.ImmutablePair;
import com.azure.cosmos.implementation.batch.BatchResponseParser;
import com.azure.cosmos.implementation.batch.PartitionKeyRangeServerBatchRequest;
import com.azure.cosmos.implementation.batch.ServerBatchRequest;
import com.azure.cosmos.implementation.batch.SinglePartitionKeyServerBatchRequest;
import com.azure.cosmos.implementation.caches.RxClientCollectionCache;
import com.azure.cosmos.implementation.caches.RxCollectionCache;
import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache;
import com.azure.cosmos.implementation.circuitBreaker.GlobalPartitionEndpointManagerForCircuitBreaker;
import com.azure.cosmos.implementation.circuitBreaker.PartitionKeyRangeWrapper;
import com.azure.cosmos.implementation.clienttelemetry.ClientTelemetry;
import com.azure.cosmos.implementation.cpu.CpuMemoryListener;
import com.azure.cosmos.implementation.cpu.CpuMemoryMonitor;
import com.azure.cosmos.implementation.directconnectivity.AddressSelector;
import com.azure.cosmos.implementation.directconnectivity.GatewayServiceConfigurationReader;
import com.azure.cosmos.implementation.directconnectivity.GlobalAddressResolver;
import com.azure.cosmos.implementation.directconnectivity.ServerStoreModel;
import com.azure.cosmos.implementation.directconnectivity.StoreClient;
import com.azure.cosmos.implementation.directconnectivity.StoreClientFactory;
import com.azure.cosmos.implementation.faultinjection.IFaultInjectorProvider;
import com.azure.cosmos.implementation.feedranges.FeedRangeEpkImpl;
import com.azure.cosmos.implementation.http.HttpClient;
import com.azure.cosmos.implementation.http.HttpClientConfig;
import com.azure.cosmos.implementation.http.HttpHeaders;
import com.azure.cosmos.implementation.http.SharedGatewayHttpClient;
import com.azure.cosmos.implementation.patch.PatchUtil;
import com.azure.cosmos.implementation.query.DocumentQueryExecutionContextFactory;
import com.azure.cosmos.implementation.query.IDocumentQueryClient;
import com.azure.cosmos.implementation.query.IDocumentQueryExecutionContext;
import com.azure.cosmos.implementation.query.Paginator;
import com.azure.cosmos.implementation.query.PartitionedQueryExecutionInfo;
import com.azure.cosmos.implementation.query.PipelinedQueryExecutionContextBase;
import com.azure.cosmos.implementation.query.QueryInfo;
import com.azure.cosmos.implementation.routing.CollectionRoutingMap;
import com.azure.cosmos.implementation.routing.PartitionKeyAndResourceTokenPair;
import com.azure.cosmos.implementation.routing.PartitionKeyInternal;
import com.azure.cosmos.implementation.routing.PartitionKeyInternalHelper;
import com.azure.cosmos.implementation.routing.PartitionKeyRangeIdentity;
import com.azure.cosmos.implementation.routing.Range;
import com.azure.cosmos.implementation.routing.RegionNameToRegionIdMap;
import com.azure.cosmos.implementation.spark.OperationContext;
import com.azure.cosmos.implementation.spark.OperationContextAndListenerTuple;
import com.azure.cosmos.implementation.spark.OperationListener;
import com.azure.cosmos.implementation.throughputControl.ThroughputControlStore;
import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal;
import com.azure.cosmos.models.CosmosAuthorizationTokenResolver;
import com.azure.cosmos.models.CosmosBatchResponse;
import com.azure.cosmos.models.CosmosChangeFeedRequestOptions;
import com.azure.cosmos.models.CosmosClientTelemetryConfig;
import com.azure.cosmos.models.CosmosContainerIdentity;
import com.azure.cosmos.models.CosmosItemIdentity;
import com.azure.cosmos.models.CosmosItemRequestOptions;
import com.azure.cosmos.models.CosmosItemResponse;
import com.azure.cosmos.models.CosmosOperationDetails;
import com.azure.cosmos.models.CosmosPatchOperations;
import com.azure.cosmos.models.CosmosQueryRequestOptions;
import com.azure.cosmos.models.FeedRange;
import com.azure.cosmos.models.FeedResponse;
import com.azure.cosmos.models.ModelBridgeInternal;
import com.azure.cosmos.models.PartitionKey;
import com.azure.cosmos.models.PartitionKeyDefinition;
import com.azure.cosmos.models.PartitionKind;
import com.azure.cosmos.models.SqlParameter;
import com.azure.cosmos.models.SqlQuerySpec;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.Exceptions;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.SignalType;
import reactor.util.concurrent.Queues;
import reactor.util.function.Tuple2;
import reactor.util.retry.Retry;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.lang.management.ManagementFactory;
import java.net.URI;
import java.net.URLEncoder;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import static com.azure.cosmos.BridgeInternal.getAltLink;
import static com.azure.cosmos.BridgeInternal.toResourceResponse;
import static com.azure.cosmos.BridgeInternal.toStoredProcedureResponse;
import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument;
import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
/**
* While this class is public, it is not part of our published public APIs.
* This is meant to be internally used only by our sdk.
*/
public class RxDocumentClientImpl implements AsyncDocumentClient, IAuthorizationTokenProvider, CpuMemoryListener,
DiagnosticsClientContext {
private final static List EMPTY_REGION_LIST = Collections.emptyList();
private final static List EMPTY_ENDPOINT_LIST = Collections.emptyList();
private final static
ImplementationBridgeHelpers.CosmosDiagnosticsHelper.CosmosDiagnosticsAccessor diagnosticsAccessor =
ImplementationBridgeHelpers.CosmosDiagnosticsHelper.getCosmosDiagnosticsAccessor();
private final static
ImplementationBridgeHelpers.FeedResponseHelper.FeedResponseAccessor feedResponseAccessor =
ImplementationBridgeHelpers.FeedResponseHelper.getFeedResponseAccessor();
private final static
ImplementationBridgeHelpers.CosmosClientTelemetryConfigHelper.CosmosClientTelemetryConfigAccessor telemetryCfgAccessor =
ImplementationBridgeHelpers.CosmosClientTelemetryConfigHelper.getCosmosClientTelemetryConfigAccessor();
private final static
ImplementationBridgeHelpers.CosmosDiagnosticsContextHelper.CosmosDiagnosticsContextAccessor ctxAccessor =
ImplementationBridgeHelpers.CosmosDiagnosticsContextHelper.getCosmosDiagnosticsContextAccessor();
private final static
ImplementationBridgeHelpers.CosmosQueryRequestOptionsHelper.CosmosQueryRequestOptionsAccessor qryOptAccessor =
ImplementationBridgeHelpers.CosmosQueryRequestOptionsHelper.getCosmosQueryRequestOptionsAccessor();
private final static
ImplementationBridgeHelpers.CosmosItemResponseHelper.CosmosItemResponseBuilderAccessor itemResponseAccessor =
ImplementationBridgeHelpers.CosmosItemResponseHelper.getCosmosItemResponseBuilderAccessor();
private static final ImplementationBridgeHelpers.CosmosChangeFeedRequestOptionsHelper.CosmosChangeFeedRequestOptionsAccessor changeFeedOptionsAccessor =
ImplementationBridgeHelpers.CosmosChangeFeedRequestOptionsHelper.getCosmosChangeFeedRequestOptionsAccessor();
private static final ImplementationBridgeHelpers.CosmosOperationDetailsHelper.CosmosOperationDetailsAccessor operationDetailsAccessor =
ImplementationBridgeHelpers.CosmosOperationDetailsHelper.getCosmosOperationDetailsAccessor();
private static final String tempMachineId = "uuid:" + UUID.randomUUID();
private static final AtomicInteger activeClientsCnt = new AtomicInteger(0);
private static final Map clientMap = new ConcurrentHashMap<>();
private static final AtomicInteger clientIdGenerator = new AtomicInteger(0);
private static final Range RANGE_INCLUDING_ALL_PARTITION_KEY_RANGES = new Range<>(
PartitionKeyInternalHelper.MinimumInclusiveEffectivePartitionKey,
PartitionKeyInternalHelper.MaximumExclusiveEffectivePartitionKey, true, false);
private static final String DUMMY_SQL_QUERY = "this is dummy and only used in creating " +
"ParallelDocumentQueryExecutioncontext, but not used";
private final static ObjectMapper mapper = Utils.getSimpleObjectMapper();
private final CosmosItemSerializer defaultCustomSerializer;
private final static Logger logger = LoggerFactory.getLogger(RxDocumentClientImpl.class);
private final String masterKeyOrResourceToken;
private final URI serviceEndpoint;
private final ConnectionPolicy connectionPolicy;
private final ConsistencyLevel consistencyLevel;
private final BaseAuthorizationTokenProvider authorizationTokenProvider;
private final UserAgentContainer userAgentContainer;
private final boolean hasAuthKeyResourceToken;
private final Configs configs;
private final boolean connectionSharingAcrossClientsEnabled;
private AzureKeyCredential credential;
private final TokenCredential tokenCredential;
private String[] tokenCredentialScopes;
private SimpleTokenCache tokenCredentialCache;
private CosmosAuthorizationTokenResolver cosmosAuthorizationTokenResolver;
AuthorizationTokenType authorizationTokenType;
private ISessionContainer sessionContainer;
private String firstResourceTokenFromPermissionFeed = StringUtils.EMPTY;
private RxClientCollectionCache collectionCache;
private RxGatewayStoreModel gatewayProxy;
private RxStoreModel storeModel;
private GlobalAddressResolver addressResolver;
private RxPartitionKeyRangeCache partitionKeyRangeCache;
private Map> resourceTokensMap;
private final boolean contentResponseOnWriteEnabled;
private final Map queryPlanCache;
private final AtomicBoolean closed = new AtomicBoolean(false);
private final int clientId;
private ClientTelemetry clientTelemetry;
private final ApiType apiType;
private final CosmosEndToEndOperationLatencyPolicyConfig cosmosEndToEndOperationLatencyPolicyConfig;
private final AtomicReference mostRecentlyCreatedDiagnostics = new AtomicReference<>(null);
// RetryPolicy retries a request when it encounters session unavailable (see ClientRetryPolicy).
// Once it exhausts all write regions it clears the session container, then it uses RxClientCollectionCache
// to resolves the request's collection name. If it differs from the session container's resource id it
// explains the session unavailable exception: somebody removed and recreated the collection. In this
// case we retry once again (with empty session token) otherwise we return the error to the client
// (see RenameCollectionAwareClientRetryPolicy)
private IRetryPolicyFactory resetSessionTokenRetryPolicy;
/**
* Compatibility mode: Allows to specify compatibility mode used by client when
* making query requests. Should be removed when application/sql is no longer
* supported.
*/
private final QueryCompatibilityMode queryCompatibilityMode = QueryCompatibilityMode.Default;
private final GlobalEndpointManager globalEndpointManager;
private final GlobalPartitionEndpointManagerForCircuitBreaker globalPartitionEndpointManagerForCircuitBreaker;
private final RetryPolicy retryPolicy;
private HttpClient reactorHttpClient;
private Function httpClientInterceptor;
private volatile boolean useMultipleWriteLocations;
// creator of TransportClient is responsible for disposing it.
private StoreClientFactory storeClientFactory;
private GatewayServiceConfigurationReader gatewayConfigurationReader;
private final DiagnosticsClientConfig diagnosticsClientConfig;
private final AtomicBoolean throughputControlEnabled;
private ThroughputControlStore throughputControlStore;
private final CosmosClientTelemetryConfig clientTelemetryConfig;
private final String clientCorrelationId;
private final SessionRetryOptions sessionRetryOptions;
private final boolean sessionCapturingOverrideEnabled;
private final boolean sessionCapturingDisabled;
private final boolean isRegionScopedSessionCapturingEnabledOnClientOrSystemConfig;
private List operationPolicies;
private AtomicReference cachedCosmosAsyncClientSnapshot;
public RxDocumentClientImpl(URI serviceEndpoint,
String masterKeyOrResourceToken,
List permissionFeed,
ConnectionPolicy connectionPolicy,
ConsistencyLevel consistencyLevel,
Configs configs,
CosmosAuthorizationTokenResolver cosmosAuthorizationTokenResolver,
AzureKeyCredential credential,
boolean sessionCapturingOverride,
boolean connectionSharingAcrossClientsEnabled,
boolean contentResponseOnWriteEnabled,
CosmosClientMetadataCachesSnapshot metadataCachesSnapshot,
ApiType apiType,
CosmosClientTelemetryConfig clientTelemetryConfig,
String clientCorrelationId,
CosmosEndToEndOperationLatencyPolicyConfig cosmosEndToEndOperationLatencyPolicyConfig,
SessionRetryOptions sessionRetryOptions,
CosmosContainerProactiveInitConfig containerProactiveInitConfig,
CosmosItemSerializer defaultCustomSerializer,
boolean isRegionScopedSessionCapturingEnabled) {
this(
serviceEndpoint,
masterKeyOrResourceToken,
permissionFeed,
connectionPolicy,
consistencyLevel,
configs,
credential,
null,
sessionCapturingOverride,
connectionSharingAcrossClientsEnabled,
contentResponseOnWriteEnabled,
metadataCachesSnapshot,
apiType,
clientTelemetryConfig,
clientCorrelationId,
cosmosEndToEndOperationLatencyPolicyConfig,
sessionRetryOptions,
containerProactiveInitConfig,
defaultCustomSerializer,
isRegionScopedSessionCapturingEnabled);
this.cosmosAuthorizationTokenResolver = cosmosAuthorizationTokenResolver;
}
public RxDocumentClientImpl(URI serviceEndpoint,
String masterKeyOrResourceToken,
List permissionFeed,
ConnectionPolicy connectionPolicy,
ConsistencyLevel consistencyLevel,
Configs configs,
CosmosAuthorizationTokenResolver cosmosAuthorizationTokenResolver,
AzureKeyCredential credential,
TokenCredential tokenCredential,
boolean sessionCapturingOverride,
boolean connectionSharingAcrossClientsEnabled,
boolean contentResponseOnWriteEnabled,
CosmosClientMetadataCachesSnapshot metadataCachesSnapshot,
ApiType apiType,
CosmosClientTelemetryConfig clientTelemetryConfig,
String clientCorrelationId,
CosmosEndToEndOperationLatencyPolicyConfig cosmosEndToEndOperationLatencyPolicyConfig,
SessionRetryOptions sessionRetryOptions,
CosmosContainerProactiveInitConfig containerProactiveInitConfig,
CosmosItemSerializer defaultCustomSerializer,
boolean isRegionScopedSessionCapturingEnabled,
List operationPolicies) {
this(
serviceEndpoint,
masterKeyOrResourceToken,
permissionFeed,
connectionPolicy,
consistencyLevel,
configs,
credential,
tokenCredential,
sessionCapturingOverride,
connectionSharingAcrossClientsEnabled,
contentResponseOnWriteEnabled,
metadataCachesSnapshot,
apiType,
clientTelemetryConfig,
clientCorrelationId,
cosmosEndToEndOperationLatencyPolicyConfig,
sessionRetryOptions,
containerProactiveInitConfig,
defaultCustomSerializer,
isRegionScopedSessionCapturingEnabled);
this.cosmosAuthorizationTokenResolver = cosmosAuthorizationTokenResolver;
this.operationPolicies = operationPolicies;
}
private RxDocumentClientImpl(URI serviceEndpoint,
String masterKeyOrResourceToken,
List permissionFeed,
ConnectionPolicy connectionPolicy,
ConsistencyLevel consistencyLevel,
Configs configs,
AzureKeyCredential credential,
TokenCredential tokenCredential,
boolean sessionCapturingOverrideEnabled,
boolean connectionSharingAcrossClientsEnabled,
boolean contentResponseOnWriteEnabled,
CosmosClientMetadataCachesSnapshot metadataCachesSnapshot,
ApiType apiType,
CosmosClientTelemetryConfig clientTelemetryConfig,
String clientCorrelationId,
CosmosEndToEndOperationLatencyPolicyConfig cosmosEndToEndOperationLatencyPolicyConfig,
SessionRetryOptions sessionRetryOptions,
CosmosContainerProactiveInitConfig containerProactiveInitConfig,
CosmosItemSerializer defaultCustomSerializer,
boolean isRegionScopedSessionCapturingEnabled) {
this(
serviceEndpoint,
masterKeyOrResourceToken,
connectionPolicy,
consistencyLevel,
configs,
credential,
tokenCredential,
sessionCapturingOverrideEnabled,
connectionSharingAcrossClientsEnabled,
contentResponseOnWriteEnabled,
metadataCachesSnapshot,
apiType,
clientTelemetryConfig,
clientCorrelationId,
cosmosEndToEndOperationLatencyPolicyConfig,
sessionRetryOptions,
containerProactiveInitConfig,
defaultCustomSerializer,
isRegionScopedSessionCapturingEnabled);
if (permissionFeed != null && permissionFeed.size() > 0) {
this.resourceTokensMap = new HashMap<>();
for (Permission permission : permissionFeed) {
String[] segments = StringUtils.split(permission.getResourceLink(),
Constants.Properties.PATH_SEPARATOR.charAt(0));
if (segments.length == 0) {
throw new IllegalArgumentException("resourceLink");
}
List partitionKeyAndResourceTokenPairs = null;
PathInfo pathInfo = new PathInfo(false, StringUtils.EMPTY, StringUtils.EMPTY, false);
if (!PathsHelper.tryParsePathSegments(permission.getResourceLink(), pathInfo, null)) {
throw new IllegalArgumentException(permission.getResourceLink());
}
partitionKeyAndResourceTokenPairs = resourceTokensMap.get(pathInfo.resourceIdOrFullName);
if (partitionKeyAndResourceTokenPairs == null) {
partitionKeyAndResourceTokenPairs = new ArrayList<>();
this.resourceTokensMap.put(pathInfo.resourceIdOrFullName, partitionKeyAndResourceTokenPairs);
}
PartitionKey partitionKey = permission.getResourcePartitionKey();
partitionKeyAndResourceTokenPairs.add(new PartitionKeyAndResourceTokenPair(
partitionKey != null ? BridgeInternal.getPartitionKeyInternal(partitionKey) : PartitionKeyInternal.Empty,
permission.getToken()));
logger.debug("Initializing resource token map , with map key [{}] , partition key [{}] and resource token [{}]",
pathInfo.resourceIdOrFullName, partitionKey != null ? partitionKey.toString() : null, permission.getToken());
}
if(this.resourceTokensMap.isEmpty()) {
throw new IllegalArgumentException("permissionFeed");
}
String firstToken = permissionFeed.get(0).getToken();
if(ResourceTokenAuthorizationHelper.isResourceToken(firstToken)) {
this.firstResourceTokenFromPermissionFeed = firstToken;
}
}
}
RxDocumentClientImpl(URI serviceEndpoint,
String masterKeyOrResourceToken,
ConnectionPolicy connectionPolicy,
ConsistencyLevel consistencyLevel,
Configs configs,
AzureKeyCredential credential,
TokenCredential tokenCredential,
boolean sessionCapturingOverrideEnabled,
boolean connectionSharingAcrossClientsEnabled,
boolean contentResponseOnWriteEnabled,
CosmosClientMetadataCachesSnapshot metadataCachesSnapshot,
ApiType apiType,
CosmosClientTelemetryConfig clientTelemetryConfig,
String clientCorrelationId,
CosmosEndToEndOperationLatencyPolicyConfig cosmosEndToEndOperationLatencyPolicyConfig,
SessionRetryOptions sessionRetryOptions,
CosmosContainerProactiveInitConfig containerProactiveInitConfig,
CosmosItemSerializer defaultCustomSerializer,
boolean isRegionScopedSessionCapturingEnabled) {
assert(clientTelemetryConfig != null);
Boolean clientTelemetryEnabled = ImplementationBridgeHelpers
.CosmosClientTelemetryConfigHelper
.getCosmosClientTelemetryConfigAccessor()
.isSendClientTelemetryToServiceEnabled(clientTelemetryConfig);
assert(clientTelemetryEnabled != null);
activeClientsCnt.incrementAndGet();
this.clientId = clientIdGenerator.incrementAndGet();
this.clientCorrelationId = Strings.isNullOrWhiteSpace(clientCorrelationId) ?
String.format("%05d",this.clientId): clientCorrelationId;
clientMap.put(serviceEndpoint.toString(), clientMap.getOrDefault(serviceEndpoint.toString(), 0) + 1);
this.diagnosticsClientConfig = new DiagnosticsClientConfig();
this.diagnosticsClientConfig.withClientId(this.clientId);
this.diagnosticsClientConfig.withActiveClientCounter(activeClientsCnt);
this.diagnosticsClientConfig.withClientMap(clientMap);
this.diagnosticsClientConfig.withConnectionSharingAcrossClientsEnabled(connectionSharingAcrossClientsEnabled);
this.diagnosticsClientConfig.withConsistency(consistencyLevel);
this.throughputControlEnabled = new AtomicBoolean(false);
this.cosmosEndToEndOperationLatencyPolicyConfig = cosmosEndToEndOperationLatencyPolicyConfig;
this.diagnosticsClientConfig.withEndToEndOperationLatencyPolicy(cosmosEndToEndOperationLatencyPolicyConfig);
this.sessionRetryOptions = sessionRetryOptions;
this.defaultCustomSerializer = defaultCustomSerializer;
logger.info(
"Initializing DocumentClient [{}] with"
+ " serviceEndpoint [{}], connectionPolicy [{}], consistencyLevel [{}]",
this.clientId, serviceEndpoint, connectionPolicy, consistencyLevel, configs.getProtocol());
try {
this.connectionSharingAcrossClientsEnabled = connectionSharingAcrossClientsEnabled;
this.configs = configs;
this.masterKeyOrResourceToken = masterKeyOrResourceToken;
this.serviceEndpoint = serviceEndpoint;
this.credential = credential;
this.tokenCredential = tokenCredential;
this.contentResponseOnWriteEnabled = contentResponseOnWriteEnabled;
this.authorizationTokenType = AuthorizationTokenType.Invalid;
if (this.credential != null) {
hasAuthKeyResourceToken = false;
this.authorizationTokenType = AuthorizationTokenType.PrimaryMasterKey;
this.authorizationTokenProvider = new BaseAuthorizationTokenProvider(this.credential);
} else if (masterKeyOrResourceToken != null && ResourceTokenAuthorizationHelper.isResourceToken(masterKeyOrResourceToken)) {
this.authorizationTokenProvider = null;
hasAuthKeyResourceToken = true;
this.authorizationTokenType = AuthorizationTokenType.ResourceToken;
} else if(masterKeyOrResourceToken != null && !ResourceTokenAuthorizationHelper.isResourceToken(masterKeyOrResourceToken)) {
this.credential = new AzureKeyCredential(this.masterKeyOrResourceToken);
hasAuthKeyResourceToken = false;
this.authorizationTokenType = AuthorizationTokenType.PrimaryMasterKey;
this.authorizationTokenProvider = new BaseAuthorizationTokenProvider(this.credential);
} else {
hasAuthKeyResourceToken = false;
this.authorizationTokenProvider = null;
if (tokenCredential != null) {
this.tokenCredentialScopes = new String[] {
// AadTokenAuthorizationHelper.AAD_AUTH_TOKEN_COSMOS_WINDOWS_SCOPE,
serviceEndpoint.getScheme() + "://" + serviceEndpoint.getHost() + "/.default"
};
this.tokenCredentialCache = new SimpleTokenCache(() -> this.tokenCredential
.getToken(new TokenRequestContext().addScopes(this.tokenCredentialScopes)));
this.authorizationTokenType = AuthorizationTokenType.AadToken;
}
}
if (connectionPolicy != null) {
this.connectionPolicy = connectionPolicy;
} else {
this.connectionPolicy = new ConnectionPolicy(DirectConnectionConfig.getDefaultConfig());
}
this.diagnosticsClientConfig.withConnectionMode(this.getConnectionPolicy().getConnectionMode());
this.diagnosticsClientConfig.withConnectionPolicy(this.connectionPolicy);
this.diagnosticsClientConfig.withMultipleWriteRegionsEnabled(this.connectionPolicy.isMultipleWriteRegionsEnabled());
this.diagnosticsClientConfig.withEndpointDiscoveryEnabled(this.connectionPolicy.isEndpointDiscoveryEnabled());
this.diagnosticsClientConfig.withPreferredRegions(this.connectionPolicy.getPreferredRegions());
this.diagnosticsClientConfig.withMachineId(tempMachineId);
this.diagnosticsClientConfig.withProactiveContainerInitConfig(containerProactiveInitConfig);
this.diagnosticsClientConfig.withSessionRetryOptions(sessionRetryOptions);
this.sessionCapturingOverrideEnabled = sessionCapturingOverrideEnabled;
boolean disableSessionCapturing = (ConsistencyLevel.SESSION != consistencyLevel && !sessionCapturingOverrideEnabled);
this.sessionCapturingDisabled = disableSessionCapturing;
this.consistencyLevel = consistencyLevel;
this.userAgentContainer = new UserAgentContainer();
String userAgentSuffix = this.connectionPolicy.getUserAgentSuffix();
if (userAgentSuffix != null && userAgentSuffix.length() > 0) {
userAgentContainer.setSuffix(userAgentSuffix);
}
this.httpClientInterceptor = null;
this.reactorHttpClient = httpClient();
this.globalEndpointManager = new GlobalEndpointManager(asDatabaseAccountManagerInternal(), this.connectionPolicy, /**/configs);
this.isRegionScopedSessionCapturingEnabledOnClientOrSystemConfig = isRegionScopedSessionCapturingEnabled;
this.sessionContainer = new SessionContainer(this.serviceEndpoint.getHost(), disableSessionCapturing);
this.globalPartitionEndpointManagerForCircuitBreaker = new GlobalPartitionEndpointManagerForCircuitBreaker(this.globalEndpointManager);
this.globalPartitionEndpointManagerForCircuitBreaker.init();
this.cachedCosmosAsyncClientSnapshot = new AtomicReference<>();
this.diagnosticsClientConfig.withPartitionLevelCircuitBreakerConfig(this.globalPartitionEndpointManagerForCircuitBreaker.getCircuitBreakerConfig());
this.retryPolicy = new RetryPolicy(
this,
this.globalEndpointManager,
this.connectionPolicy,
this.globalPartitionEndpointManagerForCircuitBreaker);
this.resetSessionTokenRetryPolicy = retryPolicy;
CpuMemoryMonitor.register(this);
this.queryPlanCache = new ConcurrentHashMap<>();
this.apiType = apiType;
this.clientTelemetryConfig = clientTelemetryConfig;
} catch (RuntimeException e) {
logger.error("unexpected failure in initializing client.", e);
close();
throw e;
}
}
@Override
public DiagnosticsClientConfig getConfig() {
return diagnosticsClientConfig;
}
@Override
public CosmosDiagnostics createDiagnostics() {
CosmosDiagnostics diagnostics =
diagnosticsAccessor.create(this, telemetryCfgAccessor.getSamplingRate(this.clientTelemetryConfig));
this.mostRecentlyCreatedDiagnostics.set(diagnostics);
return diagnostics;
}
private DatabaseAccount initializeGatewayConfigurationReader() {
this.gatewayConfigurationReader = new GatewayServiceConfigurationReader(this.globalEndpointManager);
DatabaseAccount databaseAccount = this.globalEndpointManager.getLatestDatabaseAccount();
//Database account should not be null here,
// this.globalEndpointManager.init() must have been already called
// hence asserting it
if (databaseAccount == null) {
Throwable databaseRefreshErrorSnapshot = this.globalEndpointManager.getLatestDatabaseRefreshError();
if (databaseRefreshErrorSnapshot != null) {
logger.error("Client initialization failed. Check if the endpoint is reachable and if your auth token "
+ "is valid. More info: https://aka.ms/cosmosdb-tsg-service-unavailable-java. More details: "+ databaseRefreshErrorSnapshot.getMessage(),
databaseRefreshErrorSnapshot
);
throw new RuntimeException("Client initialization failed. Check if the endpoint is reachable and if your auth token "
+ "is valid. More info: https://aka.ms/cosmosdb-tsg-service-unavailable-java. More details: "+ databaseRefreshErrorSnapshot.getMessage(),
databaseRefreshErrorSnapshot);
} else {
logger.error("Client initialization failed."
+ " Check if the endpoint is reachable and if your auth token is valid. More info: https://aka.ms/cosmosdb-tsg-service-unavailable-java");
throw new RuntimeException("Client initialization failed. Check if the endpoint is reachable and if your auth token "
+ "is valid. More info: https://aka.ms/cosmosdb-tsg-service-unavailable-java.");
}
}
this.useMultipleWriteLocations = this.connectionPolicy.isMultipleWriteRegionsEnabled() && BridgeInternal.isEnableMultipleWriteLocations(databaseAccount);
return databaseAccount;
// TODO: add support for openAsync
// https://msdata.visualstudio.com/CosmosDB/_workitems/edit/332589
}
private void resetSessionContainerIfNeeded(DatabaseAccount databaseAccount) {
boolean isRegionScopingOfSessionTokensPossible = this.isRegionScopingOfSessionTokensPossible(databaseAccount, this.useMultipleWriteLocations, this.isRegionScopedSessionCapturingEnabledOnClientOrSystemConfig);
if (isRegionScopingOfSessionTokensPossible) {
this.sessionContainer = new RegionScopedSessionContainer(this.serviceEndpoint.getHost(), this.sessionCapturingDisabled, this.globalEndpointManager);
this.diagnosticsClientConfig.withRegionScopedSessionContainerOptions((RegionScopedSessionContainer) this.sessionContainer);
}
}
private boolean isRegionScopingOfSessionTokensPossible(DatabaseAccount databaseAccount, boolean useMultipleWriteLocations, boolean isRegionScopedSessionCapturingEnabled) {
if (!isRegionScopedSessionCapturingEnabled) {
return false;
}
if (!useMultipleWriteLocations) {
return false;
}
Iterable readableLocationsIterable = databaseAccount.getReadableLocations();
Iterator readableLocationsIterator = readableLocationsIterable.iterator();
while (readableLocationsIterator.hasNext()) {
DatabaseAccountLocation readableLocation = readableLocationsIterator.next();
String normalizedReadableRegion = readableLocation.getName().toLowerCase(Locale.ROOT).trim().replace(" ", "");
if (RegionNameToRegionIdMap.getRegionId(normalizedReadableRegion) == -1) {
return false;
}
}
return true;
}
private void updateGatewayProxy() {
(this.gatewayProxy).setGatewayServiceConfigurationReader(this.gatewayConfigurationReader);
(this.gatewayProxy).setCollectionCache(this.collectionCache);
(this.gatewayProxy).setPartitionKeyRangeCache(this.partitionKeyRangeCache);
(this.gatewayProxy).setUseMultipleWriteLocations(this.useMultipleWriteLocations);
(this.gatewayProxy).setSessionContainer(this.sessionContainer);
}
public void init(CosmosClientMetadataCachesSnapshot metadataCachesSnapshot, Function httpClientInterceptor) {
try {
// TODO: add support for openAsync
// https://msdata.visualstudio.com/CosmosDB/_workitems/edit/332589
this.httpClientInterceptor = httpClientInterceptor;
if (httpClientInterceptor != null) {
this.reactorHttpClient = httpClientInterceptor.apply(httpClient());
}
this.gatewayProxy = createRxGatewayProxy(this.sessionContainer,
this.consistencyLevel,
this.queryCompatibilityMode,
this.userAgentContainer,
this.globalEndpointManager,
this.reactorHttpClient,
this.apiType);
this.globalEndpointManager.init();
DatabaseAccount databaseAccountSnapshot = this.initializeGatewayConfigurationReader();
this.resetSessionContainerIfNeeded(databaseAccountSnapshot);
if (metadataCachesSnapshot != null) {
this.collectionCache = new RxClientCollectionCache(this,
this.sessionContainer,
this.gatewayProxy,
this,
this.retryPolicy,
metadataCachesSnapshot.getCollectionInfoByNameCache(),
metadataCachesSnapshot.getCollectionInfoByIdCache()
);
} else {
this.collectionCache = new RxClientCollectionCache(this,
this.sessionContainer,
this.gatewayProxy,
this,
this.retryPolicy);
}
this.resetSessionTokenRetryPolicy = new ResetSessionTokenRetryPolicyFactory(this.sessionContainer, this.collectionCache, this.retryPolicy);
this.partitionKeyRangeCache = new RxPartitionKeyRangeCache(RxDocumentClientImpl.this,
collectionCache);
updateGatewayProxy();
clientTelemetry = new ClientTelemetry(
this,
null,
randomUuid().toString(),
ManagementFactory.getRuntimeMXBean().getName(),
connectionPolicy.getConnectionMode(),
globalEndpointManager.getLatestDatabaseAccount().getId(),
null,
null,
this.configs,
this.clientTelemetryConfig,
this,
this.connectionPolicy.getPreferredRegions());
clientTelemetry.init().thenEmpty((publisher) -> {
logger.warn(
"Initialized DocumentClient [{}] with machineId[{}]"
+ " serviceEndpoint [{}], connectionPolicy [{}], consistencyLevel [{}]",
clientId,
ClientTelemetry.getMachineId(diagnosticsClientConfig),
serviceEndpoint,
connectionPolicy,
consistencyLevel);
}).subscribe();
if (this.connectionPolicy.getConnectionMode() == ConnectionMode.GATEWAY) {
this.storeModel = this.gatewayProxy;
} else {
this.initializeDirectConnectivity();
}
this.retryPolicy.setRxCollectionCache(this.collectionCache);
ConsistencyLevel effectiveConsistencyLevel = consistencyLevel != null
? consistencyLevel
: this.getDefaultConsistencyLevelOfAccount();
boolean updatedDisableSessionCapturing =
(ConsistencyLevel.SESSION != effectiveConsistencyLevel && !sessionCapturingOverrideEnabled);
this.sessionContainer.setDisableSessionCapturing(updatedDisableSessionCapturing);
} catch (Exception e) {
logger.error("unexpected failure in initializing client.", e);
close();
throw e;
}
}
public void serialize(CosmosClientMetadataCachesSnapshot state) {
RxCollectionCache.serialize(state, this.collectionCache);
}
private void initializeDirectConnectivity() {
this.addressResolver = new GlobalAddressResolver(this,
this.reactorHttpClient,
this.globalEndpointManager,
this.configs.getProtocol(),
this,
this.collectionCache,
this.partitionKeyRangeCache,
userAgentContainer,
// TODO: GATEWAY Configuration Reader
// this.gatewayConfigurationReader,
null,
this.connectionPolicy,
this.apiType);
this.storeClientFactory = new StoreClientFactory(
this.addressResolver,
this.diagnosticsClientConfig,
this.configs,
this.connectionPolicy,
// this.maxConcurrentConnectionOpenRequests,
this.userAgentContainer,
this.connectionSharingAcrossClientsEnabled,
this.clientTelemetry,
this.globalEndpointManager);
this.globalPartitionEndpointManagerForCircuitBreaker.setGlobalAddressResolver(this.addressResolver);
this.createStoreModel(true);
}
DatabaseAccountManagerInternal asDatabaseAccountManagerInternal() {
return new DatabaseAccountManagerInternal() {
@Override
public URI getServiceEndpoint() {
return RxDocumentClientImpl.this.getServiceEndpoint();
}
@Override
public Flux getDatabaseAccountFromEndpoint(URI endpoint) {
logger.info("Getting database account endpoint from {}", endpoint);
return RxDocumentClientImpl.this.getDatabaseAccountFromEndpoint(endpoint);
}
@Override
public ConnectionPolicy getConnectionPolicy() {
return RxDocumentClientImpl.this.getConnectionPolicy();
}
};
}
RxGatewayStoreModel createRxGatewayProxy(ISessionContainer sessionContainer,
ConsistencyLevel consistencyLevel,
QueryCompatibilityMode queryCompatibilityMode,
UserAgentContainer userAgentContainer,
GlobalEndpointManager globalEndpointManager,
HttpClient httpClient,
ApiType apiType) {
return new RxGatewayStoreModel(
this,
sessionContainer,
consistencyLevel,
queryCompatibilityMode,
userAgentContainer,
globalEndpointManager,
httpClient,
apiType);
}
private HttpClient httpClient() {
HttpClientConfig httpClientConfig = new HttpClientConfig(this.configs)
.withMaxIdleConnectionTimeout(this.connectionPolicy.getIdleHttpConnectionTimeout())
.withPoolSize(this.connectionPolicy.getMaxConnectionPoolSize())
.withProxy(this.connectionPolicy.getProxy())
.withNetworkRequestTimeout(this.connectionPolicy.getHttpNetworkRequestTimeout());
if (connectionSharingAcrossClientsEnabled) {
return SharedGatewayHttpClient.getOrCreateInstance(httpClientConfig, diagnosticsClientConfig);
} else {
diagnosticsClientConfig.withGatewayHttpClientConfig(httpClientConfig.toDiagnosticsString());
return HttpClient.createFixed(httpClientConfig);
}
}
private void createStoreModel(boolean subscribeRntbdStatus) {
// EnableReadRequestsFallback, if not explicitly set on the connection policy,
// is false if the account's consistency is bounded staleness,
// and true otherwise.
StoreClient storeClient = this.storeClientFactory.createStoreClient(this,
this.addressResolver,
this.sessionContainer,
this.gatewayConfigurationReader,
this,
this.useMultipleWriteLocations,
this.sessionRetryOptions);
this.storeModel = new ServerStoreModel(storeClient);
}
@Override
public URI getServiceEndpoint() {
return this.serviceEndpoint;
}
@Override
public ConnectionPolicy getConnectionPolicy() {
return this.connectionPolicy;
}
@Override
public boolean isContentResponseOnWriteEnabled() {
return contentResponseOnWriteEnabled;
}
@Override
public ConsistencyLevel getConsistencyLevel() {
return consistencyLevel;
}
@Override
public ClientTelemetry getClientTelemetry() {
return this.clientTelemetry;
}
@Override
public String getClientCorrelationId() {
return this.clientCorrelationId;
}
@Override
public String getMachineId() {
if (this.diagnosticsClientConfig == null) {
return null;
}
return ClientTelemetry.getMachineId(diagnosticsClientConfig);
}
@Override
public String getUserAgent() {
return this.userAgentContainer.getUserAgent();
}
@Override
public CosmosDiagnostics getMostRecentlyCreatedDiagnostics() {
return mostRecentlyCreatedDiagnostics.get();
}
@Override
public Mono> createDatabase(Database database, RequestOptions options) {
DocumentClientRetryPolicy retryPolicyInstance = this.resetSessionTokenRetryPolicy.getRequestPolicy(null);
return ObservableHelper.inlineIfPossibleAsObs(() -> createDatabaseInternal(database, options, retryPolicyInstance), retryPolicyInstance);
}
private Mono> createDatabaseInternal(Database database, RequestOptions options, DocumentClientRetryPolicy retryPolicyInstance) {
try {
if (database == null) {
throw new IllegalArgumentException("Database");
}
logger.debug("Creating a Database. id: [{}]", database.getId());
validateResource(database);
Map requestHeaders = this.getRequestHeaders(options, ResourceType.Database, OperationType.Create);
Instant serializationStartTimeUTC = Instant.now();
ByteBuffer byteBuffer = database.serializeJsonToByteBuffer(CosmosItemSerializer.DEFAULT_SERIALIZER, null, false);
Instant serializationEndTimeUTC = Instant.now();
SerializationDiagnosticsContext.SerializationDiagnostics serializationDiagnostics = new SerializationDiagnosticsContext.SerializationDiagnostics(
serializationStartTimeUTC,
serializationEndTimeUTC,
SerializationDiagnosticsContext.SerializationType.DATABASE_SERIALIZATION);
RxDocumentServiceRequest request = RxDocumentServiceRequest.create(this,
OperationType.Create, ResourceType.Database, Paths.DATABASES_ROOT, byteBuffer, requestHeaders, options);
if (retryPolicyInstance != null) {
retryPolicyInstance.onBeforeSendRequest(request);
}
SerializationDiagnosticsContext serializationDiagnosticsContext = BridgeInternal.getSerializationDiagnosticsContext(request.requestContext.cosmosDiagnostics);
if (serializationDiagnosticsContext != null) {
serializationDiagnosticsContext.addSerializationDiagnostics(serializationDiagnostics);
}
return this.create(request, retryPolicyInstance, getOperationContextAndListenerTuple(options))
.map(response -> toResourceResponse(response, Database.class));
} catch (Exception e) {
logger.debug("Failure in creating a database. due to [{}]", e.getMessage(), e);
return Mono.error(e);
}
}
@Override
public Mono> deleteDatabase(String databaseLink, RequestOptions options) {
DocumentClientRetryPolicy retryPolicyInstance = this.resetSessionTokenRetryPolicy.getRequestPolicy(null);
return ObservableHelper.inlineIfPossibleAsObs(() -> deleteDatabaseInternal(databaseLink, options, retryPolicyInstance), retryPolicyInstance);
}
private Mono> deleteDatabaseInternal(String databaseLink, RequestOptions options,
DocumentClientRetryPolicy retryPolicyInstance) {
try {
if (StringUtils.isEmpty(databaseLink)) {
throw new IllegalArgumentException("databaseLink");
}
logger.debug("Deleting a Database. databaseLink: [{}]", databaseLink);
String path = Utils.joinPath(databaseLink, null);
Map requestHeaders = this.getRequestHeaders(options, ResourceType.Database, OperationType.Delete);
RxDocumentServiceRequest request = RxDocumentServiceRequest.create(this,
OperationType.Delete, ResourceType.Database, path, requestHeaders, options);
if (retryPolicyInstance != null) {
retryPolicyInstance.onBeforeSendRequest(request);
}
return this.delete(request, retryPolicyInstance, getOperationContextAndListenerTuple(options)).map(response -> toResourceResponse(response, Database.class));
} catch (Exception e) {
logger.debug("Failure in deleting a database. due to [{}]", e.getMessage(), e);
return Mono.error(e);
}
}
@Override
public Mono> readDatabase(String databaseLink, RequestOptions options) {
DocumentClientRetryPolicy retryPolicyInstance = this.resetSessionTokenRetryPolicy.getRequestPolicy(null);
return ObservableHelper.inlineIfPossibleAsObs(() -> readDatabaseInternal(databaseLink, options, retryPolicyInstance), retryPolicyInstance);
}
private Mono> readDatabaseInternal(String databaseLink, RequestOptions options, DocumentClientRetryPolicy retryPolicyInstance) {
try {
if (StringUtils.isEmpty(databaseLink)) {
throw new IllegalArgumentException("databaseLink");
}
logger.debug("Reading a Database. databaseLink: [{}]", databaseLink);
String path = Utils.joinPath(databaseLink, null);
Map requestHeaders = this.getRequestHeaders(options, ResourceType.Database, OperationType.Read);
RxDocumentServiceRequest request = RxDocumentServiceRequest.create(this,
OperationType.Read, ResourceType.Database, path, requestHeaders, options);
if (retryPolicyInstance != null) {
retryPolicyInstance.onBeforeSendRequest(request);
}
return this.read(request, retryPolicyInstance).map(response -> toResourceResponse(response, Database.class));
} catch (Exception e) {
logger.debug("Failure in reading a database. due to [{}]", e.getMessage(), e);
return Mono.error(e);
}
}
@Override
public Flux> readDatabases(QueryFeedOperationState state) {
return nonDocumentReadFeed(state, ResourceType.Database, Database.class, Paths.DATABASES_ROOT);
}
private String parentResourceLinkToQueryLink(String parentResourceLink, ResourceType resourceTypeEnum) {
switch (resourceTypeEnum) {
case Database:
return Paths.DATABASES_ROOT;
case DocumentCollection:
return Utils.joinPath(parentResourceLink, Paths.COLLECTIONS_PATH_SEGMENT);
case Document:
return Utils.joinPath(parentResourceLink, Paths.DOCUMENTS_PATH_SEGMENT);
case Offer:
return Paths.OFFERS_ROOT;
case User:
return Utils.joinPath(parentResourceLink, Paths.USERS_PATH_SEGMENT);
case ClientEncryptionKey:
return Utils.joinPath(parentResourceLink, Paths.CLIENT_ENCRYPTION_KEY_PATH_SEGMENT);
case Permission:
return Utils.joinPath(parentResourceLink, Paths.PERMISSIONS_PATH_SEGMENT);
case Attachment:
return Utils.joinPath(parentResourceLink, Paths.ATTACHMENTS_PATH_SEGMENT);
case StoredProcedure:
return Utils.joinPath(parentResourceLink, Paths.STORED_PROCEDURES_PATH_SEGMENT);
case Trigger:
return Utils.joinPath(parentResourceLink, Paths.TRIGGERS_PATH_SEGMENT);
case UserDefinedFunction:
return Utils.joinPath(parentResourceLink, Paths.USER_DEFINED_FUNCTIONS_PATH_SEGMENT);
case Conflict:
return Utils.joinPath(parentResourceLink, Paths.CONFLICTS_PATH_SEGMENT);
default:
throw new IllegalArgumentException("resource type not supported");
}
}
private OperationContextAndListenerTuple getOperationContextAndListenerTuple(CosmosQueryRequestOptions options) {
if (options == null) {
return null;
}
return qryOptAccessor.getImpl(options).getOperationContextAndListenerTuple();
}
private OperationContextAndListenerTuple getOperationContextAndListenerTuple(RequestOptions options) {
if (options == null) {
return null;
}
return options.getOperationContextAndListenerTuple();
}
private Flux> createQuery(
String parentResourceLink,
SqlQuerySpec sqlQuery,
QueryFeedOperationState state,
Class klass,
ResourceType resourceTypeEnum) {
return createQuery(parentResourceLink, sqlQuery, state, klass, resourceTypeEnum, this);
}
private Flux> createQuery(
String parentResourceLink,
SqlQuerySpec sqlQuery,
QueryFeedOperationState state,
Class klass,
ResourceType resourceTypeEnum,
DiagnosticsClientContext innerDiagnosticsFactory) {
String resourceLink = parentResourceLinkToQueryLink(parentResourceLink, resourceTypeEnum);
CosmosQueryRequestOptions nonNullQueryOptions = state.getQueryOptions();
UUID correlationActivityIdOfRequestOptions = qryOptAccessor
.getImpl(nonNullQueryOptions)
.getCorrelationActivityId();
UUID correlationActivityId = correlationActivityIdOfRequestOptions != null ?
correlationActivityIdOfRequestOptions : randomUuid();
final AtomicBoolean isQueryCancelledOnTimeout = new AtomicBoolean(false);
IDocumentQueryClient queryClient = documentQueryClientImpl(RxDocumentClientImpl.this, getOperationContextAndListenerTuple(nonNullQueryOptions));
// Trying to put this logic as low as the query pipeline
// Since for parallelQuery, each partition will have its own request, so at this point, there will be no request associate with this retry policy.
// For default document context, it already wired up InvalidPartitionExceptionRetry, but there is no harm to wire it again here
InvalidPartitionExceptionRetryPolicy invalidPartitionExceptionRetryPolicy = new InvalidPartitionExceptionRetryPolicy(
this.collectionCache,
null,
resourceLink,
ModelBridgeInternal.getPropertiesFromQueryRequestOptions(nonNullQueryOptions));
final ScopedDiagnosticsFactory diagnosticsFactory = new ScopedDiagnosticsFactory(innerDiagnosticsFactory, false);
state.registerDiagnosticsFactory(
diagnosticsFactory::reset,
diagnosticsFactory::merge);
return
ObservableHelper.fluxInlineIfPossibleAsObs(
() -> createQueryInternal(
diagnosticsFactory, resourceLink, sqlQuery, state.getQueryOptions(), klass, resourceTypeEnum, queryClient, correlationActivityId, isQueryCancelledOnTimeout),
invalidPartitionExceptionRetryPolicy
).flatMap(result -> {
diagnosticsFactory.merge(state.getDiagnosticsContextSnapshot());
return Mono.just(result);
})
.onErrorMap(throwable -> {
diagnosticsFactory.merge(state.getDiagnosticsContextSnapshot());
return throwable;
})
.doOnCancel(() -> diagnosticsFactory.merge(state.getDiagnosticsContextSnapshot()));
}
private Flux> createQueryInternal(
DiagnosticsClientContext diagnosticsClientContext,
String resourceLink,
SqlQuerySpec sqlQuery,
CosmosQueryRequestOptions options,
Class klass,
ResourceType resourceTypeEnum,
IDocumentQueryClient queryClient,
UUID activityId,
final AtomicBoolean isQueryCancelledOnTimeout) {
Flux> executionContext =
DocumentQueryExecutionContextFactory
.createDocumentQueryExecutionContextAsync(diagnosticsClientContext, queryClient, resourceTypeEnum, klass, sqlQuery,
options, resourceLink, false, activityId,
Configs.isQueryPlanCachingEnabled(), queryPlanCache, isQueryCancelledOnTimeout);
AtomicBoolean isFirstResponse = new AtomicBoolean(true);
return executionContext.flatMap(iDocumentQueryExecutionContext -> {
QueryInfo queryInfo = null;
if (iDocumentQueryExecutionContext instanceof PipelinedQueryExecutionContextBase) {
queryInfo = ((PipelinedQueryExecutionContextBase) iDocumentQueryExecutionContext).getQueryInfo();
}
QueryInfo finalQueryInfo = queryInfo;
Flux> feedResponseFlux = iDocumentQueryExecutionContext.executeAsync()
.map(tFeedResponse -> {
if (finalQueryInfo != null) {
if (finalQueryInfo.hasSelectValue()) {
ModelBridgeInternal
.addQueryInfoToFeedResponse(tFeedResponse, finalQueryInfo);
}
if (isFirstResponse.compareAndSet(true, false)) {
ModelBridgeInternal.addQueryPlanDiagnosticsContextToFeedResponse(tFeedResponse,
finalQueryInfo.getQueryPlanDiagnosticsContext());
}
}
return tFeedResponse;
});
RequestOptions requestOptions = options == null? null : ImplementationBridgeHelpers
.CosmosQueryRequestOptionsHelper
.getCosmosQueryRequestOptionsAccessor()
.toRequestOptions(options);
CosmosEndToEndOperationLatencyPolicyConfig endToEndPolicyConfig =
getEndToEndOperationLatencyPolicyConfig(requestOptions, resourceTypeEnum, OperationType.Query);
if (endToEndPolicyConfig != null && endToEndPolicyConfig.isEnabled()) {
return getFeedResponseFluxWithTimeout(
feedResponseFlux,
endToEndPolicyConfig,
options,
isQueryCancelledOnTimeout,
diagnosticsClientContext);
}
return feedResponseFlux;
// concurrency is set to Queues.SMALL_BUFFER_SIZE to
// maximize the IDocumentQueryExecutionContext publisher instances to subscribe to concurrently
// prefetch is set to 1 to minimize the no. prefetched pages (result of merged executeAsync invocations)
}, Queues.SMALL_BUFFER_SIZE, 1);
}
private static void applyExceptionToMergedDiagnosticsForQuery(
CosmosQueryRequestOptions requestOptions,
CosmosException exception,
DiagnosticsClientContext diagnosticsClientContext) {
CosmosDiagnostics mostRecentlyCreatedDiagnostics =
diagnosticsClientContext.getMostRecentlyCreatedDiagnostics();
if (mostRecentlyCreatedDiagnostics != null) {
// When reaching here, it means the query(s) has timed out based on the e2e timeout config policy
// Since all the underlying ongoing query requests will all timed out
// We just use the last cosmosDiagnostics in the scoped diagnostics factory to populate the exception
BridgeInternal.setCosmosDiagnostics(
exception,
mostRecentlyCreatedDiagnostics);
} else {
List cancelledRequestDiagnostics =
qryOptAccessor
.getCancelledRequestDiagnosticsTracker(requestOptions);
// if there is any cancelled requests, collect cosmos diagnostics
if (cancelledRequestDiagnostics != null && !cancelledRequestDiagnostics.isEmpty()) {
// combine all the cosmos diagnostics
CosmosDiagnostics aggregratedCosmosDiagnostics =
cancelledRequestDiagnostics
.stream()
.reduce((first, toBeMerged) -> {
ClientSideRequestStatistics clientSideRequestStatistics =
ImplementationBridgeHelpers
.CosmosDiagnosticsHelper
.getCosmosDiagnosticsAccessor()
.getClientSideRequestStatisticsRaw(first);
ClientSideRequestStatistics toBeMergedClientSideRequestStatistics =
ImplementationBridgeHelpers
.CosmosDiagnosticsHelper
.getCosmosDiagnosticsAccessor()
.getClientSideRequestStatisticsRaw(first);
if (clientSideRequestStatistics == null) {
return toBeMerged;
} else {
clientSideRequestStatistics.mergeClientSideRequestStatistics(toBeMergedClientSideRequestStatistics);
return first;
}
})
.get();
BridgeInternal.setCosmosDiagnostics(exception, aggregratedCosmosDiagnostics);
}
}
}
private static Flux> getFeedResponseFluxWithTimeout(
Flux> feedResponseFlux,
CosmosEndToEndOperationLatencyPolicyConfig endToEndPolicyConfig,
CosmosQueryRequestOptions requestOptions,
final AtomicBoolean isQueryCancelledOnTimeout,
DiagnosticsClientContext diagnosticsClientContext) {
Duration endToEndTimeout = endToEndPolicyConfig.getEndToEndOperationTimeout();
if (endToEndTimeout.isNegative()) {
return feedResponseFlux
.timeout(endToEndTimeout)
.onErrorMap(throwable -> {
if (throwable instanceof TimeoutException) {
CosmosException cancellationException = getNegativeTimeoutException(null, endToEndTimeout);
cancellationException.setStackTrace(throwable.getStackTrace());
isQueryCancelledOnTimeout.set(true);
applyExceptionToMergedDiagnosticsForQuery(
requestOptions, cancellationException, diagnosticsClientContext);
return cancellationException;
}
return throwable;
});
}
return feedResponseFlux
.timeout(endToEndTimeout)
.onErrorMap(throwable -> {
if (throwable instanceof TimeoutException) {
CosmosException exception = new OperationCancelledException();
exception.setStackTrace(throwable.getStackTrace());
isQueryCancelledOnTimeout.set(true);
applyExceptionToMergedDiagnosticsForQuery(requestOptions, exception, diagnosticsClientContext);
return exception;
}
return throwable;
});
}
@Override
public Flux> queryDatabases(String query, QueryFeedOperationState state) {
return queryDatabases(new SqlQuerySpec(query), state);
}
@Override
public Flux> queryDatabases(SqlQuerySpec querySpec, QueryFeedOperationState state) {
return createQuery(Paths.DATABASES_ROOT, querySpec, state, Database.class, ResourceType.Database);
}
@Override
public Mono> createCollection(String databaseLink,
DocumentCollection collection, RequestOptions options) {
DocumentClientRetryPolicy retryPolicyInstance = this.resetSessionTokenRetryPolicy.getRequestPolicy(null);
return ObservableHelper.inlineIfPossibleAsObs(() -> this.createCollectionInternal(databaseLink, collection, options, retryPolicyInstance), retryPolicyInstance);
}
private Mono> createCollectionInternal(String databaseLink,
DocumentCollection collection, RequestOptions options, DocumentClientRetryPolicy retryPolicyInstance) {
try {
if (StringUtils.isEmpty(databaseLink)) {
throw new IllegalArgumentException("databaseLink");
}
if (collection == null) {
throw new IllegalArgumentException("collection");
}
logger.debug("Creating a Collection. databaseLink: [{}], Collection id: [{}]", databaseLink,
collection.getId());
validateResource(collection);
String path = Utils.joinPath(databaseLink, Paths.COLLECTIONS_PATH_SEGMENT);
Map requestHeaders = this.getRequestHeaders(options, ResourceType.DocumentCollection, OperationType.Create);
Instant serializationStartTimeUTC = Instant.now();
ByteBuffer byteBuffer = collection.serializeJsonToByteBuffer(CosmosItemSerializer.DEFAULT_SERIALIZER, null, false);
Instant serializationEndTimeUTC = Instant.now();
SerializationDiagnosticsContext.SerializationDiagnostics serializationDiagnostics = new SerializationDiagnosticsContext.SerializationDiagnostics(
serializationStartTimeUTC,
serializationEndTimeUTC,
SerializationDiagnosticsContext.SerializationType.CONTAINER_SERIALIZATION);
RxDocumentServiceRequest request = RxDocumentServiceRequest.create(this,
OperationType.Create, ResourceType.DocumentCollection, path, byteBuffer, requestHeaders, options);
if (retryPolicyInstance != null){
retryPolicyInstance.onBeforeSendRequest(request);
}
SerializationDiagnosticsContext serializationDiagnosticsContext = BridgeInternal.getSerializationDiagnosticsContext(request.requestContext.cosmosDiagnostics);
if (serializationDiagnosticsContext != null) {
serializationDiagnosticsContext.addSerializationDiagnostics(serializationDiagnostics);
}
return this.create(request, retryPolicyInstance, getOperationContextAndListenerTuple(options)).map(response -> toResourceResponse(response, DocumentCollection.class))
.doOnNext(resourceResponse -> {
// set the session token
this.sessionContainer.setSessionToken(
request,
resourceResponse.getResource().getResourceId(),
getAltLink(resourceResponse.getResource()),
resourceResponse.getResponseHeaders());
});
} catch (Exception e) {
logger.debug("Failure in creating a collection. due to [{}]", e.getMessage(), e);
return Mono.error(e);
}
}
@Override
public Mono> replaceCollection(DocumentCollection collection,
RequestOptions options) {
DocumentClientRetryPolicy retryPolicyInstance = this.resetSessionTokenRetryPolicy.getRequestPolicy(null);
return ObservableHelper.inlineIfPossibleAsObs(() -> replaceCollectionInternal(collection, options, retryPolicyInstance), retryPolicyInstance);
}
private Mono> replaceCollectionInternal(DocumentCollection collection,
RequestOptions options, DocumentClientRetryPolicy retryPolicyInstance) {
try {
if (collection == null) {
throw new IllegalArgumentException("collection");
}
logger.debug("Replacing a Collection. id: [{}]", collection.getId());
validateResource(collection);
String path = Utils.joinPath(collection.getSelfLink(), null);
Map requestHeaders = this.getRequestHeaders(options, ResourceType.DocumentCollection, OperationType.Replace);
Instant serializationStartTimeUTC = Instant.now();
ByteBuffer byteBuffer = collection.serializeJsonToByteBuffer(CosmosItemSerializer.DEFAULT_SERIALIZER, null, false);
Instant serializationEndTimeUTC = Instant.now();
SerializationDiagnosticsContext.SerializationDiagnostics serializationDiagnostics = new SerializationDiagnosticsContext.SerializationDiagnostics(
serializationStartTimeUTC,
serializationEndTimeUTC,
SerializationDiagnosticsContext.SerializationType.CONTAINER_SERIALIZATION);
RxDocumentServiceRequest request = RxDocumentServiceRequest.create(this,
OperationType.Replace, ResourceType.DocumentCollection, path, byteBuffer, requestHeaders, options);
// TODO: .Net has some logic for updating session token which we don't
// have here
if (retryPolicyInstance != null){
retryPolicyInstance.onBeforeSendRequest(request);
}
SerializationDiagnosticsContext serializationDiagnosticsContext = BridgeInternal.getSerializationDiagnosticsContext(request.requestContext.cosmosDiagnostics);
if (serializationDiagnosticsContext != null) {
serializationDiagnosticsContext.addSerializationDiagnostics(serializationDiagnostics);
}
return this.replace(request, retryPolicyInstance).map(response -> toResourceResponse(response, DocumentCollection.class))
.doOnNext(resourceResponse -> {
if (resourceResponse.getResource() != null) {
// set the session token
this.sessionContainer.setSessionToken(
request,
resourceResponse.getResource().getResourceId(),
getAltLink(resourceResponse.getResource()),
resourceResponse.getResponseHeaders());
}
});
} catch (Exception e) {
logger.debug("Failure in replacing a collection. due to [{}]", e.getMessage(), e);
return Mono.error(e);
}
}
@Override
public Mono> deleteCollection(String collectionLink,
RequestOptions options) {
DocumentClientRetryPolicy retryPolicyInstance = this.resetSessionTokenRetryPolicy.getRequestPolicy(null);
return ObservableHelper.inlineIfPossibleAsObs(() -> deleteCollectionInternal(collectionLink, options, retryPolicyInstance), retryPolicyInstance);
}
private Mono> deleteCollectionInternal(String collectionLink,
RequestOptions options, DocumentClientRetryPolicy retryPolicyInstance) {
try {
if (StringUtils.isEmpty(collectionLink)) {
throw new IllegalArgumentException("collectionLink");
}
logger.debug("Deleting a Collection. collectionLink: [{}]", collectionLink);
String path = Utils.joinPath(collectionLink, null);
Map requestHeaders = this.getRequestHeaders(options, ResourceType.DocumentCollection, OperationType.Delete);
RxDocumentServiceRequest request = RxDocumentServiceRequest.create(this,
OperationType.Delete, ResourceType.DocumentCollection, path, requestHeaders, options);
if (retryPolicyInstance != null){
retryPolicyInstance.onBeforeSendRequest(request);
}
return this.delete(request, retryPolicyInstance, getOperationContextAndListenerTuple(options)).map(response -> toResourceResponse(response, DocumentCollection.class));
} catch (Exception e) {
logger.debug("Failure in deleting a collection, due to [{}]", e.getMessage(), e);
return Mono.error(e);
}
}
private Mono delete(RxDocumentServiceRequest request, DocumentClientRetryPolicy documentClientRetryPolicy, OperationContextAndListenerTuple operationContextAndListenerTuple) {
return populateHeadersAsync(request, RequestVerb.DELETE)
.flatMap(requestPopulated -> {
if (documentClientRetryPolicy.getRetryContext() != null && documentClientRetryPolicy.getRetryContext().getRetryCount() > 0) {
documentClientRetryPolicy.getRetryContext().updateEndTime();
}
return getStoreProxy(requestPopulated).processMessage(requestPopulated, operationContextAndListenerTuple);
});
}
private Mono deleteAllItemsByPartitionKey(RxDocumentServiceRequest request, DocumentClientRetryPolicy documentClientRetryPolicy, OperationContextAndListenerTuple operationContextAndListenerTuple) {
return populateHeadersAsync(request, RequestVerb.POST)
.flatMap(requestPopulated -> {
RxStoreModel storeProxy = this.getStoreProxy(requestPopulated);
if (documentClientRetryPolicy.getRetryContext() != null && documentClientRetryPolicy.getRetryContext().getRetryCount() > 0) {
documentClientRetryPolicy.getRetryContext().updateEndTime();
}
return storeProxy.processMessage(requestPopulated, operationContextAndListenerTuple);
});
}
private Mono read(RxDocumentServiceRequest request, DocumentClientRetryPolicy documentClientRetryPolicy) {
return populateHeadersAsync(request, RequestVerb.GET)
.flatMap(requestPopulated -> {
if (documentClientRetryPolicy.getRetryContext() != null && documentClientRetryPolicy.getRetryContext().getRetryCount() > 0) {
documentClientRetryPolicy.getRetryContext().updateEndTime();
}
return getStoreProxy(requestPopulated).processMessage(requestPopulated);
});
}
Mono readFeed(RxDocumentServiceRequest request) {
return populateHeadersAsync(request, RequestVerb.GET)
.flatMap(requestPopulated -> getStoreProxy(requestPopulated).processMessage(requestPopulated));
}
private Mono query(RxDocumentServiceRequest request) {
// If endToEndOperationLatencyPolicy is set in the query request options,
// then it should have been populated into request context already
// otherwise set them here with the client level policy
return populateHeadersAsync(request, RequestVerb.POST)
.flatMap(requestPopulated ->
this.getStoreProxy(requestPopulated).processMessage(requestPopulated)
.map(response -> {
this.captureSessionToken(requestPopulated, response);
return response;
}
));
}
@Override
public Mono> readCollection(String collectionLink,
RequestOptions options) {
DocumentClientRetryPolicy retryPolicyInstance = this.resetSessionTokenRetryPolicy.getRequestPolicy(null);
return ObservableHelper.inlineIfPossibleAsObs(() -> readCollectionInternal(collectionLink, options, retryPolicyInstance), retryPolicyInstance);
}
private Mono> readCollectionInternal(String collectionLink,
RequestOptions options, DocumentClientRetryPolicy retryPolicyInstance) {
// we are using an observable factory here
// observable will be created fresh upon subscription
// this is to ensure we capture most up to date information (e.g.,
// session)
try {
if (StringUtils.isEmpty(collectionLink)) {
throw new IllegalArgumentException("collectionLink");
}
logger.debug("Reading a Collection. collectionLink: [{}]", collectionLink);
String path = Utils.joinPath(collectionLink, null);
Map requestHeaders = this.getRequestHeaders(options, ResourceType.DocumentCollection, OperationType.Read);
RxDocumentServiceRequest request = RxDocumentServiceRequest.create(this,
OperationType.Read, ResourceType.DocumentCollection, path, requestHeaders, options);
if (retryPolicyInstance != null){
retryPolicyInstance.onBeforeSendRequest(request);
}
return this.read(request, retryPolicyInstance).map(response -> toResourceResponse(response, DocumentCollection.class));
} catch (Exception e) {
// this is only in trace level to capture what's going on
logger.debug("Failure in reading a collection, due to [{}]", e.getMessage(), e);
return Mono.error(e);
}
}
@Override
public Flux> readCollections(String databaseLink, QueryFeedOperationState state) {
if (StringUtils.isEmpty(databaseLink)) {
throw new IllegalArgumentException("databaseLink");
}
return nonDocumentReadFeed(state, ResourceType.DocumentCollection, DocumentCollection.class,
Utils.joinPath(databaseLink, Paths.COLLECTIONS_PATH_SEGMENT));
}
@Override
public Flux> queryCollections(String databaseLink, String query,
QueryFeedOperationState state) {
return createQuery(databaseLink, new SqlQuerySpec(query), state, DocumentCollection.class, ResourceType.DocumentCollection);
}
@Override
public Flux> queryCollections(String databaseLink,
SqlQuerySpec querySpec, QueryFeedOperationState state) {
return createQuery(databaseLink, querySpec, state, DocumentCollection.class, ResourceType.DocumentCollection);
}
private static String serializeProcedureParams(List objectArray) {
String[] stringArray = new String[objectArray.size()];
for (int i = 0; i < objectArray.size(); ++i) {
Object object = objectArray.get(i);
if (object instanceof JsonSerializable) {
stringArray[i] = ((JsonSerializable) object).toJson();
} else {
// POJO, ObjectNode, number, STRING or Boolean
try {
stringArray[i] = mapper.writeValueAsString(object);
} catch (IOException e) {
throw new IllegalArgumentException("Can't serialize the object into the json string", e);
}
}
}
return String.format("[%s]", StringUtils.join(stringArray, ","));
}
private static void validateResource(Resource resource) {
if (!StringUtils.isEmpty(resource.getId())) {
if (resource.getId().indexOf('/') != -1 || resource.getId().indexOf('\\') != -1 ||
resource.getId().indexOf('?') != -1 || resource.getId().indexOf('#') != -1) {
throw new IllegalArgumentException("Id contains illegal chars.");
}
if (resource.getId().endsWith(" ")) {
throw new IllegalArgumentException("Id ends with a space.");
}
}
}
private Map getRequestHeaders(RequestOptions options, ResourceType resourceType, OperationType operationType) {
Map headers = new HashMap<>();
if (this.useMultipleWriteLocations) {
headers.put(HttpConstants.HttpHeaders.ALLOW_TENTATIVE_WRITES, Boolean.TRUE.toString());
}
if (consistencyLevel != null) {
headers.put(HttpConstants.HttpHeaders.CONSISTENCY_LEVEL, consistencyLevel.toString());
}
if (options == null) {
// Corner case, if options are null, then just check this flag from CosmosClientBuilder
// If content response on write is not enabled, and operation is document write - then add minimal prefer header
// Otherwise, don't add this header, which means return the full response
if (!this.contentResponseOnWriteEnabled && resourceType.equals(ResourceType.Document) && operationType.isWriteOperation()) {
headers.put(HttpConstants.HttpHeaders.PREFER, HttpConstants.HeaderValues.PREFER_RETURN_MINIMAL);
}
return headers;
}
Map customOptions = options.getHeaders();
if (customOptions != null) {
headers.putAll(customOptions);
}
boolean contentResponseOnWriteEnabled = this.contentResponseOnWriteEnabled;
// If options has contentResponseOnWriteEnabled set to true / false, override the value from CosmosClientBuilder
if (options.isContentResponseOnWriteEnabled() != null) {
contentResponseOnWriteEnabled = options.isContentResponseOnWriteEnabled();
}
// If content response on write is not enabled, and operation is document write - then add minimal prefer header
// Otherwise, don't add this header, which means return the full response
if (!contentResponseOnWriteEnabled && resourceType.equals(ResourceType.Document) && operationType.isWriteOperation()) {
headers.put(HttpConstants.HttpHeaders.PREFER, HttpConstants.HeaderValues.PREFER_RETURN_MINIMAL);
}
if (options.getIfMatchETag() != null) {
headers.put(HttpConstants.HttpHeaders.IF_MATCH, options.getIfMatchETag());
}
if (options.getIfNoneMatchETag() != null) {
headers.put(HttpConstants.HttpHeaders.IF_NONE_MATCH, options.getIfNoneMatchETag());
}
if (options.getConsistencyLevel() != null) {
headers.put(HttpConstants.HttpHeaders.CONSISTENCY_LEVEL, options.getConsistencyLevel().toString());
}
if (options.getIndexingDirective() != null) {
headers.put(HttpConstants.HttpHeaders.INDEXING_DIRECTIVE, options.getIndexingDirective().toString());
}
if (options.getPostTriggerInclude() != null && options.getPostTriggerInclude().size() > 0) {
String postTriggerInclude = StringUtils.join(options.getPostTriggerInclude(), ",");
headers.put(HttpConstants.HttpHeaders.POST_TRIGGER_INCLUDE, postTriggerInclude);
}
if (options.getPreTriggerInclude() != null && options.getPreTriggerInclude().size() > 0) {
String preTriggerInclude = StringUtils.join(options.getPreTriggerInclude(), ",");
headers.put(HttpConstants.HttpHeaders.PRE_TRIGGER_INCLUDE, preTriggerInclude);
}
if (!Strings.isNullOrEmpty(options.getSessionToken())) {
headers.put(HttpConstants.HttpHeaders.SESSION_TOKEN, options.getSessionToken());
}
if (options.getResourceTokenExpirySeconds() != null) {
headers.put(HttpConstants.HttpHeaders.RESOURCE_TOKEN_EXPIRY,
String.valueOf(options.getResourceTokenExpirySeconds()));
}
if (options.getOfferThroughput() != null && options.getOfferThroughput() >= 0) {
headers.put(HttpConstants.HttpHeaders.OFFER_THROUGHPUT, options.getOfferThroughput().toString());
} else if (options.getOfferType() != null) {
headers.put(HttpConstants.HttpHeaders.OFFER_TYPE, options.getOfferType());
}
if (options.getOfferThroughput() == null) {
if (options.getThroughputProperties() != null) {
Offer offer = ModelBridgeInternal.getOfferFromThroughputProperties(options.getThroughputProperties());
final OfferAutoscaleSettings offerAutoscaleSettings = offer.getOfferAutoScaleSettings();
OfferAutoscaleAutoUpgradeProperties autoscaleAutoUpgradeProperties = null;
if (offerAutoscaleSettings != null) {
autoscaleAutoUpgradeProperties
= offer.getOfferAutoScaleSettings().getAutoscaleAutoUpgradeProperties();
}
if (offer.hasOfferThroughput() &&
(offerAutoscaleSettings != null && offerAutoscaleSettings.getMaxThroughput() >= 0 ||
autoscaleAutoUpgradeProperties != null &&
autoscaleAutoUpgradeProperties
.getAutoscaleThroughputProperties()
.getIncrementPercent() >= 0)) {
throw new IllegalArgumentException("Autoscale provisioned throughput can not be configured with "
+ "fixed offer");
}
if (offer.hasOfferThroughput()) {
headers.put(HttpConstants.HttpHeaders.OFFER_THROUGHPUT, String.valueOf(offer.getThroughput()));
} else if (offer.getOfferAutoScaleSettings() != null) {
headers.put(HttpConstants.HttpHeaders.OFFER_AUTOPILOT_SETTINGS,
offer.getOfferAutoScaleSettings().toJson());
}
}
}
if (options.isQuotaInfoEnabled()) {
headers.put(HttpConstants.HttpHeaders.POPULATE_QUOTA_INFO, String.valueOf(true));
}
if (options.isScriptLoggingEnabled()) {
headers.put(HttpConstants.HttpHeaders.SCRIPT_ENABLE_LOGGING, String.valueOf(true));
}
if (options.getDedicatedGatewayRequestOptions() != null) {
if (options.getDedicatedGatewayRequestOptions().getMaxIntegratedCacheStaleness() != null) {
headers.put(HttpConstants.HttpHeaders.DEDICATED_GATEWAY_PER_REQUEST_CACHE_STALENESS,
String.valueOf(Utils.getMaxIntegratedCacheStalenessInMillis(options.getDedicatedGatewayRequestOptions())));
}
if (options.getDedicatedGatewayRequestOptions().isIntegratedCacheBypassed()) {
headers.put(HttpConstants.HttpHeaders.DEDICATED_GATEWAY_PER_REQUEST_BYPASS_CACHE,
String.valueOf(options.getDedicatedGatewayRequestOptions().isIntegratedCacheBypassed()));
}
}
return headers;
}
public IRetryPolicyFactory getResetSessionTokenRetryPolicy() {
return this.resetSessionTokenRetryPolicy;
}
private Mono addPartitionKeyInformation(RxDocumentServiceRequest request,
ByteBuffer contentAsByteBuffer,
Document document,
RequestOptions options) {
Mono> collectionObs = this.collectionCache.resolveCollectionAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), request);
return collectionObs
.map(collectionValueHolder -> {
addPartitionKeyInformation(request, contentAsByteBuffer, document, options, collectionValueHolder.v, null);
return request;
});
}
private Mono addPartitionKeyInformation(RxDocumentServiceRequest request,
ByteBuffer contentAsByteBuffer,
Object document,
RequestOptions options,
Mono> collectionObs,
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker) {
return collectionObs.map(collectionValueHolder -> {
addPartitionKeyInformation(request, contentAsByteBuffer, document, options, collectionValueHolder.v, pointOperationContextForCircuitBreaker);
return request;
});
}
private void addPartitionKeyInformation(RxDocumentServiceRequest request,
ByteBuffer contentAsByteBuffer,
Object objectDoc, RequestOptions options,
DocumentCollection collection,
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker) {
PartitionKeyDefinition partitionKeyDefinition = collection.getPartitionKey();
PartitionKeyInternal partitionKeyInternal = null;
if (options != null && options.getPartitionKey() != null && options.getPartitionKey().equals(PartitionKey.NONE)){
partitionKeyInternal = ModelBridgeInternal.getNonePartitionKey(partitionKeyDefinition);
} else if (options != null && options.getPartitionKey() != null) {
partitionKeyInternal = BridgeInternal.getPartitionKeyInternal(options.getPartitionKey());
} else if (partitionKeyDefinition == null || partitionKeyDefinition.getPaths().size() == 0) {
// For backward compatibility, if collection doesn't have partition key defined, we assume all documents
// have empty value for it and user doesn't need to specify it explicitly.
partitionKeyInternal = PartitionKeyInternal.getEmpty();
} else if (contentAsByteBuffer != null || objectDoc != null) {
InternalObjectNode internalObjectNode;
if (objectDoc instanceof InternalObjectNode) {
internalObjectNode = (InternalObjectNode) objectDoc;
} else if (objectDoc instanceof ObjectNode) {
internalObjectNode = new InternalObjectNode((ObjectNode)objectDoc);
} else if (contentAsByteBuffer != null) {
contentAsByteBuffer.rewind();
internalObjectNode = new InternalObjectNode(contentAsByteBuffer);
} else {
// This is a safety check, this should not happen ever.
// If it does, it is a SDK bug
throw new IllegalStateException("ContentAsByteBuffer and objectDoc are null");
}
Instant serializationStartTime = Instant.now();
partitionKeyInternal = PartitionKeyHelper.extractPartitionKeyValueFromDocument(internalObjectNode, partitionKeyDefinition);
Instant serializationEndTime = Instant.now();
SerializationDiagnosticsContext.SerializationDiagnostics serializationDiagnostics = new SerializationDiagnosticsContext.SerializationDiagnostics(
serializationStartTime,
serializationEndTime,
SerializationDiagnosticsContext.SerializationType.PARTITION_KEY_FETCH_SERIALIZATION
);
SerializationDiagnosticsContext serializationDiagnosticsContext = BridgeInternal.getSerializationDiagnosticsContext(request.requestContext.cosmosDiagnostics);
if (serializationDiagnosticsContext != null) {
serializationDiagnosticsContext.addSerializationDiagnostics(serializationDiagnostics);
} else if (pointOperationContextForCircuitBreaker != null) {
serializationDiagnosticsContext = pointOperationContextForCircuitBreaker.getSerializationDiagnosticsContext();
if (serializationDiagnosticsContext != null) {
serializationDiagnosticsContext.addSerializationDiagnostics(serializationDiagnostics);
}
}
} else {
throw new UnsupportedOperationException("PartitionKey value must be supplied for this operation.");
}
request.setPartitionKeyInternal(partitionKeyInternal);
request.setPartitionKeyDefinition(partitionKeyDefinition);
request.getHeaders().put(HttpConstants.HttpHeaders.PARTITION_KEY, Utils.escapeNonAscii(partitionKeyInternal.toJson()));
}
private Mono>> getCreateDocumentRequest(DocumentClientRetryPolicy requestRetryPolicy,
String documentCollectionLink,
Object document,
RequestOptions options,
boolean disableAutomaticIdGeneration,
OperationType operationType,
DiagnosticsClientContext clientContextOverride,
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker) {
if (StringUtils.isEmpty(documentCollectionLink)) {
throw new IllegalArgumentException("documentCollectionLink");
}
if (document == null) {
throw new IllegalArgumentException("document");
}
Instant serializationStartTimeUTC = Instant.now();
String trackingId = null;
if (options != null) {
trackingId = options.getTrackingId();
}
ByteBuffer content = InternalObjectNode.serializeJsonToByteBuffer(document, options.getEffectiveItemSerializer(), trackingId, true);
Instant serializationEndTimeUTC = Instant.now();
SerializationDiagnosticsContext.SerializationDiagnostics serializationDiagnostics = new SerializationDiagnosticsContext.SerializationDiagnostics(
serializationStartTimeUTC,
serializationEndTimeUTC,
SerializationDiagnosticsContext.SerializationType.ITEM_SERIALIZATION);
String path = Utils.joinPath(documentCollectionLink, Paths.DOCUMENTS_PATH_SEGMENT);
Map requestHeaders = this.getRequestHeaders(options, ResourceType.Document, operationType);
RxDocumentServiceRequest request = RxDocumentServiceRequest.create(
getEffectiveClientContext(clientContextOverride),
operationType, ResourceType.Document, path, requestHeaders, options, content);
if (operationType.isWriteOperation() && options != null && options.getNonIdempotentWriteRetriesEnabled() != null && options.getNonIdempotentWriteRetriesEnabled()) {
request.setNonIdempotentWriteRetriesEnabled(true);
}
if( options != null) {
options.getMarkE2ETimeoutInRequestContextCallbackHook().set(
() -> request.requestContext.setIsRequestCancelledOnTimeout(new AtomicBoolean(true)));
request.requestContext.setExcludeRegions(options.getExcludedRegions());
request.requestContext.setKeywordIdentifiers(options.getKeywordIdentifiers());
}
SerializationDiagnosticsContext serializationDiagnosticsContext = BridgeInternal.getSerializationDiagnosticsContext(request.requestContext.cosmosDiagnostics);
if (serializationDiagnosticsContext != null) {
serializationDiagnosticsContext.addSerializationDiagnostics(serializationDiagnostics);
}
if (requestRetryPolicy != null) {
requestRetryPolicy.onBeforeSendRequest(request);
}
Mono> collectionObs = this.collectionCache.resolveCollectionAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), request);
return addPartitionKeyInformation(request, content, document, options, collectionObs, pointOperationContextForCircuitBreaker)
.zipWith(collectionObs);
}
private Mono getBatchDocumentRequest(DocumentClientRetryPolicy requestRetryPolicy,
String documentCollectionLink,
ServerBatchRequest serverBatchRequest,
RequestOptions options,
boolean disableAutomaticIdGeneration) {
checkArgument(StringUtils.isNotEmpty(documentCollectionLink), "expected non empty documentCollectionLink");
checkNotNull(serverBatchRequest, "expected non null serverBatchRequest");
Instant serializationStartTimeUTC = Instant.now();
ByteBuffer content = ByteBuffer.wrap(Utils.getUTF8Bytes(serverBatchRequest.getRequestBody()));
Instant serializationEndTimeUTC = Instant.now();
SerializationDiagnosticsContext.SerializationDiagnostics serializationDiagnostics = new SerializationDiagnosticsContext.SerializationDiagnostics(
serializationStartTimeUTC,
serializationEndTimeUTC,
SerializationDiagnosticsContext.SerializationType.ITEM_SERIALIZATION);
String path = Utils.joinPath(documentCollectionLink, Paths.DOCUMENTS_PATH_SEGMENT);
Map requestHeaders = this.getRequestHeaders(options, ResourceType.Document, OperationType.Batch);
RxDocumentServiceRequest request = RxDocumentServiceRequest.create(
this,
OperationType.Batch,
ResourceType.Document,
path,
requestHeaders,
options,
content);
if (options != null) {
options.getMarkE2ETimeoutInRequestContextCallbackHook().set(
() -> request.requestContext.setIsRequestCancelledOnTimeout(new AtomicBoolean(true)));
request.requestContext.setExcludeRegions(options.getExcludedRegions());
request.requestContext.setKeywordIdentifiers(options.getKeywordIdentifiers());
}
SerializationDiagnosticsContext serializationDiagnosticsContext = BridgeInternal.getSerializationDiagnosticsContext(request.requestContext.cosmosDiagnostics);
if (serializationDiagnosticsContext != null) {
serializationDiagnosticsContext.addSerializationDiagnostics(serializationDiagnostics);
}
if (options != null) {
request.requestContext.setExcludeRegions(options.getExcludedRegions());
request.requestContext.setKeywordIdentifiers(options.getKeywordIdentifiers());
}
// note: calling onBeforeSendRequest is a cheap operation which injects a CosmosDiagnostics
// instance into 'request' amongst other things - this way metadataDiagnosticsContext is not
// null and can be used for metadata-related telemetry (partition key range, container and server address lookups)
if (requestRetryPolicy != null) {
requestRetryPolicy.onBeforeSendRequest(request);
}
MetadataDiagnosticsContext metadataDiagnosticsContext = BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics);
request.requestContext.setPointOperationContext(
new PointOperationContextForCircuitBreaker(
new AtomicBoolean(false),
false,
documentCollectionLink,
serializationDiagnosticsContext));
return this.collectionCache.resolveCollectionAsync(metadataDiagnosticsContext, request)
.flatMap(documentCollectionValueHolder -> {
if (documentCollectionValueHolder == null || documentCollectionValueHolder.v == null) {
return Mono.error(new IllegalStateException("documentCollectionValueHolder or documentCollectionValueHolder.v cannot be null"));
}
return this.partitionKeyRangeCache.tryLookupAsync(metadataDiagnosticsContext, documentCollectionValueHolder.v.getResourceId(), null, null)
.flatMap(collectionRoutingMapValueHolder -> {
if (collectionRoutingMapValueHolder == null || collectionRoutingMapValueHolder.v == null) {
return Mono.error(new IllegalStateException("collectionRoutingMapValueHolder or collectionRoutingMapValueHolder.v cannot be null"));
}
addBatchHeaders(request, serverBatchRequest, documentCollectionValueHolder.v);
if (this.globalPartitionEndpointManagerForCircuitBreaker.isPartitionLevelCircuitBreakingApplicable(request) && options != null) {
options.setPartitionKeyDefinition(documentCollectionValueHolder.v.getPartitionKey());
addPartitionLevelUnavailableRegionsForRequest(request, options, collectionRoutingMapValueHolder.v, requestRetryPolicy);
}
return Mono.just(request);
});
});
}
private RxDocumentServiceRequest addBatchHeaders(RxDocumentServiceRequest request,
ServerBatchRequest serverBatchRequest,
DocumentCollection collection) {
if(serverBatchRequest instanceof SinglePartitionKeyServerBatchRequest) {
PartitionKey partitionKey = ((SinglePartitionKeyServerBatchRequest) serverBatchRequest).getPartitionKeyValue();
PartitionKeyInternal partitionKeyInternal;
if (partitionKey.equals(PartitionKey.NONE)) {
PartitionKeyDefinition partitionKeyDefinition = collection.getPartitionKey();
partitionKeyInternal = ModelBridgeInternal.getNonePartitionKey(partitionKeyDefinition);
} else {
// Partition key is always non-null
partitionKeyInternal = BridgeInternal.getPartitionKeyInternal(partitionKey);
}
request.setPartitionKeyInternal(partitionKeyInternal);
request.getHeaders().put(HttpConstants.HttpHeaders.PARTITION_KEY, Utils.escapeNonAscii(partitionKeyInternal.toJson()));
} else if(serverBatchRequest instanceof PartitionKeyRangeServerBatchRequest) {
request.setPartitionKeyRangeIdentity(new PartitionKeyRangeIdentity(((PartitionKeyRangeServerBatchRequest) serverBatchRequest).getPartitionKeyRangeId()));
} else {
throw new UnsupportedOperationException("Unknown Server request.");
}
request.getHeaders().put(HttpConstants.HttpHeaders.IS_BATCH_REQUEST, Boolean.TRUE.toString());
request.getHeaders().put(HttpConstants.HttpHeaders.IS_BATCH_ATOMIC, String.valueOf(serverBatchRequest.isAtomicBatch()));
request.getHeaders().put(HttpConstants.HttpHeaders.SHOULD_BATCH_CONTINUE_ON_ERROR, String.valueOf(serverBatchRequest.isShouldContinueOnError()));
request.setPartitionKeyDefinition(collection.getPartitionKey());
request.setNumberOfItemsInBatchRequest(serverBatchRequest.getOperations().size());
return request;
}
/**
* NOTE: Caller needs to consume it by subscribing to this Mono in order for the request to populate headers
* @param request request to populate headers to
* @param httpMethod http method
* @return Mono, which on subscription will populate the headers in the request passed in the argument.
*/
public Mono populateHeadersAsync(RxDocumentServiceRequest request, RequestVerb httpMethod) {
request.getHeaders().put(HttpConstants.HttpHeaders.X_DATE, Utils.nowAsRFC1123());
if (this.masterKeyOrResourceToken != null || this.resourceTokensMap != null
|| this.cosmosAuthorizationTokenResolver != null || this.credential != null) {
String resourceName = request.getResourceAddress();
String authorization = this.getUserAuthorizationToken(
resourceName, request.getResourceType(), httpMethod, request.getHeaders(),
AuthorizationTokenType.PrimaryMasterKey, request.properties);
try {
authorization = URLEncoder.encode(authorization, "UTF-8");
} catch (UnsupportedEncodingException e) {
throw new IllegalStateException("Failed to encode authtoken.", e);
}
request.getHeaders().put(HttpConstants.HttpHeaders.AUTHORIZATION, authorization);
}
if (this.apiType != null) {
request.getHeaders().put(HttpConstants.HttpHeaders.API_TYPE, this.apiType.toString());
}
this.populateCapabilitiesHeader(request);
if ((RequestVerb.POST.equals(httpMethod) || RequestVerb.PUT.equals(httpMethod))
&& !request.getHeaders().containsKey(HttpConstants.HttpHeaders.CONTENT_TYPE)) {
request.getHeaders().put(HttpConstants.HttpHeaders.CONTENT_TYPE, RuntimeConstants.MediaTypes.JSON);
}
if (RequestVerb.PATCH.equals(httpMethod) &&
!request.getHeaders().containsKey(HttpConstants.HttpHeaders.CONTENT_TYPE)) {
request.getHeaders().put(HttpConstants.HttpHeaders.CONTENT_TYPE, RuntimeConstants.MediaTypes.JSON_PATCH);
}
if (!request.getHeaders().containsKey(HttpConstants.HttpHeaders.ACCEPT)) {
request.getHeaders().put(HttpConstants.HttpHeaders.ACCEPT, RuntimeConstants.MediaTypes.JSON);
}
MetadataDiagnosticsContext metadataDiagnosticsCtx =
BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics);
if (this.requiresFeedRangeFiltering(request)) {
return request.getFeedRange()
.populateFeedRangeFilteringHeaders(
this.getPartitionKeyRangeCache(),
request,
this.collectionCache
.resolveCollectionAsync(metadataDiagnosticsCtx, request)
.flatMap(documentCollectionValueHolder -> {
if (documentCollectionValueHolder.v != null) {
request.setPartitionKeyDefinition(documentCollectionValueHolder.v.getPartitionKey());
}
return Mono.just(documentCollectionValueHolder);
})
)
.flatMap(this::populateAuthorizationHeader);
}
return this.populateAuthorizationHeader(request);
}
private void populateCapabilitiesHeader(RxDocumentServiceRequest request) {
if (!request.getHeaders().containsKey(HttpConstants.HttpHeaders.SDK_SUPPORTED_CAPABILITIES)) {
request
.getHeaders()
.put(HttpConstants.HttpHeaders.SDK_SUPPORTED_CAPABILITIES, HttpConstants.SDKSupportedCapabilities.SUPPORTED_CAPABILITIES);
}
}
private boolean requiresFeedRangeFiltering(RxDocumentServiceRequest request) {
if (request.getResourceType() != ResourceType.Document &&
request.getResourceType() != ResourceType.Conflict) {
return false;
}
if (request.hasFeedRangeFilteringBeenApplied()) {
return false;
}
switch (request.getOperationType()) {
case ReadFeed:
case Query:
case SqlQuery:
return request.getFeedRange() != null;
default:
return false;
}
}
@Override
public Mono populateAuthorizationHeader(RxDocumentServiceRequest request) {
if (request == null) {
throw new IllegalArgumentException("request");
}
if (this.authorizationTokenType == AuthorizationTokenType.AadToken) {
return AadTokenAuthorizationHelper.getAuthorizationToken(this.tokenCredentialCache)
.map(authorization -> {
request.getHeaders().put(HttpConstants.HttpHeaders.AUTHORIZATION, authorization);
return request;
});
} else {
return Mono.just(request);
}
}
@Override
public Mono populateAuthorizationHeader(HttpHeaders httpHeaders) {
if (httpHeaders == null) {
throw new IllegalArgumentException("httpHeaders");
}
if (this.authorizationTokenType == AuthorizationTokenType.AadToken) {
return AadTokenAuthorizationHelper.getAuthorizationToken(this.tokenCredentialCache)
.map(authorization -> {
httpHeaders.set(HttpConstants.HttpHeaders.AUTHORIZATION, authorization);
return httpHeaders;
});
}
return Mono.just(httpHeaders);
}
@Override
public AuthorizationTokenType getAuthorizationTokenType() {
return this.authorizationTokenType;
}
@Override
public String getUserAuthorizationToken(String resourceName,
ResourceType resourceType,
RequestVerb requestVerb,
Map headers,
AuthorizationTokenType tokenType,
Map properties) {
if (this.cosmosAuthorizationTokenResolver != null) {
return this.cosmosAuthorizationTokenResolver.getAuthorizationToken(requestVerb.toUpperCase(), resourceName, this.resolveCosmosResourceType(resourceType).toString(),
properties != null ? Collections.unmodifiableMap(properties) : null);
} else if (credential != null) {
return this.authorizationTokenProvider.generateKeyAuthorizationSignature(requestVerb, resourceName,
resourceType, headers);
} else if (masterKeyOrResourceToken != null && hasAuthKeyResourceToken && resourceTokensMap == null) {
return masterKeyOrResourceToken;
} else {
assert resourceTokensMap != null;
if(resourceType.equals(ResourceType.DatabaseAccount)) {
return this.firstResourceTokenFromPermissionFeed;
}
return ResourceTokenAuthorizationHelper.getAuthorizationTokenUsingResourceTokens(resourceTokensMap, requestVerb, resourceName, headers);
}
}
private CosmosResourceType resolveCosmosResourceType(ResourceType resourceType) {
CosmosResourceType cosmosResourceType =
ModelBridgeInternal.fromServiceSerializedFormat(resourceType.toString());
if (cosmosResourceType == null) {
return CosmosResourceType.SYSTEM;
}
return cosmosResourceType;
}
void captureSessionToken(RxDocumentServiceRequest request, RxDocumentServiceResponse response) {
this.sessionContainer.setSessionToken(request, response.getResponseHeaders());
}
private Mono create(RxDocumentServiceRequest request,
DocumentClientRetryPolicy documentClientRetryPolicy,
OperationContextAndListenerTuple operationContextAndListenerTuple) {
return populateHeadersAsync(request, RequestVerb.POST)
.flatMap(requestPopulated -> {
RxStoreModel storeProxy = this.getStoreProxy(requestPopulated);
if (documentClientRetryPolicy.getRetryContext() != null && documentClientRetryPolicy.getRetryContext().getRetryCount() > 0) {
documentClientRetryPolicy.getRetryContext().updateEndTime();
}
return storeProxy.processMessage(requestPopulated, operationContextAndListenerTuple);
});
}
private Mono upsert(RxDocumentServiceRequest request,
DocumentClientRetryPolicy documentClientRetryPolicy,
OperationContextAndListenerTuple operationContextAndListenerTuple) {
return populateHeadersAsync(request, RequestVerb.POST)
.flatMap(requestPopulated -> {
Map headers = requestPopulated.getHeaders();
// headers can never be null, since it will be initialized even when no
// request options are specified,
// hence using assertion here instead of exception, being in the private
// method
assert (headers != null);
headers.put(HttpConstants.HttpHeaders.IS_UPSERT, "true");
if (documentClientRetryPolicy.getRetryContext() != null && documentClientRetryPolicy.getRetryContext().getRetryCount() > 0) {
documentClientRetryPolicy.getRetryContext().updateEndTime();
}
return getStoreProxy(requestPopulated).processMessage(requestPopulated, operationContextAndListenerTuple)
.map(response -> {
this.captureSessionToken(requestPopulated, response);
return response;
}
);
});
}
private Mono replace(RxDocumentServiceRequest request, DocumentClientRetryPolicy documentClientRetryPolicy) {
return populateHeadersAsync(request, RequestVerb.PUT)
.flatMap(requestPopulated -> {
if (documentClientRetryPolicy.getRetryContext() != null && documentClientRetryPolicy.getRetryContext().getRetryCount() > 0) {
documentClientRetryPolicy.getRetryContext().updateEndTime();
}
return getStoreProxy(requestPopulated).processMessage(requestPopulated);
});
}
private Mono patch(RxDocumentServiceRequest request, DocumentClientRetryPolicy documentClientRetryPolicy) {
return populateHeadersAsync(request, RequestVerb.PATCH)
.flatMap(requestPopulated -> {
if (documentClientRetryPolicy.getRetryContext() != null && documentClientRetryPolicy.getRetryContext().getRetryCount() > 0) {
documentClientRetryPolicy.getRetryContext().updateEndTime();
}
return getStoreProxy(requestPopulated).processMessage(requestPopulated);
});
}
@Override
public Mono> createDocument(
String collectionLink,
Object document,
RequestOptions options,
boolean disableAutomaticIdGeneration) {
return wrapPointOperationWithAvailabilityStrategy(
ResourceType.Document,
OperationType.Create,
(opt, e2ecfg, clientCtxOverride, pointOperationContextForCircuitBreaker) -> createDocumentCore(
collectionLink,
document,
opt,
disableAutomaticIdGeneration,
e2ecfg,
clientCtxOverride,
pointOperationContextForCircuitBreaker
),
options,
options != null && options.getNonIdempotentWriteRetriesEnabled() != null && options.getNonIdempotentWriteRetriesEnabled(),
collectionLink
);
}
private Mono> createDocumentCore(
String collectionLink,
Object document,
RequestOptions options,
boolean disableAutomaticIdGeneration,
CosmosEndToEndOperationLatencyPolicyConfig endToEndPolicyConfig,
DiagnosticsClientContext clientContextOverride,
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker) {
ScopedDiagnosticsFactory scopedDiagnosticsFactory = new ScopedDiagnosticsFactory(clientContextOverride, false);
DocumentClientRetryPolicy requestRetryPolicy =
this.resetSessionTokenRetryPolicy.getRequestPolicy(scopedDiagnosticsFactory);
RequestOptions nonNullRequestOptions = options != null ? options : new RequestOptions();
if (nonNullRequestOptions.getPartitionKey() == null) {
requestRetryPolicy = new PartitionKeyMismatchRetryPolicy(collectionCache, requestRetryPolicy, collectionLink, nonNullRequestOptions);
}
DocumentClientRetryPolicy finalRetryPolicyInstance = requestRetryPolicy;
AtomicReference requestReference = new AtomicReference<>();
return handleCircuitBreakingFeedbackForPointOperation(getPointOperationResponseMonoWithE2ETimeout(
nonNullRequestOptions,
endToEndPolicyConfig,
ObservableHelper.inlineIfPossibleAsObs(() ->
createDocumentInternal(
collectionLink,
document,
nonNullRequestOptions,
disableAutomaticIdGeneration,
finalRetryPolicyInstance,
scopedDiagnosticsFactory,
requestReference,
pointOperationContextForCircuitBreaker),
requestRetryPolicy),
scopedDiagnosticsFactory
), requestReference);
}
private Mono> createDocumentInternal(
String collectionLink,
Object document,
RequestOptions options,
boolean disableAutomaticIdGeneration,
DocumentClientRetryPolicy requestRetryPolicy,
DiagnosticsClientContext clientContextOverride,
AtomicReference documentServiceRequestReference,
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker) {
try {
logger.debug("Creating a Document. collectionLink: [{}]", collectionLink);
Mono>> requestToDocumentCollectionObs = getCreateDocumentRequest(
requestRetryPolicy,
collectionLink,
document,
options,
disableAutomaticIdGeneration,
OperationType.Create,
clientContextOverride,
pointOperationContextForCircuitBreaker);
return requestToDocumentCollectionObs
.flatMap(requestToDocumentCollection -> {
RxDocumentServiceRequest request = requestToDocumentCollection.getT1();
Utils.ValueHolder documentCollectionValueHolder = requestToDocumentCollection.getT2();
if (documentCollectionValueHolder == null || documentCollectionValueHolder.v == null) {
return Mono.error(new IllegalStateException("documentCollectionValueHolder or documentCollectionValueHolder.v cannot be null"));
}
return this.partitionKeyRangeCache.tryLookupAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), documentCollectionValueHolder.v.getResourceId(), null, null)
.flatMap(collectionRoutingMapValueHolder -> {
if (collectionRoutingMapValueHolder == null || collectionRoutingMapValueHolder.v == null) {
return Mono.error(new IllegalStateException("collectionRoutingMapValueHolder or collectionRoutingMapValueHolder.v cannot be null"));
}
options.setPartitionKeyDefinition(documentCollectionValueHolder.v.getPartitionKey());
addPartitionLevelUnavailableRegionsForRequest(request, options, collectionRoutingMapValueHolder.v, requestRetryPolicy);
documentServiceRequestReference.set(request);
request.requestContext.setPointOperationContext(pointOperationContextForCircuitBreaker);
// needs to be after onBeforeSendRequest since CosmosDiagnostics instance needs to be wired
// to the RxDocumentServiceRequest instance
mergeContextInformationIntoDiagnosticsForPointRequest(request, pointOperationContextForCircuitBreaker);
return create(request, requestRetryPolicy, getOperationContextAndListenerTuple(options));
})
.map(serviceResponse -> toResourceResponse(serviceResponse, Document.class));
});
} catch (Exception e) {
logger.debug("Failure in creating a document due to [{}]", e.getMessage(), e);
return Mono.error(e);
}
}
private static Mono getPointOperationResponseMonoWithE2ETimeout(
RequestOptions requestOptions,
CosmosEndToEndOperationLatencyPolicyConfig endToEndPolicyConfig,
Mono rxDocumentServiceResponseMono,
ScopedDiagnosticsFactory scopedDiagnosticsFactory) {
requestOptions.setCosmosEndToEndLatencyPolicyConfig(endToEndPolicyConfig);
if (endToEndPolicyConfig != null && endToEndPolicyConfig.isEnabled()) {
Duration endToEndTimeout = endToEndPolicyConfig.getEndToEndOperationTimeout();
if (endToEndTimeout.isNegative()) {
CosmosDiagnostics latestCosmosDiagnosticsSnapshot = scopedDiagnosticsFactory.getMostRecentlyCreatedDiagnostics();
if (latestCosmosDiagnosticsSnapshot == null) {
scopedDiagnosticsFactory.createDiagnostics();
}
return Mono.error(getNegativeTimeoutException(scopedDiagnosticsFactory.getMostRecentlyCreatedDiagnostics(), endToEndTimeout));
}
return rxDocumentServiceResponseMono
.timeout(endToEndTimeout)
.onErrorMap(throwable -> getCancellationExceptionForPointOperations(
scopedDiagnosticsFactory,
throwable,
requestOptions.getMarkE2ETimeoutInRequestContextCallbackHook()));
}
return rxDocumentServiceResponseMono;
}
private Mono handleCircuitBreakingFeedbackForPointOperation(
Mono response,
AtomicReference requestReference) {
return response
.doOnSuccess(ignore -> {
if (this.globalPartitionEndpointManagerForCircuitBreaker.isPartitionLevelCircuitBreakingApplicable(requestReference.get())) {
RxDocumentServiceRequest succeededRequest = requestReference.get();
checkNotNull(succeededRequest.requestContext, "Argument 'succeededRequest.requestContext' must not be null!");
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker = succeededRequest.requestContext.getPointOperationContextForCircuitBreaker();
checkNotNull(pointOperationContextForCircuitBreaker, "Argument 'pointOperationContextForCircuitBreaker' must not be null!");
pointOperationContextForCircuitBreaker.setHasOperationSeenSuccess();
this.globalPartitionEndpointManagerForCircuitBreaker.handleLocationSuccessForPartitionKeyRange(succeededRequest);
}
})
.doOnError(throwable -> {
if (throwable instanceof OperationCancelledException) {
if (this.globalPartitionEndpointManagerForCircuitBreaker.isPartitionLevelCircuitBreakingApplicable(requestReference.get())) {
RxDocumentServiceRequest failedRequest = requestReference.get();
checkNotNull(failedRequest.requestContext, "Argument 'failedRequest.requestContext' must not be null!");
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker = failedRequest.requestContext.getPointOperationContextForCircuitBreaker();
checkNotNull(pointOperationContextForCircuitBreaker, "Argument 'pointOperationContextForCircuitBreaker' must not be null!");
if (pointOperationContextForCircuitBreaker.isThresholdBasedAvailabilityStrategyEnabled()) {
if (!pointOperationContextForCircuitBreaker.isRequestHedged() && pointOperationContextForCircuitBreaker.getHasOperationSeenSuccess()) {
this.handleLocationCancellationExceptionForPartitionKeyRange(failedRequest);
}
} else {
this.handleLocationCancellationExceptionForPartitionKeyRange(failedRequest);
}
}
}
})
.doFinally(signalType -> {
if (signalType != SignalType.CANCEL) {
return;
}
if (this.globalPartitionEndpointManagerForCircuitBreaker.isPartitionLevelCircuitBreakingApplicable(requestReference.get())) {
RxDocumentServiceRequest failedRequest = requestReference.get();
checkNotNull(failedRequest.requestContext, "Argument 'failedRequest.requestContext' must not be null!");
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker = failedRequest.requestContext.getPointOperationContextForCircuitBreaker();
checkNotNull(pointOperationContextForCircuitBreaker, "Argument 'pointOperationContextForCircuitBreaker' must not be null!");
// scoping the handling of CANCEL signal handling for reasons outside of end-to-end operation timeout
// to purely operations which have end-to-end operation timeout enabled
if (pointOperationContextForCircuitBreaker.isThresholdBasedAvailabilityStrategyEnabled()) {
if (!pointOperationContextForCircuitBreaker.isRequestHedged() && pointOperationContextForCircuitBreaker.getHasOperationSeenSuccess()) {
this.handleLocationCancellationExceptionForPartitionKeyRange(failedRequest);
}
}
}
});
}
private Mono> handleCircuitBreakingFeedbackForFeedOperationWithAvailabilityStrategy(Mono> response, RxDocumentServiceRequest request) {
return response
.doOnSuccess(nonTransientFeedOperationResult -> {
if (this.globalPartitionEndpointManagerForCircuitBreaker.isPartitionLevelCircuitBreakingApplicable(request)) {
if (!nonTransientFeedOperationResult.isError()) {
checkNotNull(request, "Argument 'request' cannot be null!");
checkNotNull(request.requestContext, "Argument 'request.requestContext' cannot be null!");
FeedOperationContextForCircuitBreaker feedOperationContextForCircuitBreaker
= request.requestContext.getFeedOperationContextForCircuitBreaker();
checkNotNull(feedOperationContextForCircuitBreaker, "Argument 'feedOperationContextForCircuitBreaker' cannot be null!");
feedOperationContextForCircuitBreaker.addPartitionKeyRangeWithSuccess(request.requestContext.resolvedPartitionKeyRange, request.getResourceId());
this.globalPartitionEndpointManagerForCircuitBreaker.handleLocationSuccessForPartitionKeyRange(request);
}
}
})
.doFinally(signalType -> {
if (signalType != SignalType.CANCEL) {
return;
}
if (this.globalPartitionEndpointManagerForCircuitBreaker.isPartitionLevelCircuitBreakingApplicable(request)) {
checkNotNull(request, "Argument 'request' cannot be null!");
checkNotNull(request.requestContext, "Argument 'request.requestContext' cannot be null!");
FeedOperationContextForCircuitBreaker feedOperationContextForCircuitBreaker
= request.requestContext.getFeedOperationContextForCircuitBreaker();
checkNotNull(feedOperationContextForCircuitBreaker, "Argument 'feedOperationContextForCircuitBreaker' cannot be null!");
if (!feedOperationContextForCircuitBreaker.getIsRequestHedged()
&& feedOperationContextForCircuitBreaker.isThresholdBasedAvailabilityStrategyEnabled()
&& feedOperationContextForCircuitBreaker.hasPartitionKeyRangeSeenSuccess(request.requestContext.resolvedPartitionKeyRange, request.getResourceId())) {
this.handleLocationCancellationExceptionForPartitionKeyRange(request);
}
}
});
}
private static Throwable getCancellationExceptionForPointOperations(
ScopedDiagnosticsFactory scopedDiagnosticsFactory,
Throwable throwable,
AtomicReference markE2ETimeoutInRequestContextCallbackHook) {
Throwable unwrappedException = reactor.core.Exceptions.unwrap(throwable);
if (unwrappedException instanceof TimeoutException) {
CosmosException exception = new OperationCancelledException();
exception.setStackTrace(throwable.getStackTrace());
Runnable actualCallback = markE2ETimeoutInRequestContextCallbackHook.get();
if (actualCallback != null) {
logger.trace("Calling actual Mark E2E timeout callback");
actualCallback.run();
}
// For point operations
// availabilityStrategy sits on top of e2eTimeoutPolicy
// e2eTimeoutPolicy sits on top of client retry policy
// for each e2eTimeoutPolicy wrap, we are going to create one distinct ScopedDiagnosticsFactory
// so for each scopedDiagnosticsFactory being used here, there will only be max one CosmosDiagnostics being tracked
CosmosDiagnostics lastDiagnosticsSnapshot = scopedDiagnosticsFactory.getMostRecentlyCreatedDiagnostics();
if (lastDiagnosticsSnapshot == null) {
scopedDiagnosticsFactory.createDiagnostics();
}
BridgeInternal.setCosmosDiagnostics(exception, scopedDiagnosticsFactory.getMostRecentlyCreatedDiagnostics());
return exception;
}
return throwable;
}
private static CosmosException getNegativeTimeoutException(CosmosDiagnostics cosmosDiagnostics, Duration negativeTimeout) {
checkNotNull(negativeTimeout, "Argument 'negativeTimeout' must not be null");
checkArgument(
negativeTimeout.isNegative(),
"This exception should only be used for negative timeouts");
String message = String.format("Negative timeout '%s' provided.", negativeTimeout);
CosmosException exception = new OperationCancelledException(message, null);
BridgeInternal.setSubStatusCode(exception, HttpConstants.SubStatusCodes.NEGATIVE_TIMEOUT_PROVIDED);
if (cosmosDiagnostics != null) {
BridgeInternal.setCosmosDiagnostics(exception, cosmosDiagnostics);
}
return exception;
}
@Override
public Mono> upsertDocument(String collectionLink, Object document,
RequestOptions options, boolean disableAutomaticIdGeneration) {
return wrapPointOperationWithAvailabilityStrategy(
ResourceType.Document,
OperationType.Upsert,
(opt, e2ecfg, clientCtxOverride, pointOperationContextForCircuitBreaker) -> upsertDocumentCore(
collectionLink, document, opt, disableAutomaticIdGeneration, e2ecfg, clientCtxOverride, pointOperationContextForCircuitBreaker),
options,
options != null && options.getNonIdempotentWriteRetriesEnabled() != null && options.getNonIdempotentWriteRetriesEnabled(),
collectionLink
);
}
private Mono> upsertDocumentCore(
String collectionLink,
Object document,
RequestOptions options,
boolean disableAutomaticIdGeneration,
CosmosEndToEndOperationLatencyPolicyConfig endToEndPolicyConfig,
DiagnosticsClientContext clientContextOverride,
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker) {
RequestOptions nonNullRequestOptions = options != null ? options : new RequestOptions();
ScopedDiagnosticsFactory scopedDiagnosticsFactory = new ScopedDiagnosticsFactory(clientContextOverride, false);
DocumentClientRetryPolicy requestRetryPolicy = this.resetSessionTokenRetryPolicy.getRequestPolicy(scopedDiagnosticsFactory);
if (nonNullRequestOptions.getPartitionKey() == null) {
requestRetryPolicy = new PartitionKeyMismatchRetryPolicy(collectionCache, requestRetryPolicy, collectionLink, nonNullRequestOptions);
}
DocumentClientRetryPolicy finalRetryPolicyInstance = requestRetryPolicy;
AtomicReference requestReference = new AtomicReference<>();
return handleCircuitBreakingFeedbackForPointOperation(getPointOperationResponseMonoWithE2ETimeout(
nonNullRequestOptions,
endToEndPolicyConfig,
ObservableHelper.inlineIfPossibleAsObs(
() -> upsertDocumentInternal(
collectionLink,
document,
nonNullRequestOptions,
disableAutomaticIdGeneration,
finalRetryPolicyInstance,
scopedDiagnosticsFactory,
requestReference,
pointOperationContextForCircuitBreaker),
finalRetryPolicyInstance),
scopedDiagnosticsFactory), requestReference);
}
private Mono> upsertDocumentInternal(
String collectionLink,
Object document,
RequestOptions options,
boolean disableAutomaticIdGeneration,
DocumentClientRetryPolicy retryPolicyInstance,
DiagnosticsClientContext clientContextOverride,
AtomicReference requestReference,
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker) {
try {
logger.debug("Upserting a Document. collectionLink: [{}]", collectionLink);
Mono>> requestToDocumentCollectionObs =
getCreateDocumentRequest(
retryPolicyInstance,
collectionLink,
document,
options,
disableAutomaticIdGeneration,
OperationType.Upsert,
clientContextOverride,
pointOperationContextForCircuitBreaker);
return requestToDocumentCollectionObs
.flatMap(requestToDocumentCollection -> {
RxDocumentServiceRequest request = requestToDocumentCollection.getT1();
Utils.ValueHolder documentCollectionValueHolder = requestToDocumentCollection.getT2();
if (documentCollectionValueHolder == null || documentCollectionValueHolder.v == null) {
return Mono.error(new IllegalStateException("documentCollectionValueHolder or documentCollectionValueHolder.v cannot be null"));
}
return this.partitionKeyRangeCache.tryLookupAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), documentCollectionValueHolder.v.getResourceId(), null, null)
.flatMap(collectionRoutingMapValueHolder -> {
if (collectionRoutingMapValueHolder == null || collectionRoutingMapValueHolder.v == null) {
return Mono.error(new IllegalStateException("collectionRoutingMapValueHolder or collectionRoutingMapValueHolder.v cannot be null"));
}
options.setPartitionKeyDefinition(documentCollectionValueHolder.v.getPartitionKey());
addPartitionLevelUnavailableRegionsForRequest(request, options, collectionRoutingMapValueHolder.v, retryPolicyInstance);
request.requestContext.setPointOperationContext(pointOperationContextForCircuitBreaker);
requestReference.set(request);
// needs to be after onBeforeSendRequest since CosmosDiagnostics instance needs to be wired
// to the RxDocumentServiceRequest instance
mergeContextInformationIntoDiagnosticsForPointRequest(request, pointOperationContextForCircuitBreaker);
return upsert(request, retryPolicyInstance, getOperationContextAndListenerTuple(options));
})
.map(serviceResponse -> toResourceResponse(serviceResponse, Document.class));
});
} catch (Exception e) {
logger.debug("Failure in upserting a document due to [{}]", e.getMessage(), e);
return Mono.error(e);
}
}
@Override
public Mono> replaceDocument(String documentLink, Object document,
RequestOptions options) {
String collectionLink = Utils.getCollectionName(documentLink);
return wrapPointOperationWithAvailabilityStrategy(
ResourceType.Document,
OperationType.Replace,
(opt, e2ecfg, clientCtxOverride, pointOperationContextForCircuitBreaker) -> replaceDocumentCore(
documentLink,
document,
opt,
e2ecfg,
clientCtxOverride,
pointOperationContextForCircuitBreaker),
options,
options != null && options.getNonIdempotentWriteRetriesEnabled() != null && options.getNonIdempotentWriteRetriesEnabled(),
collectionLink
);
}
private Mono> replaceDocumentCore(
String documentLink,
Object document,
RequestOptions options,
CosmosEndToEndOperationLatencyPolicyConfig endToEndPolicyConfig,
DiagnosticsClientContext clientContextOverride,
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker) {
RequestOptions nonNullRequestOptions = options != null ? options : new RequestOptions();
ScopedDiagnosticsFactory scopedDiagnosticsFactory = new ScopedDiagnosticsFactory(clientContextOverride, false);
DocumentClientRetryPolicy requestRetryPolicy =
this.resetSessionTokenRetryPolicy.getRequestPolicy(scopedDiagnosticsFactory);
if (nonNullRequestOptions.getPartitionKey() == null) {
String collectionLink = Utils.getCollectionName(documentLink);
requestRetryPolicy = new PartitionKeyMismatchRetryPolicy(
collectionCache, requestRetryPolicy, collectionLink, nonNullRequestOptions);
}
DocumentClientRetryPolicy finalRequestRetryPolicy = requestRetryPolicy;
AtomicReference requestReference = new AtomicReference<>();
return handleCircuitBreakingFeedbackForPointOperation(getPointOperationResponseMonoWithE2ETimeout(
nonNullRequestOptions,
endToEndPolicyConfig,
ObservableHelper.inlineIfPossibleAsObs(
() -> replaceDocumentInternal(
documentLink,
document,
nonNullRequestOptions,
finalRequestRetryPolicy,
endToEndPolicyConfig,
scopedDiagnosticsFactory,
requestReference,
pointOperationContextForCircuitBreaker),
requestRetryPolicy),
scopedDiagnosticsFactory), requestReference);
}
private Mono> replaceDocumentInternal(
String documentLink,
Object document,
RequestOptions options,
DocumentClientRetryPolicy retryPolicyInstance,
CosmosEndToEndOperationLatencyPolicyConfig endToEndPolicyConfig,
DiagnosticsClientContext clientContextOverride,
AtomicReference requestReference,
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker) {
try {
if (StringUtils.isEmpty(documentLink)) {
throw new IllegalArgumentException("documentLink");
}
if (document == null) {
throw new IllegalArgumentException("document");
}
Document typedDocument = Document.fromObject(document, options.getEffectiveItemSerializer());
return this.replaceDocumentInternal(
documentLink,
typedDocument,
options,
retryPolicyInstance,
clientContextOverride,
requestReference,
pointOperationContextForCircuitBreaker);
} catch (Exception e) {
logger.debug("Failure in replacing a document due to [{}]", e.getMessage());
return Mono.error(e);
}
}
@Override
public Mono> replaceDocument(Document document, RequestOptions options) {
String collectionLink = Utils.getCollectionName(document.getSelfLink());
return wrapPointOperationWithAvailabilityStrategy(
ResourceType.Document,
OperationType.Replace,
(opt, e2ecfg, clientCtxOverride, pointOperationContextForCircuitBreaker) -> replaceDocumentCore(
document,
opt,
e2ecfg,
clientCtxOverride,
pointOperationContextForCircuitBreaker
),
options,
options != null && options.getNonIdempotentWriteRetriesEnabled() != null && options.getNonIdempotentWriteRetriesEnabled(),
collectionLink
);
}
private Mono> replaceDocumentCore(
Document document,
RequestOptions options,
CosmosEndToEndOperationLatencyPolicyConfig endToEndPolicyConfig,
DiagnosticsClientContext clientContextOverride,
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker) {
DocumentClientRetryPolicy requestRetryPolicy =
this.resetSessionTokenRetryPolicy.getRequestPolicy(clientContextOverride);
if (options == null || options.getPartitionKey() == null) {
String collectionLink = document.getSelfLink();
requestRetryPolicy = new PartitionKeyMismatchRetryPolicy(
collectionCache, requestRetryPolicy, collectionLink, options);
}
DocumentClientRetryPolicy finalRequestRetryPolicy = requestRetryPolicy;
AtomicReference requestReference = new AtomicReference<>();
return handleCircuitBreakingFeedbackForPointOperation(ObservableHelper.inlineIfPossibleAsObs(
() -> replaceDocumentInternal(
document,
options,
finalRequestRetryPolicy,
endToEndPolicyConfig,
clientContextOverride,
requestReference,
pointOperationContextForCircuitBreaker),
requestRetryPolicy), requestReference);
}
private Mono> replaceDocumentInternal(
Document document,
RequestOptions options,
DocumentClientRetryPolicy retryPolicyInstance,
CosmosEndToEndOperationLatencyPolicyConfig endToEndPolicyConfig,
DiagnosticsClientContext clientContextOverride,
AtomicReference requestReference,
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker) {
try {
if (document == null) {
throw new IllegalArgumentException("document");
}
return this.replaceDocumentInternal(
document.getSelfLink(),
document,
options,
retryPolicyInstance,
clientContextOverride,
requestReference,
pointOperationContextForCircuitBreaker);
} catch (Exception e) {
logger.debug("Failure in replacing a database due to [{}]", e.getMessage());
return Mono.error(e);
}
}
private Mono> replaceDocumentInternal(
String documentLink,
Document document,
RequestOptions options,
DocumentClientRetryPolicy retryPolicyInstance,
DiagnosticsClientContext clientContextOverride,
AtomicReference requestReference,
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker) {
if (document == null) {
throw new IllegalArgumentException("document");
}
logger.debug("Replacing a Document. documentLink: [{}]", documentLink);
final String path = Utils.joinPath(documentLink, null);
final Map requestHeaders =
getRequestHeaders(options, ResourceType.Document, OperationType.Replace);
Instant serializationStartTimeUTC = Instant.now();
Consumer> onAfterSerialization = null;
if (options != null) {
String trackingId = options.getTrackingId();
if (trackingId != null && !trackingId.isEmpty()) {
onAfterSerialization = (node) -> node.put(Constants.Properties.TRACKING_ID, trackingId);
}
}
ByteBuffer content = document.serializeJsonToByteBuffer(options.getEffectiveItemSerializer(), onAfterSerialization, false);
Instant serializationEndTime = Instant.now();
SerializationDiagnosticsContext.SerializationDiagnostics serializationDiagnostics =
new SerializationDiagnosticsContext.SerializationDiagnostics(
serializationStartTimeUTC,
serializationEndTime,
SerializationDiagnosticsContext.SerializationType.ITEM_SERIALIZATION);
final RxDocumentServiceRequest request = RxDocumentServiceRequest.create(
getEffectiveClientContext(clientContextOverride),
OperationType.Replace, ResourceType.Document, path, requestHeaders, options, content);
if (options != null && options.getNonIdempotentWriteRetriesEnabled() != null && options.getNonIdempotentWriteRetriesEnabled()) {
request.setNonIdempotentWriteRetriesEnabled(true);
}
if (options != null) {
options.getMarkE2ETimeoutInRequestContextCallbackHook().set(
() -> request.requestContext.setIsRequestCancelledOnTimeout(new AtomicBoolean(true)));
request.requestContext.setExcludeRegions(options.getExcludedRegions());
request.requestContext.setKeywordIdentifiers(options.getKeywordIdentifiers());
}
SerializationDiagnosticsContext serializationDiagnosticsContext =
BridgeInternal.getSerializationDiagnosticsContext(request.requestContext.cosmosDiagnostics);
if (serializationDiagnosticsContext != null) {
serializationDiagnosticsContext.addSerializationDiagnostics(serializationDiagnostics);
}
if (retryPolicyInstance != null) {
retryPolicyInstance.onBeforeSendRequest(request);
}
Mono> collectionObs =
collectionCache.resolveCollectionAsync(
BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics),
request);
Mono requestObs =
addPartitionKeyInformation(request, content, document, options, collectionObs, pointOperationContextForCircuitBreaker);
return collectionObs
.flatMap(documentCollectionValueHolder -> {
if (documentCollectionValueHolder == null || documentCollectionValueHolder.v == null) {
return Mono.error(new IllegalStateException("documentCollectionValueHolder or documentCollectionValueHolder.v cannot be null"));
}
return this.partitionKeyRangeCache.tryLookupAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), documentCollectionValueHolder.v.getResourceId(), null, null)
.flatMap(collectionRoutingMapValueHolder -> {
if (collectionRoutingMapValueHolder == null || collectionRoutingMapValueHolder.v == null) {
return Mono.error(new IllegalStateException("collectionRoutingMapValueHolder or collectionRoutingMapValueHolder.v cannot be null"));
}
return requestObs.flatMap(req -> {
options.setPartitionKeyDefinition(documentCollectionValueHolder.v.getPartitionKey());
addPartitionLevelUnavailableRegionsForRequest(req, options, collectionRoutingMapValueHolder.v, retryPolicyInstance);
req.requestContext.setPointOperationContext(pointOperationContextForCircuitBreaker);
requestReference.set(req);
// needs to be after onBeforeSendRequest since CosmosDiagnostics instance needs to be wired
// to the RxDocumentServiceRequest instance
mergeContextInformationIntoDiagnosticsForPointRequest(request, pointOperationContextForCircuitBreaker);
return replace(request, retryPolicyInstance);
})
.map(resp -> toResourceResponse(resp, Document.class));
});
});
}
private CosmosEndToEndOperationLatencyPolicyConfig getEndToEndOperationLatencyPolicyConfig(
RequestOptions options,
ResourceType resourceType,
OperationType operationType) {
return this.getEffectiveEndToEndOperationLatencyPolicyConfig(
options != null ? options.getCosmosEndToEndLatencyPolicyConfig() : null,
resourceType,
operationType);
}
private CosmosEndToEndOperationLatencyPolicyConfig getEffectiveEndToEndOperationLatencyPolicyConfig(
CosmosEndToEndOperationLatencyPolicyConfig policyConfig,
ResourceType resourceType,
OperationType operationType) {
if (policyConfig != null) {
return policyConfig;
}
if (resourceType != ResourceType.Document) {
return null;
}
if (!operationType.isPointOperation() && Configs.isDefaultE2ETimeoutDisabledForNonPointOperations()) {
return null;
}
return this.cosmosEndToEndOperationLatencyPolicyConfig;
}
@Override
public Mono> patchDocument(String documentLink,
CosmosPatchOperations cosmosPatchOperations,
RequestOptions options) {
String collectionLink = Utils.getCollectionName(documentLink);
return wrapPointOperationWithAvailabilityStrategy(
ResourceType.Document,
OperationType.Patch,
(opt, e2ecfg, clientCtxOverride, pointOperationContextForCircuitBreaker) -> patchDocumentCore(
documentLink,
cosmosPatchOperations,
opt,
e2ecfg,
clientCtxOverride,
pointOperationContextForCircuitBreaker),
options,
options != null && options.getNonIdempotentWriteRetriesEnabled() != null && options.getNonIdempotentWriteRetriesEnabled(),
collectionLink
);
}
private Mono> patchDocumentCore(
String documentLink,
CosmosPatchOperations cosmosPatchOperations,
RequestOptions options,
CosmosEndToEndOperationLatencyPolicyConfig endToEndPolicyConfig,
DiagnosticsClientContext clientContextOverride,
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker) {
RequestOptions nonNullRequestOptions = options != null ? options : new RequestOptions();
ScopedDiagnosticsFactory scopedDiagnosticsFactory = new ScopedDiagnosticsFactory(clientContextOverride, false);
DocumentClientRetryPolicy documentClientRetryPolicy = this.resetSessionTokenRetryPolicy.getRequestPolicy(scopedDiagnosticsFactory);
AtomicReference requestReference = new AtomicReference<>();
return handleCircuitBreakingFeedbackForPointOperation(
getPointOperationResponseMonoWithE2ETimeout(
nonNullRequestOptions,
endToEndPolicyConfig,
ObservableHelper.inlineIfPossibleAsObs(
() -> patchDocumentInternal(
documentLink,
cosmosPatchOperations,
nonNullRequestOptions,
documentClientRetryPolicy,
scopedDiagnosticsFactory,
requestReference,
pointOperationContextForCircuitBreaker),
documentClientRetryPolicy),
scopedDiagnosticsFactory), requestReference);
}
private Mono> patchDocumentInternal(
String documentLink,
CosmosPatchOperations cosmosPatchOperations,
RequestOptions options,
DocumentClientRetryPolicy retryPolicyInstance,
DiagnosticsClientContext clientContextOverride,
AtomicReference requestReference,
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker) {
checkArgument(StringUtils.isNotEmpty(documentLink), "expected non empty documentLink");
checkNotNull(cosmosPatchOperations, "expected non null cosmosPatchOperations");
logger.debug("Running patch operations on Document. documentLink: [{}]", documentLink);
final String path = Utils.joinPath(documentLink, null);
final Map requestHeaders =
getRequestHeaders(options, ResourceType.Document, OperationType.Patch);
Instant serializationStartTimeUTC = Instant.now();
ByteBuffer content = ByteBuffer.wrap(
PatchUtil.serializeCosmosPatchToByteArray(cosmosPatchOperations, options));
Instant serializationEndTime = Instant.now();
SerializationDiagnosticsContext.SerializationDiagnostics serializationDiagnostics =
new SerializationDiagnosticsContext.SerializationDiagnostics(
serializationStartTimeUTC,
serializationEndTime,
SerializationDiagnosticsContext.SerializationType.ITEM_SERIALIZATION);
final RxDocumentServiceRequest request = RxDocumentServiceRequest.create(
clientContextOverride,
OperationType.Patch,
ResourceType.Document,
path,
requestHeaders,
options,
content);
if (options != null && options.getNonIdempotentWriteRetriesEnabled() != null && options.getNonIdempotentWriteRetriesEnabled()) {
request.setNonIdempotentWriteRetriesEnabled(true);
}
if (options != null) {
options.getMarkE2ETimeoutInRequestContextCallbackHook().set(
() -> request.requestContext.setIsRequestCancelledOnTimeout(new AtomicBoolean(true)));
request.requestContext.setExcludeRegions(options.getExcludedRegions());
request.requestContext.setKeywordIdentifiers(options.getKeywordIdentifiers());
}
if (retryPolicyInstance != null) {
retryPolicyInstance.onBeforeSendRequest(request);
}
SerializationDiagnosticsContext serializationDiagnosticsContext =
BridgeInternal.getSerializationDiagnosticsContext(request.requestContext.cosmosDiagnostics);
if (serializationDiagnosticsContext != null) {
serializationDiagnosticsContext.addSerializationDiagnostics(serializationDiagnostics);
}
Mono> collectionObs = collectionCache.resolveCollectionAsync(
BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), request);
// options will always have partition key info, so contentAsByteBuffer can be null and is not needed.
Mono requestObs = addPartitionKeyInformation(
request,
null,
null,
options,
collectionObs,
pointOperationContextForCircuitBreaker);
return collectionObs
.flatMap(documentCollectionValueHolder -> {
if (documentCollectionValueHolder == null || documentCollectionValueHolder.v == null) {
return Mono.error(new IllegalStateException("documentCollectionValueHolder or documentCollectionValueHolder.v cannot be null"));
}
return this.partitionKeyRangeCache.tryLookupAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), documentCollectionValueHolder.v.getResourceId(), null, null)
.flatMap(collectionRoutingMapValueHolder -> {
if (collectionRoutingMapValueHolder == null || collectionRoutingMapValueHolder.v == null) {
return Mono.error(new IllegalStateException("collectionRoutingMapValueHolder or collectionRoutingMapValueHolder.v cannot be null"));
}
return requestObs
.flatMap(req -> {
options.setPartitionKeyDefinition(documentCollectionValueHolder.v.getPartitionKey());
addPartitionLevelUnavailableRegionsForRequest(req, options, collectionRoutingMapValueHolder.v, retryPolicyInstance);
req.requestContext.setPointOperationContext(pointOperationContextForCircuitBreaker);
requestReference.set(req);
// needs to be after onBeforeSendRequest since CosmosDiagnostics instance needs to be wired
// to the RxDocumentServiceRequest instance
mergeContextInformationIntoDiagnosticsForPointRequest(request, pointOperationContextForCircuitBreaker);
return patch(request, retryPolicyInstance);
})
.map(resp -> toResourceResponse(resp, Document.class));
});
});
}
@Override
public Mono> deleteDocument(String documentLink, RequestOptions options) {
String collectionLink = Utils.getCollectionName(documentLink);
return wrapPointOperationWithAvailabilityStrategy(
ResourceType.Document,
OperationType.Delete,
(opt, e2ecfg, clientCtxOverride, pointOperationContextForCircuitBreaker) -> deleteDocumentCore(
documentLink,
null,
opt,
e2ecfg,
clientCtxOverride,
pointOperationContextForCircuitBreaker
),
options,
options != null && options.getNonIdempotentWriteRetriesEnabled() != null && options.getNonIdempotentWriteRetriesEnabled(),
collectionLink
);
}
@Override
public Mono> deleteDocument(String documentLink, InternalObjectNode internalObjectNode, RequestOptions options) {
String collectionLink = Utils.getCollectionName(documentLink);
return wrapPointOperationWithAvailabilityStrategy(
ResourceType.Document,
OperationType.Delete,
(opt, e2ecfg, clientCtxOverride, pointOperationContextForCircuitBreaker) -> deleteDocumentCore(
documentLink,
internalObjectNode,
opt,
e2ecfg,
clientCtxOverride,
pointOperationContextForCircuitBreaker),
options,
options != null && options.getNonIdempotentWriteRetriesEnabled() != null && options.getNonIdempotentWriteRetriesEnabled(),
collectionLink
);
}
private Mono> deleteDocumentCore(
String documentLink,
InternalObjectNode internalObjectNode,
RequestOptions options,
CosmosEndToEndOperationLatencyPolicyConfig endToEndPolicyConfig,
DiagnosticsClientContext clientContextOverride,
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker) {
RequestOptions nonNullRequestOptions = options != null ? options : new RequestOptions();
ScopedDiagnosticsFactory scopedDiagnosticsFactory = new ScopedDiagnosticsFactory(clientContextOverride, false);
DocumentClientRetryPolicy requestRetryPolicy =
this.resetSessionTokenRetryPolicy.getRequestPolicy(scopedDiagnosticsFactory);
AtomicReference requestReference = new AtomicReference<>();
return handleCircuitBreakingFeedbackForPointOperation(getPointOperationResponseMonoWithE2ETimeout(
nonNullRequestOptions,
endToEndPolicyConfig,
ObservableHelper.inlineIfPossibleAsObs(
() -> deleteDocumentInternal(
documentLink,
internalObjectNode,
nonNullRequestOptions,
requestRetryPolicy,
scopedDiagnosticsFactory,
requestReference,
pointOperationContextForCircuitBreaker),
requestRetryPolicy),
scopedDiagnosticsFactory), requestReference);
}
private Mono> deleteDocumentInternal(
String documentLink,
InternalObjectNode internalObjectNode,
RequestOptions options,
DocumentClientRetryPolicy retryPolicyInstance,
DiagnosticsClientContext clientContextOverride,
AtomicReference requestReference,
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker) {
try {
if (StringUtils.isEmpty(documentLink)) {
throw new IllegalArgumentException("documentLink");
}
logger.debug("Deleting a Document. documentLink: [{}]", documentLink);
String path = Utils.joinPath(documentLink, null);
Map requestHeaders = this.getRequestHeaders(options, ResourceType.Document, OperationType.Delete);
RxDocumentServiceRequest request = RxDocumentServiceRequest.create(
getEffectiveClientContext(clientContextOverride),
OperationType.Delete, ResourceType.Document, path, requestHeaders, options);
if (options != null && options.getNonIdempotentWriteRetriesEnabled() != null && options.getNonIdempotentWriteRetriesEnabled()) {
request.setNonIdempotentWriteRetriesEnabled(true);
}
if (options != null) {
options.getMarkE2ETimeoutInRequestContextCallbackHook().set(
() -> request.requestContext.setIsRequestCancelledOnTimeout(new AtomicBoolean(true)));
request.requestContext.setExcludeRegions(options.getExcludedRegions());
request.requestContext.setKeywordIdentifiers(options.getKeywordIdentifiers());
}
if (retryPolicyInstance != null) {
retryPolicyInstance.onBeforeSendRequest(request);
}
Mono> collectionObs = collectionCache.resolveCollectionAsync(
BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics),
request);
Mono requestObs = addPartitionKeyInformation(
request, null, internalObjectNode, options, collectionObs, pointOperationContextForCircuitBreaker);
return collectionObs
.flatMap(documentCollectionValueHolder -> this.partitionKeyRangeCache.tryLookupAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), documentCollectionValueHolder.v.getResourceId(), null, null)
.flatMap(collectionRoutingMapValueHolder -> {
return requestObs
.flatMap(req -> {
options.setPartitionKeyDefinition(documentCollectionValueHolder.v.getPartitionKey());
addPartitionLevelUnavailableRegionsForRequest(request, options, collectionRoutingMapValueHolder.v, retryPolicyInstance);
req.requestContext.setPointOperationContext(pointOperationContextForCircuitBreaker);
requestReference.set(req);
// needs to be after onBeforeSendRequest since CosmosDiagnostics instance needs to be wired
// to the RxDocumentServiceRequest instance
mergeContextInformationIntoDiagnosticsForPointRequest(request, pointOperationContextForCircuitBreaker);
return this.delete(req, retryPolicyInstance, getOperationContextAndListenerTuple(options));
})
.map(serviceResponse -> toResourceResponse(serviceResponse, Document.class));
}));
} catch (Exception e) {
logger.debug("Failure in deleting a document due to [{}]", e.getMessage());
return Mono.error(e);
}
}
@Override
public Mono> deleteAllDocumentsByPartitionKey(String collectionLink, PartitionKey partitionKey, RequestOptions options) {
// No ned-to-end policy / availability strategy applicable because PK Delete is a Gateway/Control-Plane operation
DocumentClientRetryPolicy requestRetryPolicy = this.resetSessionTokenRetryPolicy.getRequestPolicy(null);
return ObservableHelper.inlineIfPossibleAsObs(() -> deleteAllDocumentsByPartitionKeyInternal(collectionLink, options, requestRetryPolicy),
requestRetryPolicy);
}
private Mono> deleteAllDocumentsByPartitionKeyInternal(String collectionLink, RequestOptions options,
DocumentClientRetryPolicy retryPolicyInstance) {
try {
if (StringUtils.isEmpty(collectionLink)) {
throw new IllegalArgumentException("collectionLink");
}
logger.debug("Deleting all items by Partition Key. collectionLink: [{}]", collectionLink);
String path = Utils.joinPath(collectionLink, null);
Map requestHeaders = this.getRequestHeaders(options, ResourceType.PartitionKey, OperationType.Delete);
RxDocumentServiceRequest request = RxDocumentServiceRequest.create(this,
OperationType.Delete, ResourceType.PartitionKey, path, requestHeaders, options);
if (retryPolicyInstance != null) {
retryPolicyInstance.onBeforeSendRequest(request);
}
Mono> collectionObs = collectionCache.resolveCollectionAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), request);
Mono requestObs = addPartitionKeyInformation(request, null, null, options, collectionObs, null);
return requestObs.flatMap(req -> this
.deleteAllItemsByPartitionKey(req, retryPolicyInstance, getOperationContextAndListenerTuple(options))
.map(serviceResponse -> toResourceResponse(serviceResponse, Document.class)));
} catch (Exception e) {
logger.debug("Failure in deleting documents due to [{}]", e.getMessage());
return Mono.error(e);
}
}
@Override
public Mono> readDocument(String documentLink, RequestOptions options) {
return readDocument(documentLink, options, this);
}
private Mono> readDocument(
String documentLink,
RequestOptions options,
DiagnosticsClientContext innerDiagnosticsFactory) {
String collectionLink = Utils.getCollectionName(documentLink);
return wrapPointOperationWithAvailabilityStrategy(
ResourceType.Document,
OperationType.Read,
(opt, e2ecfg, clientCtxOverride, pointOperationContextForCircuitBreaker) -> readDocumentCore(documentLink, opt, e2ecfg, clientCtxOverride, pointOperationContextForCircuitBreaker),
options,
false,
innerDiagnosticsFactory,
collectionLink
);
}
private Mono> readDocumentCore(
String documentLink,
RequestOptions options,
CosmosEndToEndOperationLatencyPolicyConfig endToEndPolicyConfig,
DiagnosticsClientContext clientContextOverride,
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker) {
RequestOptions nonNullRequestOptions = options != null ? options : new RequestOptions();
ScopedDiagnosticsFactory scopedDiagnosticsFactory = new ScopedDiagnosticsFactory(clientContextOverride, false);
DocumentClientRetryPolicy retryPolicyInstance = this.resetSessionTokenRetryPolicy.getRequestPolicy(scopedDiagnosticsFactory);
AtomicReference requestReference = new AtomicReference<>();
return handleCircuitBreakingFeedbackForPointOperation(getPointOperationResponseMonoWithE2ETimeout(
nonNullRequestOptions,
endToEndPolicyConfig,
ObservableHelper.inlineIfPossibleAsObs(
() -> readDocumentInternal(
documentLink,
nonNullRequestOptions,
retryPolicyInstance,
scopedDiagnosticsFactory,
requestReference,
pointOperationContextForCircuitBreaker),
retryPolicyInstance),
scopedDiagnosticsFactory
), requestReference);
}
private Mono> readDocumentInternal(
String documentLink,
RequestOptions options,
DocumentClientRetryPolicy retryPolicyInstance,
DiagnosticsClientContext clientContextOverride,
AtomicReference requestReference,
PointOperationContextForCircuitBreaker pointOperationContextForCircuitBreaker) {
try {
if (StringUtils.isEmpty(documentLink)) {
throw new IllegalArgumentException("documentLink");
}
logger.debug("Reading a Document. documentLink: [{}]", documentLink);
String path = Utils.joinPath(documentLink, null);
Map requestHeaders = this.getRequestHeaders(options, ResourceType.Document, OperationType.Read);
RxDocumentServiceRequest request = RxDocumentServiceRequest.create(
getEffectiveClientContext(clientContextOverride),
OperationType.Read, ResourceType.Document, path, requestHeaders, options);
options.getMarkE2ETimeoutInRequestContextCallbackHook().set(
() -> request.requestContext.setIsRequestCancelledOnTimeout(new AtomicBoolean(true)));
request.requestContext.setExcludeRegions(options.getExcludedRegions());
request.requestContext.setKeywordIdentifiers(options.getKeywordIdentifiers());
if (retryPolicyInstance != null) {
retryPolicyInstance.onBeforeSendRequest(request);
}
Mono> collectionObs = this.collectionCache.resolveCollectionAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), request);
return collectionObs.flatMap(documentCollectionValueHolder -> {
if (documentCollectionValueHolder == null || documentCollectionValueHolder.v == null) {
return Mono.error(new IllegalStateException("documentCollectionValueHolder or documentCollectionValueHolder.v cannot be null"));
}
DocumentCollection documentCollection = documentCollectionValueHolder.v;
return this.partitionKeyRangeCache.tryLookupAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), documentCollection.getResourceId(), null, null)
.flatMap(collectionRoutingMapValueHolder -> {
if (collectionRoutingMapValueHolder == null || collectionRoutingMapValueHolder.v == null) {
return Mono.error(new IllegalStateException("collectionRoutingMapValueHolder or collectionRoutingMapValueHolder.v cannot be null"));
}
Mono requestObs = addPartitionKeyInformation(request, null, null, options, collectionObs, pointOperationContextForCircuitBreaker);
return requestObs.flatMap(req -> {
options.setPartitionKeyDefinition(documentCollection.getPartitionKey());
addPartitionLevelUnavailableRegionsForRequest(req, options, collectionRoutingMapValueHolder.v, retryPolicyInstance);
req.requestContext.setPointOperationContext(pointOperationContextForCircuitBreaker);
requestReference.set(req);
// needs to be after onBeforeSendRequest since CosmosDiagnostics instance needs to be wired
// to the RxDocumentServiceRequest instance
mergeContextInformationIntoDiagnosticsForPointRequest(request, pointOperationContextForCircuitBreaker);
return this.read(req, retryPolicyInstance)
.map(serviceResponse -> toResourceResponse(serviceResponse, Document.class));
});
});
}
);
} catch (Exception e) {
logger.debug("Failure in reading a document due to [{}]", e.getMessage());
return Mono.error(e);
}
}
@Override
public Flux> readDocuments(
String collectionLink, QueryFeedOperationState state, Class classOfT) {
if (StringUtils.isEmpty(collectionLink)) {
throw new IllegalArgumentException("collectionLink");
}
return queryDocuments(collectionLink, "SELECT * FROM r", state, classOfT);
}
@Override
public Mono> readMany(
List itemIdentityList,
String collectionLink,
QueryFeedOperationState state,
Class klass) {
final ScopedDiagnosticsFactory diagnosticsFactory = new ScopedDiagnosticsFactory(this, true);
state.registerDiagnosticsFactory(
() -> {}, // we never want to reset in readMany
(ctx) -> diagnosticsFactory.merge(ctx)
);
String resourceLink = parentResourceLinkToQueryLink(collectionLink, ResourceType.Document);
RxDocumentServiceRequest request = RxDocumentServiceRequest.create(diagnosticsFactory,
OperationType.Query,
ResourceType.Document,
collectionLink, null
);
// This should not get to backend
Mono> collectionObs =
collectionCache.resolveCollectionAsync(null, request);
return collectionObs
.flatMap(documentCollectionResourceResponse -> {
final DocumentCollection collection = documentCollectionResourceResponse.v;
if (collection == null) {
return Mono.error(new IllegalStateException("Collection cannot be null"));
}
final PartitionKeyDefinition pkDefinition = collection.getPartitionKey();
Mono> valueHolderMono = partitionKeyRangeCache
.tryLookupAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics),
collection.getResourceId(),
null,
null);
return valueHolderMono
.flatMap(collectionRoutingMapValueHolder -> {
Map> partitionRangeItemKeyMap = new HashMap<>();
CollectionRoutingMap routingMap = collectionRoutingMapValueHolder.v;
if (routingMap == null) {
return Mono.error(new IllegalStateException("Failed to get routing map."));
}
itemIdentityList
.forEach(itemIdentity -> {
//Check no partial partition keys are being used
if (pkDefinition.getKind().equals(PartitionKind.MULTI_HASH) &&
ModelBridgeInternal.getPartitionKeyInternal(itemIdentity.getPartitionKey())
.getComponents().size() != pkDefinition.getPaths().size()) {
throw new IllegalArgumentException(RMResources.PartitionKeyMismatch);
}
String effectivePartitionKeyString = PartitionKeyInternalHelper
.getEffectivePartitionKeyString(
BridgeInternal.getPartitionKeyInternal(
itemIdentity.getPartitionKey()),
pkDefinition);
//use routing map to find the partitionKeyRangeId of each
// effectivePartitionKey
PartitionKeyRange range =
routingMap.getRangeByEffectivePartitionKey(effectivePartitionKeyString);
//group the itemKeyList based on partitionKeyRangeId
if (partitionRangeItemKeyMap.get(range) == null) {
List list = new ArrayList<>();
list.add(itemIdentity);
partitionRangeItemKeyMap.put(range, list);
} else {
List pairs =
partitionRangeItemKeyMap.get(range);
pairs.add(itemIdentity);
partitionRangeItemKeyMap.put(range, pairs);
}
});
//Create the range query map that contains the query to be run for that
// partitionkeyrange
Map rangeQueryMap = getRangeQueryMap(partitionRangeItemKeyMap, collection.getPartitionKey());
// create point reads
Flux> pointReads = pointReadsForReadMany(
diagnosticsFactory,
partitionRangeItemKeyMap,
resourceLink,
state.getQueryOptions(),
klass);
// create the executable query
Flux> queries = queryForReadMany(
diagnosticsFactory,
resourceLink,
new SqlQuerySpec(DUMMY_SQL_QUERY),
state.getQueryOptions(),
klass,
ResourceType.Document,
collection,
Collections.unmodifiableMap(rangeQueryMap));
// merge results from point reads and queries
return Flux.merge(pointReads, queries)
.collectList()
// aggregating the result to construct a FeedResponse and aggregate RUs.
.map(feedList -> {
List finalList = new ArrayList<>();
HashMap headers = new HashMap<>();
ConcurrentMap aggregatedQueryMetrics = new ConcurrentHashMap<>();
Collection aggregateRequestStatistics = new DistinctClientSideRequestStatisticsCollection();
double requestCharge = 0;
for (FeedResponse page : feedList) {
ConcurrentMap pageQueryMetrics =
ModelBridgeInternal.queryMetrics(page);
if (pageQueryMetrics != null) {
pageQueryMetrics.forEach(
aggregatedQueryMetrics::putIfAbsent);
}
requestCharge += page.getRequestCharge();
finalList.addAll(page.getResults());
aggregateRequestStatistics.addAll(diagnosticsAccessor.getClientSideRequestStatistics(page.getCosmosDiagnostics()));
}
// NOTE: This CosmosDiagnostics instance intentionally isn't captured in the
// ScopedDiagnosticsFactory - and a ssuch won't be included in the diagnostics of the
// CosmosDiagnosticsContext - which is fine, because the CosmosDiagnosticsContext
// contains the "real" CosmosDiagnostics instances (which will also be used
// for diagnostics purposes - like metrics, logging etc.
// this artificial CosmosDiagnostics with the aggregated RU/s etc. si simply
// to maintain the API contract that a FeedResponse returns one CosmosDiagnostics
CosmosDiagnostics aggregatedDiagnostics = BridgeInternal.createCosmosDiagnostics(aggregatedQueryMetrics);
diagnosticsAccessor.addClientSideDiagnosticsToFeed(
aggregatedDiagnostics, aggregateRequestStatistics);
state.mergeDiagnosticsContext();
CosmosDiagnosticsContext ctx = state.getDiagnosticsContextSnapshot();
if (ctx != null) {
ctxAccessor.recordOperation(
ctx,
200,
0,
finalList.size(),
requestCharge,
aggregatedDiagnostics,
null
);
diagnosticsAccessor
.setDiagnosticsContext(
aggregatedDiagnostics,
ctx);
}
headers.put(HttpConstants.HttpHeaders.REQUEST_CHARGE, Double
.toString(requestCharge));
FeedResponse frp = BridgeInternal
.createFeedResponseWithQueryMetrics(
finalList,
headers,
aggregatedQueryMetrics,
null,
false,
false,
aggregatedDiagnostics);
return frp;
});
})
.onErrorMap(throwable -> {
if (throwable instanceof CosmosException) {
CosmosException cosmosException = (CosmosException)throwable;
CosmosDiagnostics diagnostics = cosmosException.getDiagnostics();
if (diagnostics != null) {
state.mergeDiagnosticsContext();
CosmosDiagnosticsContext ctx = state.getDiagnosticsContextSnapshot();
if (ctx != null) {
ctxAccessor.recordOperation(
ctx,
cosmosException.getStatusCode(),
cosmosException.getSubStatusCode(),
0,
cosmosException.getRequestCharge(),
diagnostics,
throwable
);
diagnosticsAccessor
.setDiagnosticsContext(
diagnostics,
state.getDiagnosticsContextSnapshot());
}
}
return cosmosException;
}
return throwable;
});
}
);
}
private Map getRangeQueryMap(
Map> partitionRangeItemKeyMap,
PartitionKeyDefinition partitionKeyDefinition) {
//TODO: Optimise this to include all types of partitionkeydefinitions. ex: c["prop1./ab"]["key1"]
Map rangeQueryMap = new HashMap<>();
List partitionKeySelectors = createPkSelectors(partitionKeyDefinition);
for(Map.Entry> entry: partitionRangeItemKeyMap.entrySet()) {
SqlQuerySpec sqlQuerySpec;
List cosmosItemIdentityList = entry.getValue();
if (cosmosItemIdentityList.size() > 1) {
if (partitionKeySelectors.size() == 1 && partitionKeySelectors.get(0).equals("[\"id\"]")) {
sqlQuerySpec = createReadManyQuerySpecPartitionKeyIdSame(cosmosItemIdentityList);
} else {
sqlQuerySpec = createReadManyQuerySpec(entry.getValue(), partitionKeySelectors);
}
// Add query for this partition to rangeQueryMap
rangeQueryMap.put(entry.getKey(), sqlQuerySpec);
}
}
return rangeQueryMap;
}
private SqlQuerySpec createReadManyQuerySpecPartitionKeyIdSame(List idPartitionKeyPairList) {
StringBuilder queryStringBuilder = new StringBuilder();
List parameters = new ArrayList<>();
queryStringBuilder.append("SELECT * FROM c WHERE c.id IN ( ");
for (int i = 0; i < idPartitionKeyPairList.size(); i++) {
CosmosItemIdentity itemIdentity = idPartitionKeyPairList.get(i);
String idValue = itemIdentity.getId();
String idParamName = "@param" + i;
parameters.add(new SqlParameter(idParamName, idValue));
queryStringBuilder.append(idParamName);
if (i < idPartitionKeyPairList.size() - 1) {
queryStringBuilder.append(", ");
}
}
queryStringBuilder.append(" )");
return new SqlQuerySpec(queryStringBuilder.toString(), parameters);
}
private SqlQuerySpec createReadManyQuerySpec(
List itemIdentities,
List partitionKeySelectors) {
StringBuilder queryStringBuilder = new StringBuilder();
List parameters = new ArrayList<>();
queryStringBuilder.append("SELECT * FROM c WHERE ( ");
int paramCount = 0;
for (int i = 0; i < itemIdentities.size(); i++) {
CosmosItemIdentity itemIdentity = itemIdentities.get(i);
PartitionKey pkValueAsPartitionKey = itemIdentity.getPartitionKey();
Object[] pkValues = ModelBridgeInternal.getPartitionKeyInternal(pkValueAsPartitionKey).toObjectArray();
List> partitionKeyParams = new ArrayList<>();
int pathCount = 0;
for (Object pkComponentValue : pkValues) {
String pkParamName = "@param" + paramCount;
partitionKeyParams.add(Arrays.asList(partitionKeySelectors.get(pathCount), pkParamName));
parameters.add(new SqlParameter(pkParamName, pkComponentValue));
paramCount++;
pathCount++;
}
String idValue = itemIdentity.getId();
String idParamName = "@param" + paramCount;
paramCount++;
parameters.add(new SqlParameter(idParamName, idValue));
queryStringBuilder.append("(");
queryStringBuilder.append("c.id = ");
queryStringBuilder.append(idParamName);
// partition key def
for (List pkParam: partitionKeyParams) {
queryStringBuilder.append(" AND ");
queryStringBuilder.append(" c");
queryStringBuilder.append(pkParam.get(0));
queryStringBuilder.append((" = "));
queryStringBuilder.append(pkParam.get(1));
}
queryStringBuilder.append(" )");
if (i < itemIdentities.size() - 1) {
queryStringBuilder.append(" OR ");
}
}
queryStringBuilder.append(" )");
return new SqlQuerySpec(queryStringBuilder.toString(), parameters);
}
private List createPkSelectors(PartitionKeyDefinition partitionKeyDefinition) {
return partitionKeyDefinition.getPaths()
.stream()
.map(pathPart -> StringUtils.substring(pathPart, 1)) // skip starting /
.map(pathPart -> StringUtils.replace(pathPart, "\"", "\\")) // escape quote
.map(part -> "[\"" + part + "\"]")
.collect(Collectors.toList());
}
private Flux> queryForReadMany(
ScopedDiagnosticsFactory diagnosticsFactory,
String parentResourceLink,
SqlQuerySpec sqlQuery,
CosmosQueryRequestOptions options,
Class klass,
ResourceType resourceTypeEnum,
DocumentCollection collection,
Map rangeQueryMap) {
if (rangeQueryMap.isEmpty()) {
return Flux.empty();
}
UUID activityId = randomUuid();
final AtomicBoolean isQueryCancelledOnTimeout = new AtomicBoolean(false);
IDocumentQueryClient queryClient = documentQueryClientImpl(RxDocumentClientImpl.this, getOperationContextAndListenerTuple(options));
Flux> executionContext =
DocumentQueryExecutionContextFactory.createReadManyQueryAsync(
diagnosticsFactory,
queryClient,
collection.getResourceId(),
sqlQuery,
rangeQueryMap,
options,
collection,
parentResourceLink,
activityId,
klass,
resourceTypeEnum,
isQueryCancelledOnTimeout);
Flux> feedResponseFlux = executionContext.flatMap(IDocumentQueryExecutionContext::executeAsync);
RequestOptions requestOptions = options == null? null : ImplementationBridgeHelpers
.CosmosQueryRequestOptionsHelper
.getCosmosQueryRequestOptionsAccessor()
.toRequestOptions(options);
CosmosEndToEndOperationLatencyPolicyConfig endToEndPolicyConfig =
getEndToEndOperationLatencyPolicyConfig(requestOptions, ResourceType.Document, OperationType.Query);
if (endToEndPolicyConfig != null && endToEndPolicyConfig.isEnabled()) {
return getFeedResponseFluxWithTimeout(
feedResponseFlux,
endToEndPolicyConfig,
options,
isQueryCancelledOnTimeout,
diagnosticsFactory);
}
return feedResponseFlux;
}
private Flux> pointReadsForReadMany(
ScopedDiagnosticsFactory diagnosticsFactory,
Map> singleItemPartitionRequestMap,
String resourceLink,
CosmosQueryRequestOptions queryRequestOptions,
Class klass) {
// if there is any factory method being passed in, use the factory method to deserializ the object
// else fallback to use the original way
// typically used by spark trying to convert into SparkRowItem
CosmosItemSerializer effectiveItemSerializer = getEffectiveItemSerializer(queryRequestOptions);
return Flux.fromIterable(singleItemPartitionRequestMap.values())
.flatMap(cosmosItemIdentityList -> {
if (cosmosItemIdentityList.size() == 1) {
CosmosItemIdentity firstIdentity = cosmosItemIdentityList.get(0);
RequestOptions requestOptions = ImplementationBridgeHelpers
.CosmosQueryRequestOptionsHelper
.getCosmosQueryRequestOptionsAccessor()
.toRequestOptions(queryRequestOptions);
requestOptions.setPartitionKey(firstIdentity.getPartitionKey());
return this.readDocument((resourceLink + firstIdentity.getId()), requestOptions, diagnosticsFactory)
.flatMap(resourceResponse -> Mono.just(
new ImmutablePair, CosmosException>(resourceResponse, null)
))
.onErrorResume(throwable -> {
Throwable unwrappedThrowable = Exceptions.unwrap(throwable);
if (unwrappedThrowable instanceof CosmosException) {
CosmosException cosmosException = (CosmosException) unwrappedThrowable;
int statusCode = cosmosException.getStatusCode();
int subStatusCode = cosmosException.getSubStatusCode();
if (statusCode == HttpConstants.StatusCodes.NOTFOUND && subStatusCode == HttpConstants.SubStatusCodes.UNKNOWN) {
return Mono.just(new ImmutablePair, CosmosException>(null, cosmosException));
}
}
return Mono.error(unwrappedThrowable);
});
}
return Mono.empty();
})
.flatMap(resourceResponseToExceptionPair -> {
ResourceResponse resourceResponse = resourceResponseToExceptionPair.getLeft();
CosmosException cosmosException = resourceResponseToExceptionPair.getRight();
FeedResponse feedResponse;
if (cosmosException != null) {
feedResponse = ModelBridgeInternal.createFeedResponse(new ArrayList<>(), cosmosException.getResponseHeaders());
diagnosticsAccessor.addClientSideDiagnosticsToFeed(
feedResponse.getCosmosDiagnostics(),
Collections.singleton(
BridgeInternal.getClientSideRequestStatics(cosmosException.getDiagnostics())));
} else {
CosmosItemResponse cosmosItemResponse =
itemResponseAccessor.createCosmosItemResponse(resourceResponse, klass, effectiveItemSerializer);
feedResponse = ModelBridgeInternal.createFeedResponse(
Arrays.asList(cosmosItemResponse.getItem()),
cosmosItemResponse.getResponseHeaders());
diagnosticsAccessor.addClientSideDiagnosticsToFeed(
feedResponse.getCosmosDiagnostics(),
Collections.singleton(
BridgeInternal.getClientSideRequestStatics(cosmosItemResponse.getDiagnostics())));
}
return Mono.just(feedResponse);
});
}
@Override
public Flux> queryDocuments(
String collectionLink, String query, QueryFeedOperationState state, Class classOfT) {
return queryDocuments(collectionLink, new SqlQuerySpec(query), state, classOfT);
}
@Override
public CosmosItemSerializer getEffectiveItemSerializer(CosmosItemSerializer requestOptionsItemSerializer) {
if (requestOptionsItemSerializer != null) {
return requestOptionsItemSerializer;
}
if (this.defaultCustomSerializer != null) {
return this.defaultCustomSerializer;
}
return CosmosItemSerializer.DEFAULT_SERIALIZER;
}
private CosmosItemSerializer getEffectiveItemSerializer(CosmosQueryRequestOptions queryRequestOptions) {
CosmosItemSerializer requestOptionsItemSerializer =
queryRequestOptions != null ? queryRequestOptions.getCustomItemSerializer() : null;
return this.getEffectiveItemSerializer(requestOptionsItemSerializer);
}
private CosmosItemSerializer getEffectiveItemSerializer(CosmosItemRequestOptions itemRequestOptions) {
CosmosItemSerializer requestOptionsItemSerializer =
itemRequestOptions != null ? itemRequestOptions.getCustomItemSerializer() : null;
return this.getEffectiveItemSerializer(requestOptionsItemSerializer);
}
private IDocumentQueryClient documentQueryClientImpl(RxDocumentClientImpl rxDocumentClientImpl, OperationContextAndListenerTuple operationContextAndListenerTuple) {
return new IDocumentQueryClient () {
@Override
public RxCollectionCache getCollectionCache() {
return RxDocumentClientImpl.this.collectionCache;
}
@Override
public RxPartitionKeyRangeCache getPartitionKeyRangeCache() {
return RxDocumentClientImpl.this.partitionKeyRangeCache;
}
@Override
public IRetryPolicyFactory getResetSessionTokenRetryPolicy() {
return RxDocumentClientImpl.this.resetSessionTokenRetryPolicy;
}
@Override
public ConsistencyLevel getDefaultConsistencyLevelAsync() {
return RxDocumentClientImpl.this.gatewayConfigurationReader.getDefaultConsistencyLevel();
}
@Override
public ConsistencyLevel getDesiredConsistencyLevelAsync() {
// TODO Auto-generated method stub
return RxDocumentClientImpl.this.consistencyLevel;
}
@Override
public Mono executeQueryAsync(RxDocumentServiceRequest request) {
if (operationContextAndListenerTuple == null) {
return RxDocumentClientImpl.this.query(request).single();
} else {
final OperationListener listener =
operationContextAndListenerTuple.getOperationListener();
final OperationContext operationContext = operationContextAndListenerTuple.getOperationContext();
request.getHeaders().put(HttpConstants.HttpHeaders.CORRELATED_ACTIVITY_ID, operationContext.getCorrelationActivityId());
listener.requestListener(operationContext, request);
return RxDocumentClientImpl.this.query(request).single().doOnNext(
response -> listener.responseListener(operationContext, response)
).doOnError(
ex -> listener.exceptionListener(operationContext, ex)
);
}
}
@Override
public QueryCompatibilityMode getQueryCompatibilityMode() {
// TODO Auto-generated method stub
return QueryCompatibilityMode.Default;
}
@Override
public Mono executeFeedOperationWithAvailabilityStrategy(
ResourceType resourceType,
OperationType operationType,
Supplier retryPolicyFactory,
RxDocumentServiceRequest req,
BiFunction, RxDocumentServiceRequest, Mono> feedOperation,
String collectionLink) {
return RxDocumentClientImpl.this.executeFeedOperationWithAvailabilityStrategy(
resourceType,
operationType,
retryPolicyFactory,
req,
feedOperation,
collectionLink);
}
@Override
public CosmosItemSerializer getEffectiveItemSerializer(CosmosQueryRequestOptions queryRequestOptions) {
return RxDocumentClientImpl.this.getEffectiveItemSerializer(queryRequestOptions);
}
@Override
public Mono readFeedAsync(RxDocumentServiceRequest request) {
// TODO Auto-generated method stub
return null;
}
@Override
public Mono populateFeedRangeHeader(RxDocumentServiceRequest request) {
if (RxDocumentClientImpl.this.requiresFeedRangeFiltering(request)) {
return request
.getFeedRange()
.populateFeedRangeFilteringHeaders(RxDocumentClientImpl.this.partitionKeyRangeCache, request, RxDocumentClientImpl.this.collectionCache.resolveCollectionAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), request))
.flatMap(ignore -> Mono.just(request));
} else {
return Mono.just(request);
}
}
@Override
public Mono addPartitionLevelUnavailableRegionsOnRequest(RxDocumentServiceRequest request, CosmosQueryRequestOptions queryRequestOptions, DocumentClientRetryPolicy documentClientRetryPolicy) {
if (RxDocumentClientImpl.this.globalPartitionEndpointManagerForCircuitBreaker.isPartitionLevelCircuitBreakingApplicable(request)) {
String collectionRid = RxDocumentClientImpl.qryOptAccessor.getCollectionRid(queryRequestOptions);
checkNotNull(collectionRid, "Argument 'collectionRid' cannot be null!");
return RxDocumentClientImpl.this.partitionKeyRangeCache.tryLookupAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), collectionRid, null, null)
.flatMap(collectionRoutingMapValueHolder -> {
if (collectionRoutingMapValueHolder.v == null) {
return Mono.error(new CollectionRoutingMapNotFoundException("Argument 'collectionRoutingMapValueHolder.v' cannot be null!"));
}
RxDocumentClientImpl.this.addPartitionLevelUnavailableRegionsForFeedRequest(request, queryRequestOptions, collectionRoutingMapValueHolder.v);
// onBeforeSendRequest uses excluded regions to know the next location endpoint
// to route the request to unavailable regions are effectively excluded regions for this request
if (documentClientRetryPolicy != null) {
documentClientRetryPolicy.onBeforeSendRequest(request);
}
return Mono.just(request);
});
} else {
return Mono.just(request);
}
}
@Override
public GlobalEndpointManager getGlobalEndpointManager() {
return RxDocumentClientImpl.this.getGlobalEndpointManager();
}
@Override
public GlobalPartitionEndpointManagerForCircuitBreaker getGlobalPartitionEndpointManagerForCircuitBreaker() {
return RxDocumentClientImpl.this.globalPartitionEndpointManagerForCircuitBreaker;
}
};
}
@Override
public Flux> queryDocuments(
String collectionLink,
SqlQuerySpec querySpec,
QueryFeedOperationState state,
Class classOfT) {
SqlQuerySpecLogger.getInstance().logQuery(querySpec);
return createQuery(collectionLink, querySpec, state, classOfT, ResourceType.Document);
}
@Override
public Flux> queryDocumentChangeFeed(
final DocumentCollection collection,
final CosmosChangeFeedRequestOptions requestOptions,
Class classOfT) {
checkNotNull(collection, "Argument 'collection' must not be null.");
ChangeFeedQueryImpl changeFeedQueryImpl = new ChangeFeedQueryImpl<>(
this,
ResourceType.Document,
classOfT,
collection.getAltLink(),
collection.getResourceId(),
requestOptions);
return changeFeedQueryImpl.executeAsync();
}
@Override
public Flux> queryDocumentChangeFeedFromPagedFlux(DocumentCollection collection, ChangeFeedOperationState state, Class classOfT) {
CosmosChangeFeedRequestOptions clonedOptions = changeFeedOptionsAccessor.clone(state.getChangeFeedOptions());
CosmosChangeFeedRequestOptionsImpl optionsImpl = changeFeedOptionsAccessor.getImpl(clonedOptions);
CosmosOperationDetails operationDetails = operationDetailsAccessor.create(optionsImpl, state.getDiagnosticsContextSnapshot());
this.operationPolicies.forEach(policy -> {
try {
policy.process(operationDetails);
} catch (RuntimeException exception) {
logger.info("The following exception was thrown by a custom policy on changeFeed operation" + exception.getMessage());
throw(exception);
}
});
ctxAccessor.setRequestOptions(state.getDiagnosticsContextSnapshot(), optionsImpl);
return queryDocumentChangeFeed(collection, clonedOptions, classOfT);
}
@Override
public Flux> readAllDocuments(
String collectionLink,
PartitionKey partitionKey,
QueryFeedOperationState state,
Class classOfT) {
if (StringUtils.isEmpty(collectionLink)) {
throw new IllegalArgumentException("collectionLink");
}
if (partitionKey == null) {
throw new IllegalArgumentException("partitionKey");
}
final CosmosQueryRequestOptions effectiveOptions =
qryOptAccessor.clone(state.getQueryOptions());
RequestOptions nonNullRequestOptions = qryOptAccessor.toRequestOptions(effectiveOptions);
CosmosEndToEndOperationLatencyPolicyConfig endToEndPolicyConfig =
nonNullRequestOptions.getCosmosEndToEndLatencyPolicyConfig();
List orderedApplicableRegionsForSpeculation = getApplicableRegionsForSpeculation(
endToEndPolicyConfig,
ResourceType.Document,
OperationType.Query,
false,
nonNullRequestOptions);
ScopedDiagnosticsFactory diagnosticsFactory = new ScopedDiagnosticsFactory(this, false);
if (orderedApplicableRegionsForSpeculation.size() < 2) {
state.registerDiagnosticsFactory(
() -> {},
(ctx) -> diagnosticsFactory.merge(ctx));
} else {
state.registerDiagnosticsFactory(
() -> diagnosticsFactory.reset(),
(ctx) -> diagnosticsFactory.merge(ctx));
}
RxDocumentServiceRequest request = RxDocumentServiceRequest.create(
diagnosticsFactory,
OperationType.Query,
ResourceType.Document,
collectionLink,
null
);
// This should not get to backend
Flux> collectionObs =
collectionCache.resolveCollectionAsync(null, request).flux();
return collectionObs.flatMap(documentCollectionResourceResponse -> {
DocumentCollection collection = documentCollectionResourceResponse.v;
if (collection == null) {
return Mono.error(new IllegalStateException("Collection cannot be null"));
}
PartitionKeyDefinition pkDefinition = collection.getPartitionKey();
List partitionKeySelectors = createPkSelectors(pkDefinition);
SqlQuerySpec querySpec = createLogicalPartitionScanQuerySpec(partitionKey, partitionKeySelectors);
String resourceLink = parentResourceLinkToQueryLink(collectionLink, ResourceType.Document);
UUID activityId = randomUuid();
final AtomicBoolean isQueryCancelledOnTimeout = new AtomicBoolean(false);
IDocumentQueryClient queryClient = documentQueryClientImpl(RxDocumentClientImpl.this, getOperationContextAndListenerTuple(state.getQueryOptions()));
// Trying to put this logic as low as the query pipeline
// Since for parallelQuery, each partition will have its own request, so at this point, there will be no request associate with this retry policy.
// For default document context, it already wired up InvalidPartitionExceptionRetry, but there is no harm to wire it again here
InvalidPartitionExceptionRetryPolicy invalidPartitionExceptionRetryPolicy = new InvalidPartitionExceptionRetryPolicy(
this.collectionCache,
null,
resourceLink,
ModelBridgeInternal.getPropertiesFromQueryRequestOptions(effectiveOptions));
Flux> innerFlux = ObservableHelper.fluxInlineIfPossibleAsObs(
() -> {
Flux> valueHolderMono = this.partitionKeyRangeCache
.tryLookupAsync(
BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics),
collection.getResourceId(),
null,
null).flux();
return valueHolderMono.flatMap(collectionRoutingMapValueHolder -> {
CollectionRoutingMap routingMap = collectionRoutingMapValueHolder.v;
if (routingMap == null) {
return Mono.error(new IllegalStateException("Failed to get routing map."));
}
String effectivePartitionKeyString = PartitionKeyInternalHelper
.getEffectivePartitionKeyString(
BridgeInternal.getPartitionKeyInternal(partitionKey),
pkDefinition);
//use routing map to find the partitionKeyRangeId of each
// effectivePartitionKey
PartitionKeyRange range =
routingMap.getRangeByEffectivePartitionKey(effectivePartitionKeyString);
return createQueryInternal(
diagnosticsFactory,
resourceLink,
querySpec,
ModelBridgeInternal.setPartitionKeyRangeIdInternal(effectiveOptions, range.getId()),
classOfT, //Document.class
ResourceType.Document,
queryClient,
activityId,
isQueryCancelledOnTimeout);
});
},
invalidPartitionExceptionRetryPolicy);
if (orderedApplicableRegionsForSpeculation.size() < 2) {
return innerFlux;
}
return innerFlux
.flatMap(result -> {
diagnosticsFactory.merge(nonNullRequestOptions);
return Mono.just(result);
})
.onErrorMap(throwable -> {
diagnosticsFactory.merge(nonNullRequestOptions);
return throwable;
})
.doOnCancel(() -> diagnosticsFactory.merge(nonNullRequestOptions));
});
}
@Override
public Map getQueryPlanCache() {
return queryPlanCache;
}
@Override
public Flux> readPartitionKeyRanges(final String collectionLink,
QueryFeedOperationState state) {
if (StringUtils.isEmpty(collectionLink)) {
throw new IllegalArgumentException("collectionLink");
}
return nonDocumentReadFeed(state, ResourceType.PartitionKeyRange, PartitionKeyRange.class,
Utils.joinPath(collectionLink, Paths.PARTITION_KEY_RANGES_PATH_SEGMENT));
}
@Override
public Flux> readPartitionKeyRanges(String collectionLink, CosmosQueryRequestOptions options) {
if (StringUtils.isEmpty(collectionLink)) {
throw new IllegalArgumentException("collectionLink");
}
return nonDocumentReadFeed(options, ResourceType.PartitionKeyRange, PartitionKeyRange.class,
Utils.joinPath(collectionLink, Paths.PARTITION_KEY_RANGES_PATH_SEGMENT));
}
private RxDocumentServiceRequest getStoredProcedureRequest(String collectionLink, StoredProcedure storedProcedure,
RequestOptions options, OperationType operationType) {
if (StringUtils.isEmpty(collectionLink)) {
throw new IllegalArgumentException("collectionLink");
}
if (storedProcedure == null) {
throw new IllegalArgumentException("storedProcedure");
}
validateResource(storedProcedure);
String path = Utils.joinPath(collectionLink, Paths.STORED_PROCEDURES_PATH_SEGMENT);
Map requestHeaders = this.getRequestHeaders(options, ResourceType.StoredProcedure, operationType);
return RxDocumentServiceRequest.create(this, operationType,
ResourceType.StoredProcedure, path, storedProcedure, requestHeaders, options);
}
private RxDocumentServiceRequest getUserDefinedFunctionRequest(String collectionLink, UserDefinedFunction udf,
RequestOptions options, OperationType operationType) {
if (StringUtils.isEmpty(collectionLink)) {
throw new IllegalArgumentException("collectionLink");
}
if (udf == null) {
throw new IllegalArgumentException("udf");
}
validateResource(udf);
String path = Utils.joinPath(collectionLink, Paths.USER_DEFINED_FUNCTIONS_PATH_SEGMENT);
Map requestHeaders = this.getRequestHeaders(options, ResourceType.UserDefinedFunction, operationType);
return RxDocumentServiceRequest.create(this,
operationType, ResourceType.UserDefinedFunction, path, udf, requestHeaders, options);
}
@Override
public Mono 