Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 1 addition & 12 deletions sdk/cosmos/azure-cosmos/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -128,18 +128,7 @@ Licensed under the MIT License.
</exclusion>
</exclusions>
</dependency>
<dependency>
<groupId>com.azure</groupId>
<artifactId>azure-core-tracing-opentelemetry</artifactId>
<scope>test</scope>
<version>1.0.0-beta.5</version> <!-- {x-version-update;com.azure:azure-core-tracing-opentelemetry;dependency} -->
<exclusions>
<exclusion>
<groupId>io.netty</groupId>
<artifactId>netty-tcnative-boringssl-static</artifactId>
</exclusion>
</exclusions>
</dependency>

<dependency>
<groupId>com.fasterxml.jackson.module</groupId>
<artifactId>jackson-module-afterburner</artifactId>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,19 +4,16 @@

import com.azure.core.annotation.ServiceClient;
import com.azure.core.credential.AzureKeyCredential;
import com.azure.core.util.Context;
import com.azure.core.util.tracing.Tracer;
import com.azure.cosmos.implementation.AsyncDocumentClient;
import com.azure.cosmos.implementation.Configs;
import com.azure.cosmos.implementation.ConnectionPolicy;
import com.azure.cosmos.implementation.CosmosAuthorizationTokenResolver;
import com.azure.cosmos.implementation.Database;
import com.azure.cosmos.implementation.HttpConstants;
import com.azure.cosmos.implementation.TracerProvider;
import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdMetrics;
import com.azure.cosmos.models.CosmosDatabaseResponse;
import com.azure.cosmos.models.CosmosDatabaseProperties;
import com.azure.cosmos.models.CosmosDatabaseRequestOptions;
import com.azure.cosmos.models.CosmosDatabaseResponse;
import com.azure.cosmos.models.CosmosPermissionProperties;
import com.azure.cosmos.models.CosmosQueryRequestOptions;
import com.azure.cosmos.models.ModelBridgeInternal;
Expand All @@ -30,9 +27,7 @@

import java.io.Closeable;
import java.util.List;
import java.util.ServiceLoader;

import static com.azure.core.util.FluxUtil.withContext;
import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount;

/**
Expand All @@ -56,7 +51,6 @@ public final class CosmosAsyncClient implements Closeable {
private final AzureKeyCredential credential;
private final boolean sessionCapturingOverride;
private final boolean enableTransportClientSharing;
private final TracerProvider tracerProvider;
private final boolean contentResponseOnWriteEnabled;

CosmosAsyncClient(CosmosClientBuilder builder) {
Expand All @@ -71,7 +65,6 @@ public final class CosmosAsyncClient implements Closeable {
this.sessionCapturingOverride = builder.isSessionCapturingOverrideEnabled();
this.enableTransportClientSharing = builder.isConnectionSharingAcrossClientsEnabled();
this.contentResponseOnWriteEnabled = builder.isContentResponseOnWriteEnabled();
this.tracerProvider = new TracerProvider(ServiceLoader.load(Tracer.class));
this.asyncDocumentClient = new AsyncDocumentClient.Builder()
.withServiceEndpoint(this.serviceEndpoint)
.withMasterKeyOrResourceToken(this.keyOrResourceToken)
Expand Down Expand Up @@ -201,14 +194,8 @@ boolean isContentResponseOnWriteEnabled() {
* @return a {@link Mono} containing the cosmos database response with the created or existing database or
* an error.
*/
public Mono<CosmosDatabaseResponse> createDatabaseIfNotExists(CosmosDatabaseProperties databaseProperties) {
if(!getTracerProvider().isEnabled()) {
CosmosAsyncDatabase database = getDatabase(databaseProperties.getId());
return createDatabaseIfNotExistsInternal(database.read(), database, null, null);
}

return withContext(context -> createDatabaseIfNotExistsInternal(getDatabase(databaseProperties.getId()),
null, context));
Mono<CosmosDatabaseResponse> createDatabaseIfNotExists(CosmosDatabaseProperties databaseProperties) {
return createDatabaseIfNotExistsInternal(getDatabase(databaseProperties.getId()));
}

/**
Expand All @@ -222,12 +209,21 @@ public Mono<CosmosDatabaseResponse> createDatabaseIfNotExists(CosmosDatabaseProp
* an error.
*/
public Mono<CosmosDatabaseResponse> createDatabaseIfNotExists(String id) {
if(!getTracerProvider().isEnabled()) {
CosmosAsyncDatabase database = getDatabase(id);
return createDatabaseIfNotExistsInternal(database.read(), database, null,null);
}
return createDatabaseIfNotExistsInternal(getDatabase(id));
}

return withContext(context -> createDatabaseIfNotExistsInternal(getDatabase(id), null, context));
private Mono<CosmosDatabaseResponse> createDatabaseIfNotExistsInternal(CosmosAsyncDatabase database) {
return database.read().onErrorResume(exception -> {
final Throwable unwrappedException = Exceptions.unwrap(exception);
if (unwrappedException instanceof CosmosException) {
final CosmosException cosmosException = (CosmosException) unwrappedException;
if (cosmosException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) {
return createDatabase(new CosmosDatabaseProperties(database.getId()),
new CosmosDatabaseRequestOptions());
}
}
return Mono.error(unwrappedException);
});
}

/**
Expand All @@ -244,13 +240,19 @@ public Mono<CosmosDatabaseResponse> createDatabaseIfNotExists(String id) {
* @return the mono.
*/
public Mono<CosmosDatabaseResponse> createDatabaseIfNotExists(String id, ThroughputProperties throughputProperties) {
if(!getTracerProvider().isEnabled()) {
CosmosAsyncDatabase database = getDatabase(id);
return createDatabaseIfNotExistsInternal(database.read(), database, throughputProperties, null);
}

return withContext(context -> createDatabaseIfNotExistsInternal(getDatabase(id),
throughputProperties, context));
return this.getDatabase(id).read().onErrorResume(exception -> {
final Throwable unwrappedException = Exceptions.unwrap(exception);
if (unwrappedException instanceof CosmosException) {
final CosmosException cosmosException = (CosmosException) unwrappedException;
if (cosmosException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) {
CosmosDatabaseRequestOptions options = new CosmosDatabaseRequestOptions();
ModelBridgeInternal.setThroughputProperties(options, throughputProperties);
return createDatabase(new CosmosDatabaseProperties(id),
options);
}
}
return Mono.error(unwrappedException);
});
}

/**
Expand All @@ -272,12 +274,9 @@ public Mono<CosmosDatabaseResponse> createDatabase(CosmosDatabaseProperties data
}
Database wrappedDatabase = new Database();
wrappedDatabase.setId(databaseProperties.getId());
if(!getTracerProvider().isEnabled()) {
return createDatabaseInternal(wrappedDatabase, options);
}

final CosmosDatabaseRequestOptions requestOptions = options;
return withContext(context -> createDatabaseInternal(wrappedDatabase, requestOptions, context));
return asyncDocumentClient.createDatabase(wrappedDatabase, ModelBridgeInternal.toRequestOptions(options))
.map(databaseResourceResponse -> ModelBridgeInternal.createCosmosDatabaseResponse(databaseResourceResponse))
.single();
}

/**
Expand Down Expand Up @@ -332,13 +331,9 @@ public Mono<CosmosDatabaseResponse> createDatabase(CosmosDatabaseProperties data
ModelBridgeInternal.setThroughputProperties(options, throughputProperties);
Database wrappedDatabase = new Database();
wrappedDatabase.setId(databaseProperties.getId());
if (!getTracerProvider().isEnabled()) {
return createDatabaseInternal(wrappedDatabase, options);
}


final CosmosDatabaseRequestOptions requestOptions = options;
return withContext(context -> createDatabaseInternal(wrappedDatabase, requestOptions, context));
return asyncDocumentClient.createDatabase(wrappedDatabase, ModelBridgeInternal.toRequestOptions(options))
.map(databaseResourceResponse -> ModelBridgeInternal.createCosmosDatabaseResponse(databaseResourceResponse))
.single();
}

/**
Expand Down Expand Up @@ -402,15 +397,13 @@ public Mono<CosmosDatabaseResponse> createDatabase(String id, ThroughputProperti
*/
CosmosPagedFlux<CosmosDatabaseProperties> readAllDatabases(CosmosQueryRequestOptions options) {
return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> {
String spanName = "readAllDatabases";
pagedFluxOptions.setTracerInformation(this.tracerProvider, spanName, this.serviceEndpoint, null);
setContinuationTokenAndMaxItemCount(pagedFluxOptions, options);
return getDocClientWrapper().readDatabases(options)
.map(response ->
BridgeInternal.createFeedResponse(
ModelBridgeInternal.getCosmosDatabasePropertiesFromV2Results(response.getResults()),
response.getResponseHeaders()));
}, this.tracerProvider.isEnabled());
});
}

/**
Expand Down Expand Up @@ -439,7 +432,7 @@ public CosmosPagedFlux<CosmosDatabaseProperties> readAllDatabases() {
* @return a {@link CosmosPagedFlux} containing one or several feed response pages of read databases or an error.
*/
public CosmosPagedFlux<CosmosDatabaseProperties> queryDatabases(String query, CosmosQueryRequestOptions options) {
return queryDatabasesInternal(new SqlQuerySpec(query), options);
return queryDatabases(new SqlQuerySpec(query), options);
}

/**
Expand All @@ -454,7 +447,13 @@ public CosmosPagedFlux<CosmosDatabaseProperties> queryDatabases(String query, Co
* @return a {@link CosmosPagedFlux} containing one or several feed response pages of read databases or an error.
*/
public CosmosPagedFlux<CosmosDatabaseProperties> queryDatabases(SqlQuerySpec querySpec, CosmosQueryRequestOptions options) {
return queryDatabasesInternal(querySpec, options);
return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> {
setContinuationTokenAndMaxItemCount(pagedFluxOptions, options);
return getDocClientWrapper().queryDatabases(querySpec, options)
.map(response -> BridgeInternal.createFeedResponse(
ModelBridgeInternal.getCosmosDatabasePropertiesFromV2Results(response.getResults()),
response.getResponseHeaders()));
});
}

/**
Expand All @@ -474,77 +473,4 @@ public CosmosAsyncDatabase getDatabase(String id) {
public void close() {
asyncDocumentClient.close();
}

TracerProvider getTracerProvider(){
return this.tracerProvider;
}

private CosmosPagedFlux<CosmosDatabaseProperties> queryDatabasesInternal(SqlQuerySpec querySpec, CosmosQueryRequestOptions options){
return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> {
String spanName = "queryDatabases";
pagedFluxOptions.setTracerInformation(this.tracerProvider, spanName, this.serviceEndpoint, null);
setContinuationTokenAndMaxItemCount(pagedFluxOptions, options);
return getDocClientWrapper().queryDatabases(querySpec, options)
.map(response -> BridgeInternal.createFeedResponse(
ModelBridgeInternal.getCosmosDatabasePropertiesFromV2Results(response.getResults()),
response.getResponseHeaders()));
}, this.tracerProvider.isEnabled());
}


private Mono<CosmosDatabaseResponse> createDatabaseIfNotExistsInternal(CosmosAsyncDatabase database,
ThroughputProperties throughputProperties, Context context) {
String spanName = "createDatabaseIfNotExists." + database.getId();
Context nestedContext = context.addData(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.COSMOS_CALL_DEPTH_VAL);
Mono<CosmosDatabaseResponse> responseMono = createDatabaseIfNotExistsInternal(database.readInternal(new CosmosDatabaseRequestOptions(), nestedContext), database, throughputProperties, nestedContext);
return tracerProvider.traceEnabledCosmosResponsePublisher(responseMono,
context,
spanName,
database.getId(),
this.serviceEndpoint);
}

private Mono<CosmosDatabaseResponse> createDatabaseIfNotExistsInternal(Mono<CosmosDatabaseResponse> responseMono, CosmosAsyncDatabase database, ThroughputProperties throughputProperties, Context context) {
return responseMono.onErrorResume(exception -> {
final Throwable unwrappedException = Exceptions.unwrap(exception);
if (unwrappedException instanceof CosmosException) {
final CosmosException cosmosException = (CosmosException) unwrappedException;
if (cosmosException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) {
CosmosDatabaseRequestOptions requestOptions = new CosmosDatabaseRequestOptions();
if(throughputProperties != null) {
ModelBridgeInternal.setThroughputProperties(requestOptions, throughputProperties);
}

if (context != null) {
Database wrappedDatabase = new Database();
wrappedDatabase.setId(database.getId());
return createDatabaseInternal(wrappedDatabase,
requestOptions, context);
}

return createDatabase(new CosmosDatabaseProperties(database.getId()),
requestOptions);
}
}
return Mono.error(unwrappedException);
});
}


private Mono<CosmosDatabaseResponse> createDatabaseInternal(Database database, CosmosDatabaseRequestOptions options,
Context context) {
String spanName = "createDatabase." + database.getId();
Mono<CosmosDatabaseResponse> responseMono = createDatabaseInternal(database, options);
return tracerProvider.traceEnabledCosmosResponsePublisher(responseMono,
context,
spanName,
database.getId(),
this.serviceEndpoint);
}

private Mono<CosmosDatabaseResponse> createDatabaseInternal(Database database, CosmosDatabaseRequestOptions options) {
return asyncDocumentClient.createDatabase(database, ModelBridgeInternal.toRequestOptions(options))
.map(databaseResourceResponse -> ModelBridgeInternal.createCosmosDatabaseResponse(databaseResourceResponse))
.single();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -2,16 +2,13 @@
// Licensed under the MIT License.
package com.azure.cosmos;

import com.azure.core.util.Context;
import com.azure.cosmos.implementation.Paths;
import com.azure.cosmos.implementation.RequestOptions;
import com.azure.cosmos.models.CosmosConflictResponse;
import com.azure.cosmos.models.CosmosConflictRequestOptions;
import com.azure.cosmos.models.ModelBridgeInternal;
import reactor.core.publisher.Mono;

import static com.azure.core.util.FluxUtil.withContext;

/**
* Read and delete conflicts
*/
Expand Down Expand Up @@ -67,11 +64,9 @@ public Mono<CosmosConflictResponse> read(CosmosConflictRequestOptions options) {
options = new CosmosConflictRequestOptions();
}
RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options);
if (!this.container.getDatabase().getClient().getTracerProvider().isEnabled()) {
return readInternal(requestOptions);
}
return this.container.getDatabase().getDocClientWrapper().readConflict(getLink(), requestOptions)
.map(response -> ModelBridgeInternal.createCosmosConflictResponse(response)).single();

return withContext(context -> readInternal(requestOptions, context));
}

/**
Expand All @@ -90,11 +85,8 @@ public Mono<CosmosConflictResponse> delete(CosmosConflictRequestOptions options)
options = new CosmosConflictRequestOptions();
}
RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options);
if (!this.container.getDatabase().getClient().getTracerProvider().isEnabled()) {
return deleteInternal(requestOptions);
}

return withContext(context -> deleteInternal(requestOptions, context));
return this.container.getDatabase().getDocClientWrapper().deleteConflict(getLink(), requestOptions)
.map(response -> ModelBridgeInternal.createCosmosConflictResponse(response)).single();
}

String getURIPathSegment() {
Expand All @@ -114,33 +106,4 @@ String getLink() {
builder.append(getId());
return builder.toString();
}

private Mono<CosmosConflictResponse> readInternal(RequestOptions options, Context context) {
String spanName = "readConflict." + getId();
Mono<CosmosConflictResponse> responseMono = this.readInternal(options);
return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context,
spanName,
this.container.getDatabase().getId(),
this.container.getDatabase().getClient().getServiceEndpoint());

}

private Mono<CosmosConflictResponse> readInternal(RequestOptions options) {
return this.container.getDatabase().getDocClientWrapper().readConflict(getLink(), options)
.map(response -> ModelBridgeInternal.createCosmosConflictResponse(response)).single();
}

private Mono<CosmosConflictResponse> deleteInternal(RequestOptions options, Context context) {
String spanName = "deleteConflict." + getId();
Mono<CosmosConflictResponse> responseMono = deleteInternal(options);
return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context,
spanName,
this.container.getDatabase().getId(),
this.container.getDatabase().getClient().getServiceEndpoint());
}

private Mono<CosmosConflictResponse> deleteInternal(RequestOptions options) {
return this.container.getDatabase().getDocClientWrapper().deleteConflict(getLink(), options)
.map(response -> ModelBridgeInternal.createCosmosConflictResponse(response)).single();
}
}
Loading