From 089307a3a8167f96dd38ecbf3c3a1ca6a1ca32b2 Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Wed, 8 Apr 2020 10:36:47 -0400 Subject: [PATCH 01/40] initial commit --- sdk/cosmos/azure-cosmos/pom.xml | 26 ++- .../com/azure/cosmos/CosmosAsyncClient.java | 155 +++++++++++++++--- .../com/azure/cosmos/CosmosAsyncDatabase.java | 114 ++++++++++++- .../com/azure/cosmos/CosmosClientBuilder.java | 14 ++ .../cosmos/implementation/TracerProvider.java | 116 +++++++++++++ .../src/main/java/module-info.java | 10 +- .../com/azure/cosmos/CosmosTracingSample.java | 88 ++++++++++ 7 files changed, 490 insertions(+), 33 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java create mode 100644 sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracingSample.java diff --git a/sdk/cosmos/azure-cosmos/pom.xml b/sdk/cosmos/azure-cosmos/pom.xml index 14909afd49ec..c56da8757b2f 100644 --- a/sdk/cosmos/azure-cosmos/pom.xml +++ b/sdk/cosmos/azure-cosmos/pom.xml @@ -82,7 +82,26 @@ Licensed under the MIT License. - + + io.opentelemetry + opentelemetry-exporters-jaeger + 0.2.0 + + + io.grpc + grpc-okhttp + 1.23.0 + + + io.opentelemetry + opentelemetry-sdk + 0.2.0 + + + com.azure + azure-core-tracing-opentelemetry + 1.0.0-beta.4 + commons-io commons-io @@ -119,7 +138,7 @@ Licensed under the MIT License. com.azure azure-core - 1.4.0 + 1.4.0-beta.1 io.netty @@ -333,6 +352,9 @@ Licensed under the MIT License. io.netty:netty-transport-native-epoll com.google.code.findbugs:jsr305 + io.opentelemetry:opentelemetry-sdk + io.opentelemetry:opentelemetry-exporters-jaeger + io.grpc:grpc-okhttp diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index e28517af0d70..1bd2dc54cab8 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -3,10 +3,14 @@ package com.azure.cosmos; import com.azure.core.annotation.ServiceClient; +import com.azure.core.util.Context; +import com.azure.core.util.FluxUtil; +import com.azure.core.util.tracing.ProcessKind; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.Configs; 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.CosmosAsyncDatabaseResponse; import com.azure.cosmos.models.CosmosDatabaseProperties; @@ -20,10 +24,19 @@ import io.micrometer.core.instrument.MeterRegistry; import reactor.core.Exceptions; import reactor.core.publisher.Mono; +import reactor.core.publisher.Signal; import java.io.Closeable; +import java.io.PrintWriter; +import java.io.StringWriter; +import java.util.HashMap; import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; +import static com.azure.core.util.FluxUtil.withContext; import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; /** @@ -48,6 +61,7 @@ public final class CosmosAsyncClient implements Closeable { private final CosmosKeyCredential cosmosKeyCredential; private final boolean sessionCapturingOverride; private final boolean enableTransportClientSharing; + private final TracerProvider tracerProvider; CosmosAsyncClient(CosmosClientBuilder builder) { this.configs = builder.configs(); @@ -60,6 +74,7 @@ public final class CosmosAsyncClient implements Closeable { this.cosmosKeyCredential = builder.getKeyCredential(); this.sessionCapturingOverride = builder.isSessionCapturingOverrideEnabled(); this.enableTransportClientSharing = builder.isConnectionReuseAcrossClientsEnabled(); + this.tracerProvider = builder.getTracerProvider(); this.asyncDocumentClient = new AsyncDocumentClient.Builder() .withServiceEndpoint(this.serviceEndpoint) .withMasterKeyOrResourceToken(this.keyOrResourceToken) @@ -189,20 +204,6 @@ public Mono createDatabaseIfNotExists(String id) { return createDatabaseIfNotExistsInternal(getDatabase(id)); } - private Mono createDatabaseIfNotExistsInternal(CosmosAsyncDatabase database) { - return database.read().onErrorResume(exception -> { - final Throwable unwrappedException = Exceptions.unwrap(exception); - if (unwrappedException instanceof CosmosClientException) { - final CosmosClientException cosmosClientException = (CosmosClientException) unwrappedException; - if (cosmosClientException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) { - return createDatabase(new CosmosDatabaseProperties(database.getId()), - new CosmosDatabaseRequestOptions()); - } - } - return Mono.error(unwrappedException); - }); - } - /** * Creates a database. *

@@ -222,10 +223,14 @@ public Mono createDatabase(CosmosDatabaseProperties } Database wrappedDatabase = new Database(); wrappedDatabase.setId(databaseSettings.getId()); - return asyncDocumentClient.createDatabase(wrappedDatabase, ModelBridgeInternal.toRequestOptions(options)) - .map(databaseResourceResponse -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(databaseResourceResponse, - this)) - .single(); + final CosmosDatabaseRequestOptions requestOptions = options; + return withContext(context -> createDatabase(wrappedDatabase, requestOptions, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } /** @@ -280,10 +285,14 @@ public Mono createDatabase(CosmosDatabaseProperties ModelBridgeInternal.setOfferThroughput(options, throughput); Database wrappedDatabase = new Database(); wrappedDatabase.setId(databaseSettings.getId()); - return asyncDocumentClient.createDatabase(wrappedDatabase, ModelBridgeInternal.toRequestOptions(options)) - .map(databaseResourceResponse -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(databaseResourceResponse, - this)) - .single(); + final CosmosDatabaseRequestOptions requestOptions = options; + return withContext(context -> createDatabase(wrappedDatabase, requestOptions, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } /** @@ -410,4 +419,106 @@ public CosmosAsyncDatabase getDatabase(String id) { public void close() { asyncDocumentClient.close(); } + + TracerProvider getTracerProvider(){ + return this.tracerProvider; + } + + private Mono createDatabaseIfNotExistsInternal(CosmosAsyncDatabase database) { + return withContext(context -> createDatabaseIfNotExistsInternal(database, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); + } + + private Mono createDatabaseIfNotExistsInternal(CosmosAsyncDatabase database, + Context context) { + final boolean isTracingEnabled = tracerProvider.isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "createDatabaseIfNotExistsInternal." + database.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, database.getId()); + put(TracerProvider.DB_URL, serviceEndpoint); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return database.read().doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(tracerProvider.startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).onErrorResume(exception -> { + final Throwable unwrappedException = Exceptions.unwrap(exception); + if (unwrappedException instanceof CosmosClientException) { + final CosmosClientException cosmosClientException = (CosmosClientException) unwrappedException; + if (cosmosClientException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) { + return createDatabase(new CosmosDatabaseProperties(database.getId()), + new CosmosDatabaseRequestOptions()); + } + } + return Mono.error(unwrappedException); + }).doOnSuccess(signal -> { + if (isTracingEnabled) { + tracerProvider.endSpan(parentContext.get(), Signal.complete()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + tracerProvider.endSpan(parentContext.get(), Signal.error(throwable)); + } + }); + } + + private Mono createDatabase(Database database, CosmosDatabaseRequestOptions options, Context context){ + final boolean isTracingEnabled = tracerProvider.isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "createDatabase." + database.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, database.getId()); + put(TracerProvider.DB_URL, serviceEndpoint); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return asyncDocumentClient.createDatabase(database, ModelBridgeInternal.toRequestOptions(options)).doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if(!callerFunc.isPresent()) { + parentContext.set(tracerProvider.startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(signal -> { + if (isTracingEnabled) { + tracerProvider.endSpan(parentContext.get(), Signal.complete()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + Map errorAttributes = new HashMap() {{ + put(TracerProvider.ERROR_MSG, throwable.getMessage()); + put(TracerProvider.ERROR_TYPE,throwable.getClass().getName()); + StringWriter errorStack = new StringWriter(); + throwable.printStackTrace(new PrintWriter(errorStack)); + put(TracerProvider.ERROR_STACK, errorStack.toString()); + }}; + tracerProvider.endSpan(parentContext.get(), Signal.error(throwable)); + } + }).map(databaseResourceResponse -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(databaseResourceResponse, + this)) + .single(); + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java index f84535a1ee77..676e1dd103dd 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java @@ -2,10 +2,15 @@ // Licensed under the MIT License. package com.azure.cosmos; +import com.azure.core.util.Context; +import com.azure.core.util.FluxUtil; +import com.azure.core.util.tracing.ProcessKind; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.Offer; import com.azure.cosmos.implementation.Paths; +import com.azure.cosmos.implementation.TracerProvider; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.models.CosmosAsyncContainerResponse; import com.azure.cosmos.models.CosmosAsyncDatabaseResponse; import com.azure.cosmos.models.CosmosAsyncUserResponse; @@ -16,12 +21,19 @@ import com.azure.cosmos.models.FeedOptions; import com.azure.cosmos.models.ModelBridgeInternal; import com.azure.cosmos.models.SqlQuerySpec; -import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.util.CosmosPagedFlux; import com.azure.cosmos.util.UtilBridgeInternal; import reactor.core.Exceptions; import reactor.core.publisher.Mono; +import reactor.core.publisher.Signal; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; + +import static com.azure.core.util.FluxUtil.withContext; import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; /** @@ -76,8 +88,48 @@ public Mono read(CosmosDatabaseRequestOptions optio if (options == null) { options = new CosmosDatabaseRequestOptions(); } - return getDocClientWrapper().readDatabase(getLink(), ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single(); + final CosmosDatabaseRequestOptions requestOptions = options; + return withContext(context -> read(requestOptions, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); + } + + private Mono read(CosmosDatabaseRequestOptions options, Context context) { + final boolean isTracingEnabled = client.getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "read." + this.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, getId()); + put(TracerProvider.DB_URL, client.getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return getDocClientWrapper().readDatabase(getLink(), ModelBridgeInternal.toRequestOptions(options)).doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(client.getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(signal -> { + if (isTracingEnabled) { + client.getTracerProvider().endSpan(parentContext.get(), Signal.complete()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + client.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable)); + } + }).map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single(); } /** @@ -107,8 +159,49 @@ public Mono delete(CosmosDatabaseRequestOptions opt if (options == null) { options = new CosmosDatabaseRequestOptions(); } - return getDocClientWrapper().deleteDatabase(getLink(), ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single(); + + final CosmosDatabaseRequestOptions requestOptions = options; + return withContext(context -> delete(requestOptions, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); + } + + public Mono delete(CosmosDatabaseRequestOptions options, Context context) { + final boolean isTracingEnabled = client.getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "delete." + this.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, getId()); + put(TracerProvider.DB_URL, client.getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return getDocClientWrapper().deleteDatabase(getLink(), ModelBridgeInternal.toRequestOptions(options)).doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(client.getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(signal -> { + if (isTracingEnabled) { + client.getTracerProvider().endSpan(parentContext.get(), Signal.complete()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + client.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable)); + } + }).map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single(); } /* CosmosAsyncContainer operations */ @@ -177,12 +270,23 @@ public Mono createContainer( if (options == null) { options = new CosmosContainerRequestOptions(); } + return getDocClientWrapper() .createCollection(this.getLink(), ModelBridgeInternal.getV2Collection(containerProperties), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, this)).single(); } + private Mono createContainer( + CosmosContainerProperties containerProperties, + CosmosContainerRequestOptions options, + Context Context) { + return getDocClientWrapper() + .createCollection(this.getLink(), ModelBridgeInternal.getV2Collection(containerProperties), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, this)).single(); + } + /** * Creates a document container. *

diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java index 270f6b51e2d4..e21feb07a508 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java @@ -3,11 +3,14 @@ package com.azure.cosmos; import com.azure.core.annotation.ServiceClientBuilder; +import com.azure.core.util.tracing.Tracer; import com.azure.cosmos.implementation.Configs; +import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.models.Permission; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import java.util.List; +import java.util.ServiceLoader; /** * Helper class to buildAsyncClient {@link CosmosAsyncClient} instances @@ -38,11 +41,13 @@ public class CosmosClientBuilder { private CosmosKeyCredential cosmosKeyCredential; private boolean sessionCapturingOverrideEnabled; private boolean connectionReuseAcrossClientsEnabled; + private TracerProvider tracerProvider; /** * Instantiates a new Cosmos client builder. */ public CosmosClientBuilder() { + this.tracerProvider = new TracerProvider(ServiceLoader.load(Tracer.class)); } /** @@ -290,6 +295,15 @@ public CosmosClientBuilder keyCredential(CosmosKeyCredential cosmosKeyCredential return this; } + /** + * Gets the tracer provider + * + * @return tracerProvider + */ + public TracerProvider getTracerProvider() { + return this.tracerProvider; + } + /** * Builds a cosmos configuration object with the provided properties * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java new file mode 100644 index 000000000000..0ffc8f1a4ad2 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java @@ -0,0 +1,116 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +package com.azure.cosmos.implementation; + +import com.azure.core.util.Context; +import com.azure.core.util.logging.ClientLogger; +import com.azure.core.util.tracing.ProcessKind; +import com.azure.core.util.tracing.Tracer; +import com.azure.cosmos.CosmosClientException; +import com.azure.cosmos.models.CosmosAsyncDatabaseResponse; +import com.azure.cosmos.models.CosmosResponse; +import reactor.core.publisher.Signal; + +import java.io.PrintWriter; +import java.io.StringWriter; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +public class TracerProvider { + private final ClientLogger logger = new ClientLogger(TracerProvider.class); + private final List tracers = new ArrayList<>(); + public final static String DB_TYPE_VALUE = "cosmosdb"; + public final static String DB_TYPE = "db.type"; + public final static String DB_INSTANCE = "db.instance"; + public final static String DB_URL = "db.url"; + public static final String DB_STATEMENT = "db.statement"; + public static final String ERROR_MSG = "error.msg"; + public static final String ERROR_TYPE = "error.type"; + public static final String ERROR_STACK = "error.stack"; + public static final String MASTER_CALL = "masterCall"; + public static final String NESTED_CALL = "nestedCall"; + + public static final Object ATTRIBUTE_MAP = "span-attributes"; + + + public TracerProvider(Iterable tracers) { + Objects.requireNonNull(tracers, "'tracers' cannot be null."); + tracers.forEach(e -> this.tracers.add(e)); + } + + public boolean isEnabled() { + return tracers.size() > 0; + } + + /** + * For each tracer plugged into the SDK a new tracing span is created. + * + * The {@code context} will be checked for containing information about a parent span. If a parent span is found the + * new span will be added as a child, otherwise the span will be created and added to the context and any downstream + * start calls will use the created span as the parent. + * + * @param context Additional metadata that is passed through the call stack. + * @return An updated context object. + */ + public Context startSpan(String methodName, Context context, ProcessKind processKind) { + Context local = Objects.requireNonNull(context, "'context' cannot be null."); + Objects.requireNonNull(processKind, "'processKind' cannot be null."); + + for (Tracer tracer : tracers) { + local = tracer.start(methodName, local, processKind); + } + + return local; + } + + /** + * Given a context containing the current tracing span the span is marked completed with status info from + * {@link Signal}. For each tracer plugged into the SDK the current tracing span is marked as completed. + * @param context Additional metadata that is passed through the call stack. + * @param signal The signal indicates the status and contains the metadata we need to end the tracing span. + */ + public void endSpan(Context context, Signal signal) { + Objects.requireNonNull(context, "'context' cannot be null."); + Objects.requireNonNull(signal, "'signal' cannot be null."); + + switch (signal.getType()) { + case ON_COMPLETE: + end("success", null, context); + break; + case ON_ERROR: + String errorCondition = ""; + Throwable throwable = null; + if (signal.hasError()) { + // The last status available is on error, this contains the thrown error. + throwable = signal.getThrowable(); + + if (throwable instanceof CosmosClientException) { + CosmosClientException exception = (CosmosClientException) throwable; + // confirm ? + errorCondition =exception.getError().getErrorDetails(); + } + } + end(errorCondition, throwable, context); + break; + default: + // ON_SUBSCRIBE and ON_NEXT don't have the information to end the span so just return. + break; + } + } + + private void end(String statusMessage, Throwable throwable, Context context) { + for (Tracer tracer : tracers) { + if(throwable != null) { + tracer.setAttribute(TracerProvider.ERROR_MSG, throwable.getMessage(), context); + tracer.setAttribute(TracerProvider.ERROR_TYPE, throwable.getClass().getName(), context); + StringWriter errorStack = new StringWriter(); + throwable.printStackTrace(new PrintWriter(errorStack)); + tracer.setAttribute(TracerProvider.ERROR_STACK, errorStack.toString(), context); + } + tracer.end(statusMessage, throwable, context); + } + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/module-info.java b/sdk/cosmos/azure-cosmos/src/main/java/module-info.java index 06c610563564..3fe48a84f536 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/module-info.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/module-info.java @@ -3,8 +3,7 @@ module com.azure.cosmos { - requires transitive com.azure.core; - + requires transitive azure.core; requires com.fasterxml.jackson.datatype.jsr310; requires io.netty.transport; requires io.netty.handler; @@ -24,13 +23,16 @@ requires jdk.management; requires micrometer.core; requires java.logging; + requires reactor.core; + requires org.reactivestreams; + requires com.fasterxml.jackson.databind; // public API surface area exports com.azure.cosmos; exports com.azure.cosmos.models; // exporting some packages specifically for Jackson - opens com.azure.cosmos to com.fasterxml.jackson.databind; + // opens com.azure.cosmos to com.fasterxml.jackson.databind; opens com.azure.cosmos.models to com.fasterxml.jackson.databind; opens com.azure.cosmos.implementation to com.fasterxml.jackson.databind, java.logging; opens com.azure.cosmos.implementation.routing to com.fasterxml.jackson.databind; @@ -45,4 +47,4 @@ opens com.azure.cosmos.implementation.query.orderbyquery to com.fasterxml.jackson.databind; uses com.azure.cosmos.implementation.guava25.base.PatternCompiler; -} \ No newline at end of file +} diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracingSample.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracingSample.java new file mode 100644 index 000000000000..5e5651827b35 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracingSample.java @@ -0,0 +1,88 @@ +//package com.azure.cosmos; +//import com.azure.cosmos.models.CosmosAsyncDatabaseResponse; +//import com.azure.cosmos.models.CosmosDatabaseProperties; +//import com.azure.cosmos.models.CosmosDatabaseRequestOptions; +//import io.grpc.ManagedChannel; +//import io.grpc.ManagedChannelBuilder; +//import io.opentelemetry.OpenTelemetry; +//import io.opentelemetry.context.Scope; +//import io.opentelemetry.exporters.jaeger.JaegerGrpcSpanExporter; +//import io.opentelemetry.sdk.trace.TracerSdkFactory; +//import io.opentelemetry.sdk.trace.export.SimpleSpansProcessor; +//import io.opentelemetry.trace.Span; +//import io.opentelemetry.trace.Tracer; +//import reactor.core.publisher.Mono; +//import reactor.util.context.Context; +// +//import java.util.logging.Logger; +// +//import static com.azure.core.util.tracing.Tracer.PARENT_SPAN_KEY; +//import static java.util.logging.Logger.getLogger; +// +//public class CosmosTracingSample { +// final static String ENDPOINT = "https://naveencosmosdb.documents.azure.com:443/"; +// final static String KEY = "IWlgE5b53sHJW8jVmtX7xQ3xN7TMAYyG2cIAzicS6ZKT24GLRPR1oImiIxWAUvsjn0BtuDoygGjX3TsEK6yc0A=="; +// private static final Logger LOGGER = getLogger("Sample"); +// private static final Tracer TRACER; +// private static final TracerSdkFactory TRACER_SDK_PROVIDER; +// static { +// TRACER_SDK_PROVIDER = configureOpenTelemetryAndJaegerExporter(); +// TRACER = TRACER_SDK_PROVIDER.get("Sample1"); +// } +// public static void main(String[] args) { +// doAppConfigClientWork(); +// TRACER_SDK_PROVIDER.shutdown(); +// LOGGER.info("=== Tracer Shutdown ==="); +// } +// +// private static void doAppConfigClientWork() { +// CosmosAsyncClient client = new CosmosClientBuilder() +// .endpoint(ENDPOINT) +// .key(KEY) +// .buildAsyncClient(); +// +// LOGGER.info("=== Start user scoped span ==="); +// +// Span span = TRACER.spanBuilder("user-parent-span").startSpan(); +// try (final Scope scope = TRACER.withSpan(span)) { +// Context traceContext = Context.of(PARENT_SPAN_KEY, TRACER.getCurrentSpan()); +// +// System.out.println("CosmosTracingSample.doCosmosClientWork "+scope); +// // Create database if not exists +// Mono databaseIfNotExists = client.createDatabase(new CosmosDatabaseProperties("passengers")) +// .subscriberContext(traceContext); +// databaseIfNotExists.flatMap(databaseResponse -> { +// CosmosAsyncDatabase database = databaseResponse.getDatabase(); +// System.out.println("Checking database... " + database.getId() + " completed!\n"); +// return Mono.empty(); +// }).block(); +// } finally { +// span.end(); +// client.close(); +// } +// +// } +//// static TracerSdkProvider configureOpenTelemetryAndJaegerExporter() { +//// // logger.info("=== Runing With JaegerExporter ==="); +//// ManagedChannel channel = ManagedChannelBuilder.forAddress("localhost", 14250).usePlaintext().build(); +//// JaegerGrpcSpanExporter exporter = JaegerGrpcSpanExporter.newBuilder() +//// .setChannel(channel) +//// .setServiceName("Sample") +//// .setDeadlineMs(Long.MIN_VALUE) +//// .build(); +//// TracerSdkProvider tracerSdkFactory = (TracerSdkProvider) OpenTelemetry.getTracerProvider(); +//// tracerSdkFactory.addSpanProcessor(SimpleSpansProcessor.newBuilder(exporter).build()); +//// return tracerSdkFactory; +//// } +// static TracerSdkFactory configureOpenTelemetryAndJaegerExporter() { +// // logger.info("=== Runing With JaegerExporter ==="); +// ManagedChannel channel = ManagedChannelBuilder.forAddress("localhost", 14250).usePlaintext().build(); +// JaegerGrpcSpanExporter exporter = JaegerGrpcSpanExporter.newBuilder() +// .setChannel(channel) +// .setServiceName("Sample1") +// .build(); +// TracerSdkFactory tracerSdkFactory = (TracerSdkFactory) OpenTelemetry.getTracerFactory(); +// tracerSdkFactory.addSpanProcessor(SimpleSpansProcessor.newBuilder(exporter).build()); +// return tracerSdkFactory; +// } +//} From 175222e616cf0ff1026ff70735fcc36377d05054 Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Wed, 8 Apr 2020 10:46:01 -0400 Subject: [PATCH 02/40] changing sample --- .../src/test/java/com/azure/cosmos/CosmosTracingSample.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracingSample.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracingSample.java index 5e5651827b35..a50072377f60 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracingSample.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracingSample.java @@ -20,8 +20,8 @@ //import static java.util.logging.Logger.getLogger; // //public class CosmosTracingSample { -// final static String ENDPOINT = "https://naveencosmosdb.documents.azure.com:443/"; -// final static String KEY = "IWlgE5b53sHJW8jVmtX7xQ3xN7TMAYyG2cIAzicS6ZKT24GLRPR1oImiIxWAUvsjn0BtuDoygGjX3TsEK6yc0A=="; +// final static String ENDPOINT = ""; +// final static String KEY = ""; // private static final Logger LOGGER = getLogger("Sample"); // private static final Tracer TRACER; // private static final TracerSdkFactory TRACER_SDK_PROVIDER; From fe03e665cc2fdbef20b5aab7961a690854d8f48c Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Mon, 13 Apr 2020 12:29:48 -0400 Subject: [PATCH 03/40] adding query api tracer support --- .../com/azure/cosmos/CosmosAsyncClient.java | 25 +- .../azure/cosmos/CosmosAsyncContainer.java | 585 +++++++++++++++--- .../com/azure/cosmos/CosmosAsyncDatabase.java | 446 +++++++++---- .../com/azure/cosmos/CosmosAsyncScripts.java | 44 +- .../com/azure/cosmos/CosmosAsyncUser.java | 10 +- .../CosmosPagedFluxOptions.java | 41 ++ .../cosmos/implementation/TracerProvider.java | 13 +- .../azure/cosmos/util/CosmosPagedFlux.java | 67 +- 8 files changed, 993 insertions(+), 238 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index 1bd2dc54cab8..9197b6073932 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -27,8 +27,6 @@ import reactor.core.publisher.Signal; import java.io.Closeable; -import java.io.PrintWriter; -import java.io.StringWriter; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -343,6 +341,8 @@ public Mono createDatabase(String id, int throughpu */ public CosmosPagedFlux readAllDatabases(FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName = "readAllDatabases"; + pagedFluxOptions.setTracerInformation(this.tracerProvider, spanName, this.serviceEndpoint); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().readDatabases(options) .map(response -> @@ -394,6 +394,8 @@ public CosmosPagedFlux queryDatabases(String query, Fe */ public CosmosPagedFlux queryDatabases(SqlQuerySpec querySpec, FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName = "queryDatabases." + querySpec.getQueryText(); + pagedFluxOptions.setTracerInformation(this.tracerProvider, spanName, this.serviceEndpoint); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().queryDatabases(querySpec, options) .map(response -> BridgeInternal.createFeedResponse( @@ -468,13 +470,13 @@ private Mono createDatabaseIfNotExistsInternal(Cosm } } return Mono.error(unwrappedException); - }).doOnSuccess(signal -> { + }).doOnSuccess(response -> { if (isTracingEnabled) { - tracerProvider.endSpan(parentContext.get(), Signal.complete()); + tracerProvider.endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); } }).doOnError(throwable -> { if (isTracingEnabled) { - tracerProvider.endSpan(parentContext.get(), Signal.error(throwable)); + tracerProvider.endSpan(parentContext.get(), Signal.error(throwable), 0); } }); } @@ -502,20 +504,13 @@ private Mono createDatabase(Database database, Cosm context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); } } - }).doOnSuccess(signal -> { + }).doOnSuccess(response -> { if (isTracingEnabled) { - tracerProvider.endSpan(parentContext.get(), Signal.complete()); + tracerProvider.endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); } }).doOnError(throwable -> { if (isTracingEnabled) { - Map errorAttributes = new HashMap() {{ - put(TracerProvider.ERROR_MSG, throwable.getMessage()); - put(TracerProvider.ERROR_TYPE,throwable.getClass().getName()); - StringWriter errorStack = new StringWriter(); - throwable.printStackTrace(new PrintWriter(errorStack)); - put(TracerProvider.ERROR_STACK, errorStack.toString()); - }}; - tracerProvider.endSpan(parentContext.get(), Signal.error(throwable)); + tracerProvider.endSpan(parentContext.get(), Signal.error(throwable), 0); } }).map(databaseResourceResponse -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(databaseResourceResponse, this)) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index f33665cdbc54..2dbe77aeabcc 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -2,12 +2,16 @@ // Licensed under the MIT License. package com.azure.cosmos; +import com.azure.core.util.Context; +import com.azure.core.util.FluxUtil; +import com.azure.core.util.tracing.ProcessKind; import com.azure.cosmos.implementation.CosmosItemProperties; import com.azure.cosmos.implementation.Document; import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.Offer; import com.azure.cosmos.implementation.Paths; import com.azure.cosmos.implementation.RequestOptions; +import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.models.CosmosAsyncContainerResponse; import com.azure.cosmos.models.CosmosAsyncItemResponse; import com.azure.cosmos.models.CosmosConflictProperties; @@ -23,9 +27,16 @@ import com.azure.cosmos.util.UtilBridgeInternal; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.core.publisher.Signal; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; +import static com.azure.core.util.FluxUtil.withContext; import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; /** @@ -83,8 +94,15 @@ public Mono read(CosmosContainerRequestOptions opt if (options == null) { options = new CosmosContainerRequestOptions(); } - return database.getDocClientWrapper().readCollection(getLink(), ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); + + final CosmosContainerRequestOptions requestOptions = options; + return withContext(context -> read(requestOptions, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } /** @@ -102,8 +120,50 @@ public Mono delete(CosmosContainerRequestOptions o if (options == null) { options = new CosmosContainerRequestOptions(); } + + final CosmosContainerRequestOptions requestOptions = options; + return withContext(context -> delete(requestOptions, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); + } + + private Mono delete(CosmosContainerRequestOptions options, Context context) { + final boolean isTracingEnabled = database.getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "deleteContainer." + this.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, getId()); + put(TracerProvider.DB_URL, database.getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + return database.getDocClientWrapper().deleteCollection(getLink(), ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(database.getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }).map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); } /** @@ -156,9 +216,54 @@ public Mono replace( if (options == null) { options = new CosmosContainerRequestOptions(); } + + final CosmosContainerRequestOptions requestOptions = options; + return withContext(context -> replace(containerProperties, requestOptions, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); + } + + private Mono replace(CosmosContainerProperties containerProperties, + CosmosContainerRequestOptions options, + Context context) { + final boolean isTracingEnabled = database.getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "replaceContainer." + this.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, getId()); + put(TracerProvider.DB_URL, database.getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + return database.getDocClientWrapper() - .replaceCollection(ModelBridgeInternal.getV2Collection(containerProperties), ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); + .replaceCollection(ModelBridgeInternal.getV2Collection(containerProperties), ModelBridgeInternal.toRequestOptions(options)) + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + + parentContext.set(database.getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }).map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); } /* CosmosAsyncItem operations */ @@ -203,7 +308,6 @@ public Mono> createItem( return createItem(item, options); } - /** * Create an item. * @@ -216,16 +320,58 @@ public Mono> createItem(T item, CosmosItemRequest if (options == null) { options = new CosmosItemRequestOptions(); } + + final CosmosItemRequestOptions requestOptions = options; + return withContext(context -> createItem(item, requestOptions, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); + } + + private Mono> createItem(T item, CosmosItemRequestOptions options, Context context) { + final boolean isTracingEnabled = database.getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "createItem." + this.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, getId()); + put(TracerProvider.DB_URL, database.getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + @SuppressWarnings("unchecked") Class itemType = (Class) item.getClass(); RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); return database.getDocClientWrapper() - .createDocument(getLink(), - item, - requestOptions, - true) - .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) - .single(); + .createDocument(getLink(), + item, + requestOptions, + true) + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(database.getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }).map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) + .single(); } /** @@ -259,14 +405,15 @@ public Mono> upsertItem(T item, CosmosItemRequest if (options == null) { options = new CosmosItemRequestOptions(); } - @SuppressWarnings("unchecked") - Class itemType = (Class) item.getClass(); - return this.getDatabase().getDocClientWrapper() - .upsertDocument(this.getLink(), item, - ModelBridgeInternal.toRequestOptions(options), - true) - .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) - .single(); + + final CosmosItemRequestOptions requestOptions = options; + return withContext(context -> upsertItem(item, requestOptions, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } /** @@ -300,6 +447,9 @@ public CosmosPagedFlux readAllItems(Class classType) { */ public CosmosPagedFlux readAllItems(FeedOptions options, Class classType) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName = "readAllItems." + this.getId(); + pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), spanName, + this.getDatabase().getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDatabase().getDocClientWrapper().readDocuments(getLink(), options).map( response -> prepareFeedResponse(response, classType)); @@ -379,6 +529,9 @@ public CosmosPagedFlux queryItems(SqlQuerySpec querySpec, FeedOptions opt private CosmosPagedFlux queryItemsInternal( SqlQuerySpec sqlQuerySpec, FeedOptions feedOptions, Class classType) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName = "queryItems." + this.getId() + "." + sqlQuerySpec.getQueryText(); + pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), spanName, + this.getDatabase().getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, feedOptions); return getDatabase().getDocClientWrapper() .queryDocuments(CosmosAsyncContainer.this.getLink(), sqlQuerySpec, feedOptions) @@ -394,7 +547,6 @@ private FeedResponse prepareFeedResponse(FeedResponse response, ModelBridgeInternal.queryMetrics(response)); } - /** * Reads an item. *

@@ -432,12 +584,16 @@ public Mono> readItem( if (options == null) { options = new CosmosItemRequestOptions(); } + ModelBridgeInternal.setPartitionKey(options, partitionKey); RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); - return this.getDatabase().getDocClientWrapper() - .readDocument(getItemLink(itemId), requestOptions) - .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) - .single(); + return withContext(context -> readItem(itemId, partitionKey, requestOptions, itemType, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } /** @@ -481,11 +637,14 @@ public Mono> replaceItem( ModelBridgeInternal.setPartitionKey(options, partitionKey); @SuppressWarnings("unchecked") Class itemType = (Class) item.getClass(); - return this.getDatabase() - .getDocClientWrapper() - .replaceDocument(getItemLink(itemId), doc, ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) - .single(); + final CosmosItemRequestOptions requestOptions = options; + return withContext(context -> replaceItem(itemType, itemId, doc, requestOptions,context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } /** @@ -523,11 +682,13 @@ public Mono> deleteItem( } ModelBridgeInternal.setPartitionKey(options, partitionKey); RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); - return this.getDatabase() - .getDocClientWrapper() - .deleteDocument(getItemLink(itemId), requestOptions) - .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponseWithObjectType(response)) - .single(); + return withContext(context -> deleteItem(itemId, requestOptions, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } private String getItemLink(String itemId) { @@ -561,6 +722,9 @@ public CosmosAsyncScripts getScripts() { */ public CosmosPagedFlux readAllConflicts(FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName = "readAllConflicts." + this.getId(); + pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), spanName, + this.getDatabase().getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper().readConflicts(getLink(), options) .map(response -> BridgeInternal.createFeedResponse( @@ -590,6 +754,9 @@ public CosmosPagedFlux queryConflicts(String query) { */ public CosmosPagedFlux queryConflicts(String query, FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName = "queryConflicts." + this.getId() + "." + query; + pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), spanName, + this.getDatabase().getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper().queryConflicts(getLink(), query, options) .map(response -> BridgeInternal.createFeedResponse( @@ -614,23 +781,13 @@ public CosmosAsyncConflict getConflict(String id) { * @return a {@link Mono} containing throughput or an error. */ public Mono readProvisionedThroughput() { - return this.read() - .flatMap(cosmosContainerResponse -> - database.getDocClientWrapper() - .queryOffers("select * from c where c.offerResourceId = '" - + cosmosContainerResponse.getProperties() - .getResourceId() + "'", new FeedOptions()) - .single()) - .flatMap(offerFeedResponse -> { - if (offerFeedResponse.getResults().isEmpty()) { - return Mono.error( - BridgeInternal.createCosmosClientException(HttpConstants.StatusCodes.BADREQUEST, - "No offers found for the resource")); - } - return database.getDocClientWrapper() - .readOffer(offerFeedResponse.getResults().get(0).getSelfLink()) - .single(); - }).map(cosmosOfferResponse -> cosmosOfferResponse.getResource().getThroughput()); + return withContext(context -> readProvisionedThroughput(context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } /** @@ -642,23 +799,13 @@ public Mono readProvisionedThroughput() { * @return a {@link Mono} containing throughput or an error. */ public Mono replaceProvisionedThroughput(int requestUnitsPerSecond) { - return this.read() - .flatMap(cosmosContainerResponse -> - database.getDocClientWrapper() - .queryOffers("select * from c where c.offerResourceId = '" - + cosmosContainerResponse.getProperties() - .getResourceId() + "'", new FeedOptions()) - .single()) - .flatMap(offerFeedResponse -> { - if (offerFeedResponse.getResults().isEmpty()) { - return Mono.error( - BridgeInternal.createCosmosClientException(HttpConstants.StatusCodes.BADREQUEST, - "No offers found for the resource")); - } - Offer offer = offerFeedResponse.getResults().get(0); - offer.setThroughput(requestUnitsPerSecond); - return database.getDocClientWrapper().replaceOffer(offer).single(); - }).map(offerResourceResponse -> offerResourceResponse.getResource().getThroughput()); + return withContext(context -> replaceProvisionedThroughput(requestUnitsPerSecond, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } /** @@ -681,4 +828,304 @@ String getParentLink() { String getLink() { return this.link; } + + private Mono> deleteItem( + String itemId, + RequestOptions requestOptions, + Context context) { + final boolean isTracingEnabled = database.getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "deleteItem." + this.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, getId()); + put(TracerProvider.DB_URL, database.getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return this.getDatabase() + .getDocClientWrapper() + .deleteDocument(getItemLink(itemId), requestOptions) + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(database.getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }).map(response -> ModelBridgeInternal.createCosmosAsyncItemResponseWithObjectType(response)) + .single(); + } + + private Mono> replaceItem( + Class itemType, + String itemId, + Document doc, + CosmosItemRequestOptions options, + Context context) { + final boolean isTracingEnabled = database.getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "replaceItem." + this.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, getId()); + put(TracerProvider.DB_URL, database.getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return this.getDatabase() + .getDocClientWrapper() + .replaceDocument(getItemLink(itemId), doc, ModelBridgeInternal.toRequestOptions(options)) + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(database.getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }).map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) + .single(); + } + + private Mono> upsertItem(T item, CosmosItemRequestOptions options, Context context) { + final boolean isTracingEnabled = database.getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "upsertItem." + this.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, getId()); + put(TracerProvider.DB_URL, database.getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + @SuppressWarnings("unchecked") + Class itemType = (Class) item.getClass(); + return this.getDatabase().getDocClientWrapper() + .upsertDocument(this.getLink(), item, + ModelBridgeInternal.toRequestOptions(options), + true) + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(database.getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }).map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) + .single(); + } + + private Mono> readItem( + String itemId, PartitionKey partitionKey, + RequestOptions requestOptions, Class itemType, + Context context) { + final boolean isTracingEnabled = database.getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "readItem." + this.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, getId()); + put(TracerProvider.DB_URL, database.getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return this.getDatabase().getDocClientWrapper() + .readDocument(getItemLink(itemId), requestOptions) + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(database.getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }).map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) + .single(); + } + + private Mono read(CosmosContainerRequestOptions options, Context context) { + final boolean isTracingEnabled = database.getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "readContainer." + this.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, getId()); + put(TracerProvider.DB_URL, database.getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return database.getDocClientWrapper().readCollection(getLink(), ModelBridgeInternal.toRequestOptions(options)) + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(database.getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }).map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); + } + + private Mono readProvisionedThroughput(Context context) { + final boolean isTracingEnabled = this.getDatabase().getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "readProvisionedThroughput." + this.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, getId()); + put(TracerProvider.DB_URL, getDatabase().getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return this.read() + .flatMap(cosmosContainerResponse -> + database.getDocClientWrapper() + .queryOffers("select * from c where c.offerResourceId = '" + + cosmosContainerResponse.getProperties() + .getResourceId() + "'", new FeedOptions()) + .single()) + .flatMap(offerFeedResponse -> { + if (offerFeedResponse.getResults().isEmpty()) { + return Mono.error( + BridgeInternal.createCosmosClientException(HttpConstants.StatusCodes.BADREQUEST, + "No offers found for the resource")); + } + return database.getDocClientWrapper() + .readOffer(offerFeedResponse.getResults().get(0).getSelfLink()) + .single(); + }).map(cosmosOfferResponse -> cosmosOfferResponse.getResource().getThroughput()).doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if(!callerFunc.isPresent()) { + parentContext.set(getDatabase().getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), 200); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }); + } + + private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Context context) { + final boolean isTracingEnabled = this.getDatabase().getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "replaceProvisionedThroughput." + this.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, getId()); + put(TracerProvider.DB_URL, getDatabase().getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return this.read() + .flatMap(cosmosContainerResponse -> + database.getDocClientWrapper() + .queryOffers("select * from c where c.offerResourceId = '" + + cosmosContainerResponse.getProperties() + .getResourceId() + "'", new FeedOptions()) + .single()) + .flatMap(offerFeedResponse -> { + if (offerFeedResponse.getResults().isEmpty()) { + return Mono.error( + BridgeInternal.createCosmosClientException(HttpConstants.StatusCodes.BADREQUEST, + "No offers found for the resource")); + } + Offer offer = offerFeedResponse.getResults().get(0); + offer.setThroughput(requestUnitsPerSecond); + return database.getDocClientWrapper().replaceOffer(offer).single(); + }).map(offerResourceResponse -> offerResourceResponse.getResource().getThroughput()).doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if(!callerFunc.isPresent()) { + parentContext.set(getDatabase().getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), 200); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }); + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java index 676e1dd103dd..5614a9cf8ae2 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java @@ -98,39 +98,7 @@ public Mono read(CosmosDatabaseRequestOptions optio }); } - private Mono read(CosmosDatabaseRequestOptions options, Context context) { - final boolean isTracingEnabled = client.getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; - String spanName = "read." + this.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, getId()); - put(TracerProvider.DB_URL, client.getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - return getDocClientWrapper().readDatabase(getLink(), ModelBridgeInternal.toRequestOptions(options)).doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(client.getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(signal -> { - if (isTracingEnabled) { - client.getTracerProvider().endSpan(parentContext.get(), Signal.complete()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - client.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable)); - } - }).map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single(); - } /** * Deletes a database. @@ -170,40 +138,6 @@ public Mono delete(CosmosDatabaseRequestOptions opt }); } - public Mono delete(CosmosDatabaseRequestOptions options, Context context) { - final boolean isTracingEnabled = client.getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; - String spanName = "delete." + this.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, getId()); - put(TracerProvider.DB_URL, client.getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - - return getDocClientWrapper().deleteDatabase(getLink(), ModelBridgeInternal.toRequestOptions(options)).doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(client.getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(signal -> { - if (isTracingEnabled) { - client.getTracerProvider().endSpan(parentContext.get(), Signal.complete()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - client.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable)); - } - }).map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single(); - } - /* CosmosAsyncContainer operations */ /** @@ -271,21 +205,24 @@ public Mono createContainer( options = new CosmosContainerRequestOptions(); } - return getDocClientWrapper() - .createCollection(this.getLink(), ModelBridgeInternal.getV2Collection(containerProperties), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, this)).single(); + final CosmosContainerRequestOptions requestOptions = options; + return withContext(context -> + { + Optional parentSpan = context.getData("user-span-name"); + if(parentSpan.isPresent()) { + System.out.println("here"); + } + return createContainer(containerProperties, requestOptions, context); + }).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } - private Mono createContainer( - CosmosContainerProperties containerProperties, - CosmosContainerRequestOptions options, - Context Context) { - return getDocClientWrapper() - .createCollection(this.getLink(), ModelBridgeInternal.getV2Collection(containerProperties), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, this)).single(); - } + /** * Creates a document container. @@ -432,15 +369,12 @@ public Mono createContainerIfNotExists( private Mono createContainerIfNotExistsInternal( CosmosContainerProperties containerProperties, CosmosAsyncContainer container, CosmosContainerRequestOptions options) { - return container.read(options).onErrorResume(exception -> { - final Throwable unwrappedException = Exceptions.unwrap(exception); - if (unwrappedException instanceof CosmosClientException) { - final CosmosClientException cosmosClientException = (CosmosClientException) unwrappedException; - if (cosmosClientException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) { - return createContainer(containerProperties, options); - } + return withContext(context -> createContainerIfNotExistsInternal(containerProperties, container, options, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); } - return Mono.error(unwrappedException); + return reactorContext.put(TracerProvider.MASTER_CALL, true); }); } @@ -457,6 +391,9 @@ private Mono createContainerIfNotExistsInternal( */ public CosmosPagedFlux readAllContainers(FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName = "readAllContainers." + this.getId(); + pagedFluxOptions.setTracerInformation(this.getClient().getTracerProvider(), spanName, + this.getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().readCollections(getLink(), options) .map(response -> BridgeInternal.createFeedResponse( @@ -539,6 +476,9 @@ public CosmosPagedFlux queryContainers(SqlQuerySpec q */ public CosmosPagedFlux queryContainers(SqlQuerySpec querySpec, FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName = "queryContainers." + this.getId() + "." + querySpec.getQueryText(); + pagedFluxOptions.setTracerInformation(this.getClient().getTracerProvider(), spanName, + this.getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().queryCollections(getLink(), querySpec, options) .map(response -> BridgeInternal.createFeedResponse( @@ -618,6 +558,9 @@ public CosmosPagedFlux readAllUsers() { */ public CosmosPagedFlux readAllUsers(FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName = "readAllUsers." + this.getId(); + pagedFluxOptions.setTracerInformation(this.getClient().getTracerProvider(), spanName, + this.getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().readUsers(getLink(), options) .map(response -> BridgeInternal.createFeedResponse( @@ -686,6 +629,9 @@ public CosmosPagedFlux queryUsers(SqlQuerySpec querySpec) */ public CosmosPagedFlux queryUsers(SqlQuerySpec querySpec, FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName = "queryUsers." + this.getId() + "." + querySpec.getQueryText(); + pagedFluxOptions.setTracerInformation(this.getClient().getTracerProvider(), spanName, + this.getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().queryUsers(getLink(), querySpec, options) .map(response -> BridgeInternal.createFeedResponseWithQueryMetrics( @@ -710,29 +656,13 @@ public CosmosAsyncUser getUser(String id) { * @return a {@link Mono} containing throughput or an error. */ public Mono readProvisionedThroughput() { - return this.read() - .flatMap(cosmosDatabaseResponse -> getDocClientWrapper() - .queryOffers("select * from c where c.offerResourceId = '" - + cosmosDatabaseResponse - .getProperties() - .getResourceId() + "'", - new FeedOptions()) - .single() - .flatMap(offerFeedResponse -> { - if (offerFeedResponse.getResults().isEmpty()) { - return Mono.error(BridgeInternal - .createCosmosClientException( - HttpConstants.StatusCodes.BADREQUEST, - "No offers found for the resource")); - } - return getDocClientWrapper() - .readOffer(offerFeedResponse.getResults() - .get(0) - .getSelfLink()) - .single(); - }).map(cosmosContainerResponse1 -> cosmosContainerResponse1 - .getResource() - .getThroughput())); + return withContext(context -> readProvisionedThroughput(context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } /** @@ -744,27 +674,13 @@ public Mono readProvisionedThroughput() { * @return a {@link Mono} containing throughput or an error. */ public Mono replaceProvisionedThroughput(int requestUnitsPerSecond) { - return this.read() - .flatMap(cosmosDatabaseResponse -> this.getDocClientWrapper() - .queryOffers("select * from c where c.offerResourceId = '" - + cosmosDatabaseResponse.getProperties() - .getResourceId() - + "'", new FeedOptions()) - .single() - .flatMap(offerFeedResponse -> { - if (offerFeedResponse.getResults().isEmpty()) { - return Mono.error(BridgeInternal - .createCosmosClientException( - HttpConstants.StatusCodes.BADREQUEST, - "No offers found for the resource")); - } - Offer offer = offerFeedResponse.getResults().get(0); - offer.setThroughput(requestUnitsPerSecond); - return this.getDocClientWrapper().replaceOffer(offer) - .single(); - }).map(offerResourceResponse -> offerResourceResponse - .getResource() - .getThroughput())); + return withContext(context -> replaceProvisionedThroughput(requestUnitsPerSecond, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } CosmosAsyncClient getClient() { @@ -787,4 +703,272 @@ String getLink() { return this.link; } + private Mono createContainerIfNotExistsInternal( + CosmosContainerProperties containerProperties, + CosmosAsyncContainer container, + CosmosContainerRequestOptions options, + Context context) { + final boolean isTracingEnabled = client.getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "createContainerIfNotExistsInternal." + containerProperties.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, getId()); + put(TracerProvider.DB_URL, client.getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return container.read(options).onErrorResume(exception -> { + final Throwable unwrappedException = Exceptions.unwrap(exception); + if (unwrappedException instanceof CosmosClientException) { + final CosmosClientException cosmosClientException = (CosmosClientException) unwrappedException; + if (cosmosClientException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) { + return createContainer(containerProperties, options); + } + } + return Mono.error(unwrappedException); + }).doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(client.getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + client.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + client.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }); + } + + private Mono createContainer( + CosmosContainerProperties containerProperties, + CosmosContainerRequestOptions options, + Context context) { + final boolean isTracingEnabled = client.getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "createContainer." + containerProperties.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, getId()); + put(TracerProvider.DB_URL, client.getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return getDocClientWrapper() + .createCollection(this.getLink(), ModelBridgeInternal.getV2Collection(containerProperties), + ModelBridgeInternal.toRequestOptions(options)) + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(client.getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + client.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + client.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }).map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, this)).single(); + } + + private Mono read(CosmosDatabaseRequestOptions options, Context context) { + final boolean isTracingEnabled = client.getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "readDatabase." + this.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, getId()); + put(TracerProvider.DB_URL, client.getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return getDocClientWrapper().readDatabase(getLink(), ModelBridgeInternal.toRequestOptions(options)).doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(client.getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + client.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + client.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }).map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single(); + } + + private Mono delete(CosmosDatabaseRequestOptions options, Context context) { + final boolean isTracingEnabled = client.getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "deleteDatabase." + this.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, getId()); + put(TracerProvider.DB_URL, client.getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return getDocClientWrapper().deleteDatabase(getLink(), ModelBridgeInternal.toRequestOptions(options)).doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(client.getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + client.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + client.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }).map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single(); + } + + private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Context context) { + final boolean isTracingEnabled = this.client.getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "replaceOffer." + this.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, getId()); + put(TracerProvider.DB_URL, client.getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return this.read() + .flatMap(cosmosDatabaseResponse -> this.getDocClientWrapper() + .queryOffers("select * from c where c.offerResourceId = '" + + cosmosDatabaseResponse.getProperties() + .getResourceId() + + "'", new FeedOptions()) + .single() + .flatMap(offerFeedResponse -> { + if (offerFeedResponse.getResults().isEmpty()) { + return Mono.error(BridgeInternal + .createCosmosClientException( + HttpConstants.StatusCodes.BADREQUEST, + "No offers found for the resource")); + } + Offer offer = offerFeedResponse.getResults().get(0); + offer.setThroughput(requestUnitsPerSecond); + return this.getDocClientWrapper().replaceOffer(offer) + .single(); + }).map(offerResourceResponse -> offerResourceResponse + .getResource() + .getThroughput())) + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(this.client.getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + this.client.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), + 200); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + this.client.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }); + } + + private Mono readProvisionedThroughput(Context context) { + final boolean isTracingEnabled = this.client.getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "readProvisionedThroughput." + this.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, getId()); + put(TracerProvider.DB_URL, client.getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return this.read() + .flatMap(cosmosDatabaseResponse -> getDocClientWrapper() + .queryOffers("select * from c where c.offerResourceId = '" + + cosmosDatabaseResponse + .getProperties() + .getResourceId() + "'", + new FeedOptions()) + .single() + .flatMap(offerFeedResponse -> { + if (offerFeedResponse.getResults().isEmpty()) { + return Mono.error(BridgeInternal + .createCosmosClientException( + HttpConstants.StatusCodes.BADREQUEST, + "No offers found for the resource")); + } + return getDocClientWrapper() + .readOffer(offerFeedResponse.getResults() + .get(0) + .getSelfLink()) + .single(); + }).map(cosmosContainerResponse1 -> cosmosContainerResponse1 + .getResource() + .getThroughput())).doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(this.client.getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), + ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + this.client.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), + 200); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + this.client.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }); + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java index 7c066bbf3a1a..0e4e0bedd5dc 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java @@ -86,12 +86,16 @@ public Mono createStoredProcedure( * In case of failure the {@link CosmosPagedFlux} will error. * * @param options the feed options. - * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the read cosmos stored + * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the read cosmos stored * procedures * properties or an error. */ public CosmosPagedFlux readAllStoredProcedures(FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName = "readAllStoredProcedures." + this.container.getId(); + pagedFluxOptions.setTracerInformation(this.container.getDatabase().getClient().getTracerProvider(), + spanName, + this.container.getDatabase().getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper() .readStoredProcedures(container.getLink(), options) @@ -110,7 +114,7 @@ public CosmosPagedFlux readAllStoredProcedures( * * @param query the the query. * @param options the feed options. - * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the obtained stored + * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the obtained stored * procedures or * an error. */ @@ -129,7 +133,7 @@ public CosmosPagedFlux queryStoredProcedures( * * @param querySpec the SQL query specification. * @param options the feed options. - * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the obtained stored + * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the obtained stored * procedures or * an error. */ @@ -137,6 +141,10 @@ public CosmosPagedFlux queryStoredProcedures( SqlQuerySpec querySpec, FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName = "queryStoredProcedures." + this.container.getId() + "." + querySpec.getQueryText(); + pagedFluxOptions.setTracerInformation(this.container.getDatabase().getClient().getTracerProvider(), + spanName, + this.container.getDatabase().getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper() .queryStoredProcedures(container.getLink(), querySpec, options) @@ -189,12 +197,16 @@ public Mono createUserDefinedFunction( * In case of failure the {@link CosmosPagedFlux} will error. * * @param options the feed options. - * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the read user defined + * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the read user defined * functions or an * error. */ public CosmosPagedFlux readAllUserDefinedFunctions(FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName = "readAllUserDefinedFunctions." + this.container.getId(); + pagedFluxOptions.setTracerInformation(this.container.getDatabase().getClient().getTracerProvider(), + spanName, + this.container.getDatabase().getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper() .readUserDefinedFunctions(container.getLink(), options) @@ -208,13 +220,13 @@ public CosmosPagedFlux readAllUserDefinedFu * Query for user defined functions in the container. *

* After subscription the operation will be performed. - * The {@link CosmosPagedFlux} will contain one or several feed response pages of the obtained user defined + * The {@link CosmosPagedFlux} will contain one or several feed response pages of the obtained user defined * functions. * In case of failure the {@link CosmosPagedFlux} will error. * * @param query the query. * @param options the feed options. - * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the obtained user defined + * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the obtained user defined * functions * or an error. */ @@ -228,13 +240,13 @@ public CosmosPagedFlux queryUserDefinedFunc * Query for user defined functions in the container. *

* After subscription the operation will be performed. - * The {@link CosmosPagedFlux} will contain one or several feed response pages of the obtained user defined + * The {@link CosmosPagedFlux} will contain one or several feed response pages of the obtained user defined * functions. * In case of failure the {@link CosmosPagedFlux} will error. * * @param querySpec the SQL query specification. * @param options the feed options. - * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the obtained user defined + * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the obtained user defined * functions * or an error. */ @@ -242,6 +254,10 @@ public CosmosPagedFlux queryUserDefinedFunc SqlQuerySpec querySpec, FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName = "queryUserDefinedFunctions." + this.container.getId() + "." + querySpec.getQueryText(); + pagedFluxOptions.setTracerInformation(this.container.getDatabase().getClient().getTracerProvider(), + spanName, + this.container.getDatabase().getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper() .queryUserDefinedFunctions(container.getLink(), querySpec, options) @@ -286,17 +302,21 @@ public Mono createTrigger(CosmosTriggerProperties pr * Reads all triggers in a container *

* After subscription the operation will be performed. - * The {@link CosmosPagedFlux} will contain one or several feed response pages of the read cosmos trigger + * The {@link CosmosPagedFlux} will contain one or several feed response pages of the read cosmos trigger * properties. * In case of failure the {@link CosmosPagedFlux} will error. * * @param options the feed options. - * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the read cosmos rigger + * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the read cosmos rigger * properties or * an error. */ public CosmosPagedFlux readAllTriggers(FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName = "readAllTriggers." + this.container.getId(); + pagedFluxOptions.setTracerInformation(this.container.getDatabase().getClient().getTracerProvider(), + spanName, + this.container.getDatabase().getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper() .readTriggers(container.getLink(), options) @@ -338,6 +358,10 @@ public CosmosPagedFlux queryTriggers( SqlQuerySpec querySpec, FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName = "queryTriggers." + this.container.getId(); + pagedFluxOptions.setTracerInformation(this.container.getDatabase().getClient().getTracerProvider(), + spanName, + this.container.getDatabase().getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper() .queryTriggers(container.getLink(), querySpec, options) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java index 1692ceaf0b21..7bf12cdaad14 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java @@ -141,11 +141,15 @@ public Mono upsertPermission( * In case of failure the {@link CosmosPagedFlux} will error. * * @param options the feed options. - * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the read permissions or an + * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the read permissions or an * error. */ public CosmosPagedFlux readAllPermissions(FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName = "readAllPermissions." + this.getId(); + pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), + spanName, + this.getDatabase().getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDatabase().getDocClientWrapper() .readPermissions(getLink(), options) @@ -184,6 +188,10 @@ public CosmosPagedFlux queryPermissions(String query */ public CosmosPagedFlux queryPermissions(String query, FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName = "queryPermissions." + this.getId() + "." + query; + pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), + spanName, + this.getDatabase().getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDatabase().getDocClientWrapper() .queryPermissions(getLink(), query, options) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosPagedFluxOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosPagedFluxOptions.java index c4407b68bfee..eb6340bba9ba 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosPagedFluxOptions.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosPagedFluxOptions.java @@ -5,6 +5,9 @@ import com.azure.cosmos.util.CosmosPagedFlux; +import java.util.HashMap; +import java.util.Map; + /** * Specifies paging options for Cosmos Paged Flux implementation. * @see CosmosPagedFlux @@ -13,6 +16,10 @@ public class CosmosPagedFluxOptions { private String requestContinuation; private Integer maxItemCount; + private Map tracingAttributes; + private TracerProvider tracerProvider; + private String tracerSpanName; + public CosmosPagedFluxOptions() {} @@ -57,4 +64,38 @@ public CosmosPagedFluxOptions setMaxItemCount(Integer maxItemCount) { this.maxItemCount = maxItemCount; return this; } + + /** + * Gets the tracing attributes + * @return tracingAttributes + */ + public Map getTracingAttributes() { + return this.tracingAttributes; + } + + /** + * Gets the tracer provider + * @return tracerProvider + */ + public TracerProvider getTracerProvider() { + return this.tracerProvider; + } + + /** + * Gets the tracer span name + * @return tracerSpanName + */ + public String getTracerSpanName() { + return tracerSpanName; + } + + public void setTracerInformation(TracerProvider tracerProvider, String tracerSpanName, String serviceUrl) { + this.tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_URL, serviceUrl); + put(TracerProvider.DB_STATEMENT, tracerSpanName); + }}; + this.tracerSpanName = tracerSpanName; + this.tracerProvider = tracerProvider; + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java index 0ffc8f1a4ad2..f3381af9504f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java @@ -72,16 +72,15 @@ public Context startSpan(String methodName, Context context, ProcessKind process * @param context Additional metadata that is passed through the call stack. * @param signal The signal indicates the status and contains the metadata we need to end the tracing span. */ - public void endSpan(Context context, Signal signal) { + public void endSpan(Context context, Signal signal, int statusCode) { Objects.requireNonNull(context, "'context' cannot be null."); Objects.requireNonNull(signal, "'signal' cannot be null."); switch (signal.getType()) { case ON_COMPLETE: - end("success", null, context); + end(statusCode, null, context); break; case ON_ERROR: - String errorCondition = ""; Throwable throwable = null; if (signal.hasError()) { // The last status available is on error, this contains the thrown error. @@ -90,10 +89,10 @@ public void endSpan(Context context, Signal signal) { if (throwable instanceof CosmosClientException) { CosmosClientException exception = (CosmosClientException) throwable; // confirm ? - errorCondition =exception.getError().getErrorDetails(); + statusCode = exception.getStatusCode(); } } - end(errorCondition, throwable, context); + end(statusCode, throwable, context); break; default: // ON_SUBSCRIBE and ON_NEXT don't have the information to end the span so just return. @@ -101,7 +100,7 @@ public void endSpan(Context context, Signal signal) { } } - private void end(String statusMessage, Throwable throwable, Context context) { + private void end(int statusCode, Throwable throwable, Context context) { for (Tracer tracer : tracers) { if(throwable != null) { tracer.setAttribute(TracerProvider.ERROR_MSG, throwable.getMessage(), context); @@ -110,7 +109,7 @@ private void end(String statusMessage, Throwable throwable, Context context) { throwable.printStackTrace(new PrintWriter(errorStack)); tracer.setAttribute(TracerProvider.ERROR_STACK, errorStack.toString(), context); } - tracer.end(statusMessage, throwable, context); + tracer.end(statusCode, throwable, context); } } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java index ca9f20352f3a..1527d08267b9 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java @@ -3,13 +3,23 @@ package com.azure.cosmos.util; +import com.azure.core.http.HttpHeader; +import com.azure.core.util.Context; +import com.azure.core.util.FluxUtil; import com.azure.core.util.IterableStream; import com.azure.core.util.paging.ContinuablePagedFlux; +import com.azure.core.util.tracing.ProcessKind; import com.azure.cosmos.implementation.CosmosPagedFluxOptions; +import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.models.FeedResponse; +import io.netty.handler.codec.http.HttpConstants; import reactor.core.CoreSubscriber; import reactor.core.publisher.Flux; +import reactor.core.publisher.Signal; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; /** @@ -37,8 +47,13 @@ public final class CosmosPagedFlux extends ContinuablePagedFlux> byPage() { CosmosPagedFluxOptions cosmosPagedFluxOptions = new CosmosPagedFluxOptions(); - - return this.optionsFluxFunction.apply(cosmosPagedFluxOptions); + return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } @Override @@ -46,7 +61,13 @@ public Flux> byPage(String continuationToken) { CosmosPagedFluxOptions cosmosPagedFluxOptions = new CosmosPagedFluxOptions(); cosmosPagedFluxOptions.setRequestContinuation(continuationToken); - return this.optionsFluxFunction.apply(cosmosPagedFluxOptions); + return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } @Override @@ -54,7 +75,13 @@ public Flux> byPage(int preferredPageSize) { CosmosPagedFluxOptions cosmosPagedFluxOptions = new CosmosPagedFluxOptions(); cosmosPagedFluxOptions.setMaxItemCount(preferredPageSize); - return this.optionsFluxFunction.apply(cosmosPagedFluxOptions); + return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } @Override @@ -63,7 +90,13 @@ public Flux> byPage(String continuationToken, int preferredPageS cosmosPagedFluxOptions.setRequestContinuation(continuationToken); cosmosPagedFluxOptions.setMaxItemCount(preferredPageSize); - return this.optionsFluxFunction.apply(cosmosPagedFluxOptions); + return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } /** @@ -84,4 +117,28 @@ public void subscribe(CoreSubscriber coreSubscriber) { return Flux.fromIterable(elements); }).subscribe(coreSubscriber); } + + private Flux> byPage(CosmosPagedFluxOptions pagedFluxOptions, Context context) { + final AtomicReference parentContext = new AtomicReference<>(Context.NONE); + + return this.optionsFluxFunction.apply(pagedFluxOptions).doOnSubscribe(ignoredValue -> { + if (pagedFluxOptions.getTracerProvider().isEnabled()) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(pagedFluxOptions.getTracerProvider().startSpan(pagedFluxOptions.getTracerSpanName(), + context.addData(TracerProvider.ATTRIBUTE_MAP, pagedFluxOptions.getTracingAttributes()), ProcessKind.DATABASE)); + } + } + }).doOnEach(responseSignal -> { + if (pagedFluxOptions.getTracerProvider().isEnabled()) { + pagedFluxOptions.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), 200); + } + }).doOnError(throwable -> { + if (pagedFluxOptions.getTracerProvider().isEnabled()) { + pagedFluxOptions.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }); + } } From e2e8e9e468a8a8c038ea4c1cf70e33414395b775 Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Tue, 14 Apr 2020 15:25:12 -0400 Subject: [PATCH 04/40] reverting change for jaeger ui local testing --- sdk/cosmos/azure-cosmos/pom.xml | 24 +----------------------- 1 file changed, 1 insertion(+), 23 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/pom.xml b/sdk/cosmos/azure-cosmos/pom.xml index 29c9ba150e58..c36e195c1e86 100644 --- a/sdk/cosmos/azure-cosmos/pom.xml +++ b/sdk/cosmos/azure-cosmos/pom.xml @@ -82,26 +82,7 @@ Licensed under the MIT License. - - io.opentelemetry - opentelemetry-exporters-jaeger - 0.2.0 - - - io.grpc - grpc-okhttp - 1.23.0 - - - io.opentelemetry - opentelemetry-sdk - 0.2.0 - - - com.azure - azure-core-tracing-opentelemetry - 1.0.0-beta.4 - + commons-io commons-io @@ -352,9 +333,6 @@ Licensed under the MIT License. io.netty:netty-transport-native-epoll com.google.code.findbugs:jsr305 - io.opentelemetry:opentelemetry-sdk - io.opentelemetry:opentelemetry-exporters-jaeger - io.grpc:grpc-okhttp From c2fc40f3a35df2381104f88817209804491ae329 Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Wed, 15 Apr 2020 15:11:27 -0400 Subject: [PATCH 05/40] adding tracers in trigger udf and sproc --- .../com/azure/cosmos/CosmosAsyncClient.java | 65 +++-- .../azure/cosmos/CosmosAsyncContainer.java | 44 +-- .../com/azure/cosmos/CosmosAsyncDatabase.java | 28 +- .../com/azure/cosmos/CosmosAsyncScripts.java | 170 ++++++++++- .../cosmos/CosmosAsyncStoredProcedure.java | 275 +++++++++++++++--- .../com/azure/cosmos/CosmosAsyncTrigger.java | 172 ++++++++++- .../CosmosAsyncUserDefinedFunction.java | 167 ++++++++++- 7 files changed, 780 insertions(+), 141 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index fdad4fd4f0eb..00daf223c0ab 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -18,7 +18,6 @@ import com.azure.cosmos.models.CosmosPermissionProperties; import com.azure.cosmos.models.FeedOptions; import com.azure.cosmos.models.ModelBridgeInternal; -import com.azure.cosmos.implementation.Permission; import com.azure.cosmos.models.SqlQuerySpec; import com.azure.cosmos.util.CosmosPagedFlux; import com.azure.cosmos.util.UtilBridgeInternal; @@ -451,17 +450,7 @@ private Mono createDatabaseIfNotExistsInternal(Cosm put(TracerProvider.DB_STATEMENT, spanName); }}; - return database.read().doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(tracerProvider.startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).onErrorResume(exception -> { + return database.read().onErrorResume(exception -> { final Throwable unwrappedException = Exceptions.unwrap(exception); if (unwrappedException instanceof CosmosClientException) { final CosmosClientException cosmosClientException = (CosmosClientException) unwrappedException; @@ -471,6 +460,16 @@ private Mono createDatabaseIfNotExistsInternal(Cosm } } return Mono.error(unwrappedException); + }).doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(tracerProvider.startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } }).doOnSuccess(response -> { if (isTracingEnabled) { tracerProvider.endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); @@ -495,26 +494,28 @@ private Mono createDatabase(Database database, Cosm put(TracerProvider.DB_STATEMENT, spanName); }}; - return asyncDocumentClient.createDatabase(database, ModelBridgeInternal.toRequestOptions(options)).doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if(!callerFunc.isPresent()) { - parentContext.set(tracerProvider.startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + return asyncDocumentClient.createDatabase(database, ModelBridgeInternal.toRequestOptions(options)) + .map(databaseResourceResponse -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(databaseResourceResponse, + this)) + .single() + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(tracerProvider.startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - tracerProvider.endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - tracerProvider.endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }).map(databaseResourceResponse -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(databaseResourceResponse, - this)) - .single(); + }).doOnSuccess(response -> { + if (isTracingEnabled) { + tracerProvider.endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + tracerProvider.endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 2dbe77aeabcc..88d0114ec93e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -848,7 +848,8 @@ private Mono> deleteItem( return this.getDatabase() .getDocClientWrapper() .deleteDocument(getItemLink(itemId), requestOptions) - .doOnSubscribe(ignoredValue -> { + .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponseWithObjectType(response)) + .single().doOnSubscribe(ignoredValue -> { if (isTracingEnabled) { reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); @@ -860,14 +861,14 @@ private Mono> deleteItem( } }).doOnSuccess(response -> { if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), + response.getStatusCode()); } }).doOnError(throwable -> { if (isTracingEnabled) { database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); } - }).map(response -> ModelBridgeInternal.createCosmosAsyncItemResponseWithObjectType(response)) - .single(); + }); } private Mono> replaceItem( @@ -891,7 +892,8 @@ private Mono> replaceItem( return this.getDatabase() .getDocClientWrapper() .replaceDocument(getItemLink(itemId), doc, ModelBridgeInternal.toRequestOptions(options)) - .doOnSubscribe(ignoredValue -> { + .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) + .single().doOnSubscribe(ignoredValue -> { if (isTracingEnabled) { reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); @@ -903,14 +905,14 @@ private Mono> replaceItem( } }).doOnSuccess(response -> { if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), + response.getStatusCode()); } }).doOnError(throwable -> { if (isTracingEnabled) { database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); } - }).map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) - .single(); + }); } private Mono> upsertItem(T item, CosmosItemRequestOptions options, Context context) { @@ -932,7 +934,8 @@ private Mono> upsertItem(T item, CosmosItemReques .upsertDocument(this.getLink(), item, ModelBridgeInternal.toRequestOptions(options), true) - .doOnSubscribe(ignoredValue -> { + .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) + .single().doOnSubscribe(ignoredValue -> { if (isTracingEnabled) { reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); @@ -944,14 +947,14 @@ private Mono> upsertItem(T item, CosmosItemReques } }).doOnSuccess(response -> { if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), + response.getStatusCode()); } }).doOnError(throwable -> { if (isTracingEnabled) { database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); } - }).map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) - .single(); + }); } private Mono> readItem( @@ -972,7 +975,8 @@ private Mono> readItem( return this.getDatabase().getDocClientWrapper() .readDocument(getItemLink(itemId), requestOptions) - .doOnSubscribe(ignoredValue -> { + .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) + .single().doOnSubscribe(ignoredValue -> { if (isTracingEnabled) { reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); @@ -984,14 +988,14 @@ private Mono> readItem( } }).doOnSuccess(response -> { if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), + response.getStatusCode()); } }).doOnError(throwable -> { if (isTracingEnabled) { database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); } - }).map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) - .single(); + }); } private Mono read(CosmosContainerRequestOptions options, Context context) { @@ -1018,15 +1022,17 @@ private Mono read(CosmosContainerRequestOptions op context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); } } - }).doOnSuccess(response -> { + }).map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single() + .doOnSuccess(response -> { if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), + response.getStatusCode()); } }).doOnError(throwable -> { if (isTracingEnabled) { database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); } - }).map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); + }); } private Mono readProvisionedThroughput(Context context) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java index 5614a9cf8ae2..0e430db5cbce 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java @@ -779,15 +779,17 @@ private Mono createContainer( context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); } } - }).doOnSuccess(response -> { + }).map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, this)).single() + .doOnSuccess(response -> { if (isTracingEnabled) { - client.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + client.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), + response.getStatusCode()); } }).doOnError(throwable -> { if (isTracingEnabled) { client.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); } - }).map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, this)).single(); + }); } private Mono read(CosmosDatabaseRequestOptions options, Context context) { @@ -847,15 +849,17 @@ private Mono delete(CosmosDatabaseRequestOptions op context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); } } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - client.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - client.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }).map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single(); + }).map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single() + .doOnSuccess(response -> { + if (isTracingEnabled) { + client.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), + response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + client.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }); } private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Context context) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java index 0e4e0bedd5dc..55afad42f265 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java @@ -2,7 +2,11 @@ // Licensed under the MIT License. package com.azure.cosmos; +import com.azure.core.util.Context; +import com.azure.core.util.FluxUtil; +import com.azure.core.util.tracing.ProcessKind; import com.azure.cosmos.implementation.StoredProcedure; +import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.implementation.Trigger; import com.azure.cosmos.implementation.UserDefinedFunction; import com.azure.cosmos.models.CosmosAsyncStoredProcedureResponse; @@ -18,7 +22,15 @@ import com.azure.cosmos.util.CosmosPagedFlux; import com.azure.cosmos.util.UtilBridgeInternal; import reactor.core.publisher.Mono; +import reactor.core.publisher.Signal; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; + +import static com.azure.core.util.FluxUtil.withContext; import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; /** @@ -71,12 +83,18 @@ public Mono createStoredProcedure( StoredProcedure sProc = new StoredProcedure(); sProc.setId(properties.getId()); sProc.setBody(properties.getBody()); - return database.getDocClientWrapper() - .createStoredProcedure(container.getLink(), sProc, ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, this.container)) - .single(); + final CosmosStoredProcedureRequestOptions requestOptions = options; + return withContext(context -> createStoredProcedure(sProc, requestOptions, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } + + /** * Reads all cosmos stored procedures in a container. *

@@ -164,7 +182,6 @@ public CosmosAsyncStoredProcedure getStoredProcedure(String id) { return new CosmosAsyncStoredProcedure(id, this.container); } - /* UDF Operations */ /** @@ -184,9 +201,13 @@ public Mono createUserDefinedFunction( udf.setId(properties.getId()); udf.setBody(properties.getBody()); - return database.getDocClientWrapper() - .createUserDefinedFunction(container.getLink(), udf, null) - .map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, this.container)).single(); + return withContext(context -> createUserDefinedFunction(udf, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } /** @@ -290,12 +311,13 @@ public CosmosAsyncUserDefinedFunction getUserDefinedFunction(String id) { * @return an {@link Mono} containing the single resource response with the created trigger or an error. */ public Mono createTrigger(CosmosTriggerProperties properties) { - Trigger trigger = new Trigger(ModelBridgeInternal.toJsonFromJsonSerializable(properties)); - - return database.getDocClientWrapper() - .createTrigger(container.getLink(), trigger, null) - .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, this.container)) - .single(); + return withContext(context -> createTrigger(properties, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } /** @@ -381,4 +403,124 @@ public CosmosAsyncTrigger getTrigger(String id) { return new CosmosAsyncTrigger(id, this.container); } + private Mono createStoredProcedure(StoredProcedure sProc, + CosmosStoredProcedureRequestOptions options, + Context context) { + final boolean isTracingEnabled = this.container.getDatabase().getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "createStoredProcedure." + container.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, database.getId()); + put(TracerProvider.DB_URL, container.getDatabase().getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return database.getDocClientWrapper() + .createStoredProcedure(container.getLink(), sProc, ModelBridgeInternal.toRequestOptions(options)).doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(this.container.getDatabase().getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, this.container)) + .single().doOnSuccess(response -> { + if (isTracingEnabled) { + this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.error(throwable), 0); + } + }); + } + + private Mono createUserDefinedFunction( + UserDefinedFunction udf, + Context context) { + final boolean isTracingEnabled = this.container.getDatabase().getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "createUserDefinedFunction." + container.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, database.getId()); + put(TracerProvider.DB_URL, container.getDatabase().getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return database.getDocClientWrapper() + .createUserDefinedFunction(container.getLink(), udf, null).map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, + this.container)).single().doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(this.container.getDatabase().getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.error(throwable), 0); + } + }); + } + + private Mono createTrigger(CosmosTriggerProperties properties, Context context) { + final boolean isTracingEnabled = this.container.getDatabase().getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "createTrigger." + container.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, database.getId()); + put(TracerProvider.DB_URL, container.getDatabase().getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + Trigger trigger = new Trigger(ModelBridgeInternal.toJsonFromJsonSerializable(properties)); + return database.getDocClientWrapper() + .createTrigger(container.getLink(), trigger, null) + .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, this.container)) + .single() + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(this.container.getDatabase().getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.error(throwable), 0); + } + }); + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java index 9efa1778e3b8..2b84be5b9d8a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java @@ -2,13 +2,26 @@ // Licensed under the MIT License. package com.azure.cosmos; +import com.azure.core.util.Context; +import com.azure.core.util.FluxUtil; +import com.azure.core.util.tracing.ProcessKind; import com.azure.cosmos.implementation.Paths; import com.azure.cosmos.implementation.StoredProcedure; +import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.models.CosmosAsyncStoredProcedureResponse; import com.azure.cosmos.models.CosmosStoredProcedureProperties; import com.azure.cosmos.models.CosmosStoredProcedureRequestOptions; import com.azure.cosmos.models.ModelBridgeInternal; import reactor.core.publisher.Mono; +import reactor.core.publisher.Signal; + +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; + +import static com.azure.core.util.FluxUtil.withContext; /** * The type Cosmos async stored procedure. @@ -70,19 +83,20 @@ public Mono read() { * @return an {@link Mono} containing the single resource response with the read stored procedure or an error. */ public Mono read(CosmosStoredProcedureRequestOptions options) { - if (options == null) { - options = new CosmosStoredProcedureRequestOptions(); - } - return cosmosContainer.getDatabase().getDocClientWrapper().readStoredProcedure(getLink(), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer)).single(); + return withContext(context -> read(options, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } /** * Deletes a stored procedure by the stored procedure link. *

* After subscription the operation will be performed. - * The {@link Mono} upon successful completion will contain a single resource response for the deleted stored + * The {@link Mono} upon successful completion will contain a single resource response for the deleted stored * procedure. * In case of failure the {@link Mono} will error. * @@ -96,7 +110,7 @@ public Mono delete() { * Deletes a stored procedure by the stored procedure link. *

* After subscription the operation will be performed. - * The {@link Mono} upon successful completion will contain a single resource response for the deleted stored + * The {@link Mono} upon successful completion will contain a single resource response for the deleted stored * procedure. * In case of failure the {@link Mono} will error. * @@ -104,21 +118,20 @@ public Mono delete() { * @return an {@link Mono} containing the single resource response for the deleted stored procedure or an error. */ public Mono delete(CosmosStoredProcedureRequestOptions options) { - if (options == null) { - options = new CosmosStoredProcedureRequestOptions(); - } - return cosmosContainer.getDatabase() - .getDocClientWrapper() - .deleteStoredProcedure(getLink(), ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer)) - .single(); + return withContext(context -> delete(options, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } /** * Executes a stored procedure by the stored procedure link. *

* After subscription the operation will be performed. - * The {@link Mono} upon successful completion will contain a single resource response with the stored procedure + * The {@link Mono} upon successful completion will contain a single resource response with the stored procedure * response. * In case of failure the {@link Mono} will error. * @@ -128,21 +141,20 @@ public Mono delete(CosmosStoredProcedureRequ */ public Mono execute(Object[] procedureParams, CosmosStoredProcedureRequestOptions options) { - if (options == null) { - options = new CosmosStoredProcedureRequestOptions(); - } - return cosmosContainer.getDatabase() - .getDocClientWrapper() - .executeStoredProcedure(getLink(), ModelBridgeInternal.toRequestOptions(options), procedureParams) - .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer, this.id)) - .single(); + return withContext(context -> execute(procedureParams, options, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } /** * Replaces a stored procedure. *

* After subscription the operation will be performed. - * The {@link Mono} upon successful completion will contain a single resource response with the replaced stored + * The {@link Mono} upon successful completion will contain a single resource response with the replaced stored * procedure. * In case of failure the {@link Mono} will error. * @@ -157,7 +169,7 @@ public Mono replace(CosmosStoredProcedurePro * Replaces a stored procedure. *

* After subscription the operation will be performed. - * The {@link Mono} upon successful completion will contain a single resource response with the replaced stored + * The {@link Mono} upon successful completion will contain a single resource response with the replaced stored * procedure. * In case of failure the {@link Mono} will error. * @@ -167,15 +179,13 @@ public Mono replace(CosmosStoredProcedurePro */ public Mono replace(CosmosStoredProcedureProperties storedProcedureSettings, CosmosStoredProcedureRequestOptions options) { - if (options == null) { - options = new CosmosStoredProcedureRequestOptions(); - } - return cosmosContainer.getDatabase() - .getDocClientWrapper() - .replaceStoredProcedure(new StoredProcedure(ModelBridgeInternal.toJsonFromJsonSerializable(storedProcedureSettings)), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer)) - .single(); + return withContext(context -> replace(storedProcedureSettings, options, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } String getURIPathSegment() { @@ -195,4 +205,197 @@ String getLink() { builder.append(id()); return builder.toString(); } + + private Mono read(CosmosStoredProcedureRequestOptions options, + Context context) { + final boolean isTracingEnabled = this.cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "readStoredProcedure." + cosmosContainer.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, cosmosContainer.getDatabase().getId()); + put(TracerProvider.DB_URL, cosmosContainer.getDatabase().getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + if (options == null) { + options = new CosmosStoredProcedureRequestOptions(); + } + return cosmosContainer.getDatabase().getDocClientWrapper().readStoredProcedure(getLink(), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer)).single() + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(this.cosmosContainer.getDatabase().getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), + ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + this.cosmosContainer.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.complete(), + 200); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + this.cosmosContainer.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.error(throwable), 0); + } + }); + } + + private Mono delete(CosmosStoredProcedureRequestOptions options, + Context context) { + final boolean isTracingEnabled = this.cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "deleteStoredProcedure." + cosmosContainer.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, cosmosContainer.getDatabase().getId()); + put(TracerProvider.DB_URL, cosmosContainer.getDatabase().getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + if (options == null) { + options = new CosmosStoredProcedureRequestOptions(); + } + if (options == null) { + options = new CosmosStoredProcedureRequestOptions(); + } + return cosmosContainer.getDatabase() + .getDocClientWrapper() + .deleteStoredProcedure(getLink(), ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer)) + .single() + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(this.cosmosContainer.getDatabase().getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), + ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + this.cosmosContainer.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.complete(), + 200); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + this.cosmosContainer.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.error(throwable), 0); + } + }); + } + + private Mono execute(Object[] procedureParams, + CosmosStoredProcedureRequestOptions options, + Context context) { + final boolean isTracingEnabled = this.cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "executeStoredProcedure." + cosmosContainer.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, cosmosContainer.getDatabase().getId()); + put(TracerProvider.DB_URL, cosmosContainer.getDatabase().getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + if (options == null) { + options = new CosmosStoredProcedureRequestOptions(); + } + return cosmosContainer.getDatabase() + .getDocClientWrapper() + .executeStoredProcedure(getLink(), ModelBridgeInternal.toRequestOptions(options), procedureParams) + .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer, this.id)) + .single() + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(this.cosmosContainer.getDatabase().getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), + ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + this.cosmosContainer.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.complete(), + 200); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + this.cosmosContainer.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.error(throwable), 0); + } + }); + } + + private Mono replace(CosmosStoredProcedureProperties storedProcedureSettings, + CosmosStoredProcedureRequestOptions options, + Context context) { + final boolean isTracingEnabled = this.cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "replaceStoredProcedure." + cosmosContainer.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, cosmosContainer.getDatabase().getId()); + put(TracerProvider.DB_URL, cosmosContainer.getDatabase().getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + if (options == null) { + options = new CosmosStoredProcedureRequestOptions(); + } + return cosmosContainer.getDatabase() + .getDocClientWrapper() + .replaceStoredProcedure(new StoredProcedure(ModelBridgeInternal.toJsonFromJsonSerializable(storedProcedureSettings)), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer)) + .single() + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(this.cosmosContainer.getDatabase().getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), + ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + this.cosmosContainer.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.complete(), + 200); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + this.cosmosContainer.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.error(throwable), 0); + } + }); + } + } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java index 456b522f26aa..be0e6d4039eb 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java @@ -2,12 +2,25 @@ // Licensed under the MIT License. package com.azure.cosmos; +import com.azure.core.util.Context; +import com.azure.core.util.FluxUtil; +import com.azure.core.util.tracing.ProcessKind; import com.azure.cosmos.implementation.Paths; +import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.implementation.Trigger; import com.azure.cosmos.models.CosmosAsyncTriggerResponse; import com.azure.cosmos.models.CosmosTriggerProperties; import com.azure.cosmos.models.ModelBridgeInternal; import reactor.core.publisher.Mono; +import reactor.core.publisher.Signal; + +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; + +import static com.azure.core.util.FluxUtil.withContext; /** * The type Cosmos async trigger. This contains methods to operate on a cosmos trigger asynchronously @@ -52,14 +65,15 @@ CosmosAsyncTrigger setId(String id) { * @return an {@link Mono} containing the single resource response for the read cosmos trigger or an error. */ public Mono read() { - return container.getDatabase() - .getDocClientWrapper() - .readTrigger(getLink(), null) - .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, container)) - .single(); + return withContext(context -> read(context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } - /** * Replaces a cosmos trigger. *

@@ -71,11 +85,13 @@ public Mono read() { * @return an {@link Mono} containing the single resource response with the replaced cosmos trigger or an error. */ public Mono replace(CosmosTriggerProperties triggerSettings) { - return container.getDatabase() - .getDocClientWrapper() - .replaceTrigger(new Trigger(ModelBridgeInternal.toJsonFromJsonSerializable(triggerSettings)), null) - .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, container)) - .single(); + return withContext(context -> replace(triggerSettings, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } /** @@ -88,11 +104,13 @@ public Mono replace(CosmosTriggerProperties triggerS * @return an {@link Mono} containing the single resource response for the deleted cosmos trigger or an error. */ public Mono delete() { - return container.getDatabase() - .getDocClientWrapper() - .deleteTrigger(getLink(), null) - .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, container)) - .single(); + return withContext(context -> delete(context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } String getURIPathSegment() { @@ -113,4 +131,126 @@ String getLink() { return builder.toString(); } + private Mono read(Context context) { + final boolean isTracingEnabled = this.container.getDatabase().getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "readTrigger." + container.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, container.getDatabase().getId()); + put(TracerProvider.DB_URL, container.getDatabase().getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return container.getDatabase() + .getDocClientWrapper() + .readTrigger(getLink(), null) + .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, container)) + .single() + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(this.container.getDatabase().getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.error(throwable), 0); + } + }); + } + + private Mono replace(CosmosTriggerProperties triggerSettings, Context context) { + final boolean isTracingEnabled = this.container.getDatabase().getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "replaceTrigger." + container.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, container.getDatabase().getId()); + put(TracerProvider.DB_URL, container.getDatabase().getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return container.getDatabase() + .getDocClientWrapper() + .replaceTrigger(new Trigger(ModelBridgeInternal.toJsonFromJsonSerializable(triggerSettings)), null) + .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, container)) + .single() + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(this.container.getDatabase().getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.error(throwable), 0); + } + }); + } + + private Mono delete(Context context) { + final boolean isTracingEnabled = this.container.getDatabase().getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "deleteTrigger." + container.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, container.getDatabase().getId()); + put(TracerProvider.DB_URL, container.getDatabase().getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return container.getDatabase() + .getDocClientWrapper() + .deleteTrigger(getLink(), null) + .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, container)) + .single() + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(this.container.getDatabase().getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.error(throwable), 0); + } + }); + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java index 3823f1067a0b..27d517bd1d7e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java @@ -2,12 +2,25 @@ // Licensed under the MIT License. package com.azure.cosmos; +import com.azure.core.util.Context; +import com.azure.core.util.FluxUtil; +import com.azure.core.util.tracing.ProcessKind; import com.azure.cosmos.implementation.Paths; +import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.implementation.UserDefinedFunction; import com.azure.cosmos.models.CosmosAsyncUserDefinedFunctionResponse; import com.azure.cosmos.models.CosmosUserDefinedFunctionProperties; import com.azure.cosmos.models.ModelBridgeInternal; import reactor.core.publisher.Mono; +import reactor.core.publisher.Signal; + +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; + +import static com.azure.core.util.FluxUtil.withContext; /** * The type Cosmos async user defined function. @@ -54,8 +67,13 @@ CosmosAsyncUserDefinedFunction setId(String id) { * @return an {@link Mono} containing the single resource response for the read user defined function or an error. */ public Mono read() { - return container.getDatabase().getDocClientWrapper().readUserDefinedFunction(getLink(), null) - .map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, container)).single(); + return withContext(context -> read(context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } /** @@ -71,11 +89,13 @@ public Mono read() { * or an error. */ public Mono replace(CosmosUserDefinedFunctionProperties udfSettings) { - return container.getDatabase() - .getDocClientWrapper() - .replaceUserDefinedFunction(new UserDefinedFunction(ModelBridgeInternal.toJsonFromJsonSerializable(udfSettings)), null) - .map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, container)) - .single(); + return withContext(context -> replace(udfSettings, context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } /** @@ -90,11 +110,13 @@ public Mono replace(CosmosUserDefinedFun * an error. */ public Mono delete() { - return container.getDatabase() - .getDocClientWrapper() - .deleteUserDefinedFunction(this.getLink(), null) - .map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, container)) - .single(); + return withContext(context -> delete(context)).subscriberContext(reactorContext -> { + Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); + if (master.isPresent()) { + reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); + } + return reactorContext.put(TracerProvider.MASTER_CALL, true); + }); } String getURIPathSegment() { @@ -114,4 +136,125 @@ String getLink() { builder.append(getId()); return builder.toString(); } + + private Mono read(Context context) { + final boolean isTracingEnabled = this.container.getDatabase().getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "readUDF." + container.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, container.getDatabase().getId()); + put(TracerProvider.DB_URL, container.getDatabase().getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return container.getDatabase().getDocClientWrapper().readUserDefinedFunction(getLink(), null) + .map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, container)).single() + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(this.container.getDatabase().getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.error(throwable), 0); + } + }); + } + + private Mono replace(CosmosUserDefinedFunctionProperties udfSettings, + Context context) { + final boolean isTracingEnabled = this.container.getDatabase().getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "replaceUDF." + container.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, container.getDatabase().getId()); + put(TracerProvider.DB_URL, container.getDatabase().getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return container.getDatabase() + .getDocClientWrapper() + .replaceUserDefinedFunction(new UserDefinedFunction(ModelBridgeInternal.toJsonFromJsonSerializable(udfSettings)), null) + .map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, container)) + .single() + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(this.container.getDatabase().getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.error(throwable), 0); + } + }); + } + + private Mono delete(Context context) { + final boolean isTracingEnabled = this.container.getDatabase().getClient().getTracerProvider().isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + String spanName = "deleteUDF." + container.getId(); + Map tracingAttributes = new HashMap() {{ + put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); + put(TracerProvider.DB_INSTANCE, container.getDatabase().getId()); + put(TracerProvider.DB_URL, container.getDatabase().getClient().getServiceEndpoint()); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + + return container.getDatabase() + .getDocClientWrapper() + .deleteUserDefinedFunction(this.getLink(), null) + .map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, container)) + .single() + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(this.container.getDatabase().getClient().getTracerProvider().startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), + Signal.error(throwable), 0); + } + }); + } } From e614db5d38ff2f4947900fa0174774f8c13f37fe Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Tue, 21 Apr 2020 17:13:25 -0400 Subject: [PATCH 06/40] adding test case and refactoring code --- .../com/azure/cosmos/CosmosAsyncClient.java | 61 +--- .../azure/cosmos/CosmosAsyncContainer.java | 299 +++--------------- .../com/azure/cosmos/CosmosAsyncDatabase.java | 188 ++--------- .../com/azure/cosmos/CosmosAsyncScripts.java | 158 +++------ .../cosmos/CosmosAsyncStoredProcedure.java | 139 +------- .../com/azure/cosmos/CosmosAsyncTrigger.java | 95 +----- .../CosmosAsyncUserDefinedFunction.java | 91 +----- .../cosmos/implementation/TracerProvider.java | 110 ++++++- .../src/main/java/module-info.java | 1 + .../com/azure/cosmos/CosmosTracerTest.java | 265 ++++++++++++++++ .../directconnectivity/ReflectionUtils.java | 5 + 11 files changed, 527 insertions(+), 885 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index 00daf223c0ab..ccb4a72386ec 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -4,8 +4,6 @@ import com.azure.core.annotation.ServiceClient; import com.azure.core.util.Context; -import com.azure.core.util.FluxUtil; -import com.azure.core.util.tracing.ProcessKind; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.Configs; import com.azure.cosmos.implementation.Database; @@ -24,15 +22,12 @@ import io.micrometer.core.instrument.MeterRegistry; import reactor.core.Exceptions; import reactor.core.publisher.Mono; -import reactor.core.publisher.Signal; import java.io.Closeable; import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.Objects; import java.util.Optional; -import java.util.concurrent.atomic.AtomicReference; import static com.azure.core.util.FluxUtil.withContext; import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; @@ -438,19 +433,14 @@ private Mono createDatabaseIfNotExistsInternal(Cosm private Mono createDatabaseIfNotExistsInternal(CosmosAsyncDatabase database, Context context) { - final boolean isTracingEnabled = tracerProvider.isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "createDatabaseIfNotExistsInternal." + database.getId(); - Map tracingAttributes = new HashMap() {{ + Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); put(TracerProvider.DB_INSTANCE, database.getId()); put(TracerProvider.DB_URL, serviceEndpoint); put(TracerProvider.DB_STATEMENT, spanName); }}; - - return database.read().onErrorResume(exception -> { + Mono responseMono = database.read().onErrorResume(exception -> { final Throwable unwrappedException = Exceptions.unwrap(exception); if (unwrappedException instanceof CosmosClientException) { final CosmosClientException cosmosClientException = (CosmosClientException) unwrappedException; @@ -460,32 +450,11 @@ private Mono createDatabaseIfNotExistsInternal(Cosm } } return Mono.error(unwrappedException); - }).doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(tracerProvider.startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - tracerProvider.endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - tracerProvider.endSpan(parentContext.get(), Signal.error(throwable), 0); - } }); + return tracerProvider.traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); } private Mono createDatabase(Database database, CosmosDatabaseRequestOptions options, Context context){ - final boolean isTracingEnabled = tracerProvider.isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "createDatabase." + database.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -494,28 +463,10 @@ private Mono createDatabase(Database database, Cosm put(TracerProvider.DB_STATEMENT, spanName); }}; - return asyncDocumentClient.createDatabase(database, ModelBridgeInternal.toRequestOptions(options)) + Mono responseMono = asyncDocumentClient.createDatabase(database, ModelBridgeInternal.toRequestOptions(options)) .map(databaseResourceResponse -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(databaseResourceResponse, this)) - .single() - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(tracerProvider.startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - tracerProvider.endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - tracerProvider.endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }); + .single(); + return tracerProvider.traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 88d0114ec93e..929857fbfeb0 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -3,8 +3,6 @@ package com.azure.cosmos; import com.azure.core.util.Context; -import com.azure.core.util.FluxUtil; -import com.azure.core.util.tracing.ProcessKind; import com.azure.cosmos.implementation.CosmosItemProperties; import com.azure.cosmos.implementation.Document; import com.azure.cosmos.implementation.HttpConstants; @@ -27,13 +25,10 @@ import com.azure.cosmos.util.UtilBridgeInternal; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; -import reactor.core.publisher.Signal; import java.util.HashMap; import java.util.Map; -import java.util.Objects; import java.util.Optional; -import java.util.concurrent.atomic.AtomicReference; import java.util.stream.Collectors; import static com.azure.core.util.FluxUtil.withContext; @@ -132,10 +127,6 @@ public Mono delete(CosmosContainerRequestOptions o } private Mono delete(CosmosContainerRequestOptions options, Context context) { - final boolean isTracingEnabled = database.getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "deleteContainer." + this.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -144,26 +135,11 @@ private Mono delete(CosmosContainerRequestOptions put(TracerProvider.DB_STATEMENT, spanName); }}; - return database.getDocClientWrapper().deleteCollection(getLink(), ModelBridgeInternal.toRequestOptions(options)) - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(database.getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }).map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); + Mono responseMono = database.getDocClientWrapper().deleteCollection(getLink(), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); + return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, + tracingAttributes, context, spanName); } /** @@ -230,10 +206,6 @@ public Mono replace( private Mono replace(CosmosContainerProperties containerProperties, CosmosContainerRequestOptions options, Context context) { - final boolean isTracingEnabled = database.getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "replaceContainer." + this.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -242,28 +214,12 @@ private Mono replace(CosmosContainerProperties con put(TracerProvider.DB_STATEMENT, spanName); }}; - return database.getDocClientWrapper() - .replaceCollection(ModelBridgeInternal.getV2Collection(containerProperties), ModelBridgeInternal.toRequestOptions(options)) - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - - parentContext.set(database.getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }).map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); + Mono responseMono = database.getDocClientWrapper() + .replaceCollection(ModelBridgeInternal.getV2Collection(containerProperties), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); + return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, + tracingAttributes, context, spanName); } /* CosmosAsyncItem operations */ @@ -332,10 +288,6 @@ public Mono> createItem(T item, CosmosItemRequest } private Mono> createItem(T item, CosmosItemRequestOptions options, Context context) { - final boolean isTracingEnabled = database.getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "createItem." + this.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -347,31 +299,14 @@ private Mono> createItem(T item, CosmosItemReques @SuppressWarnings("unchecked") Class itemType = (Class) item.getClass(); RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); - return database.getDocClientWrapper() + Mono> responseMono = database.getDocClientWrapper() .createDocument(getLink(), item, requestOptions, true) - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(database.getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }).map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) + .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) .single(); + return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, tracingAttributes, context,spanName ); } /** @@ -833,10 +768,6 @@ private Mono> deleteItem( String itemId, RequestOptions requestOptions, Context context) { - final boolean isTracingEnabled = database.getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "deleteItem." + this.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -845,30 +776,14 @@ private Mono> deleteItem( put(TracerProvider.DB_STATEMENT, spanName); }}; - return this.getDatabase() + Mono> responseMono = this.getDatabase() .getDocClientWrapper() .deleteDocument(getItemLink(itemId), requestOptions) .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponseWithObjectType(response)) - .single().doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(database.getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), - response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }); + .single(); + return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, + tracingAttributes, + context, spanName); } private Mono> replaceItem( @@ -877,10 +792,6 @@ private Mono> replaceItem( Document doc, CosmosItemRequestOptions options, Context context) { - final boolean isTracingEnabled = database.getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "replaceItem." + this.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -889,37 +800,17 @@ private Mono> replaceItem( put(TracerProvider.DB_STATEMENT, spanName); }}; - return this.getDatabase() + Mono> responseMono = this.getDatabase() .getDocClientWrapper() .replaceDocument(getItemLink(itemId), doc, ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) - .single().doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(database.getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), - response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }); + .single(); + return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, + tracingAttributes, + context, spanName); } private Mono> upsertItem(T item, CosmosItemRequestOptions options, Context context) { - final boolean isTracingEnabled = database.getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "upsertItem." + this.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -930,41 +821,20 @@ private Mono> upsertItem(T item, CosmosItemReques @SuppressWarnings("unchecked") Class itemType = (Class) item.getClass(); - return this.getDatabase().getDocClientWrapper() + Mono> responseMono = this.getDatabase().getDocClientWrapper() .upsertDocument(this.getLink(), item, ModelBridgeInternal.toRequestOptions(options), true) .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) - .single().doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(database.getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), - response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }); + .single(); + return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, tracingAttributes, + context, spanName); } private Mono> readItem( String itemId, PartitionKey partitionKey, RequestOptions requestOptions, Class itemType, Context context) { - final boolean isTracingEnabled = database.getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "readItem." + this.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -973,36 +843,15 @@ private Mono> readItem( put(TracerProvider.DB_STATEMENT, spanName); }}; - return this.getDatabase().getDocClientWrapper() + Mono> responseMono = this.getDatabase().getDocClientWrapper() .readDocument(getItemLink(itemId), requestOptions) .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) - .single().doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(database.getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), - response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }); + .single(); + return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, tracingAttributes, + context, spanName); } private Mono read(CosmosContainerRequestOptions options, Context context) { - final boolean isTracingEnabled = database.getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "readContainer." + this.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -1011,35 +860,15 @@ private Mono read(CosmosContainerRequestOptions op put(TracerProvider.DB_STATEMENT, spanName); }}; - return database.getDocClientWrapper().readCollection(getLink(), ModelBridgeInternal.toRequestOptions(options)) - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(database.getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single() - .doOnSuccess(response -> { - if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), - response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - database.getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }); + Mono responseMono = database.getDocClientWrapper().readCollection(getLink(), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); + return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, + tracingAttributes, + context, spanName); } private Mono readProvisionedThroughput(Context context) { - final boolean isTracingEnabled = this.getDatabase().getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "readProvisionedThroughput." + this.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -1048,7 +877,7 @@ private Mono readProvisionedThroughput(Context context) { put(TracerProvider.DB_STATEMENT, spanName); }}; - return this.read() + Mono responseMono = this.read() .flatMap(cosmosContainerResponse -> database.getDocClientWrapper() .queryOffers("select * from c where c.offerResourceId = '" @@ -1064,32 +893,12 @@ private Mono readProvisionedThroughput(Context context) { return database.getDocClientWrapper() .readOffer(offerFeedResponse.getResults().get(0).getSelfLink()) .single(); - }).map(cosmosOfferResponse -> cosmosOfferResponse.getResource().getThroughput()).doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if(!callerFunc.isPresent()) { - parentContext.set(getDatabase().getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), 200); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }); + }).map(cosmosOfferResponse -> cosmosOfferResponse.getResource().getThroughput()); + return this.getDatabase().getClient().getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono + , tracingAttributes, context, spanName); } private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Context context) { - final boolean isTracingEnabled = this.getDatabase().getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "replaceProvisionedThroughput." + this.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -1098,7 +907,7 @@ private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Co put(TracerProvider.DB_STATEMENT, spanName); }}; - return this.read() + Mono responseMono = this.read() .flatMap(cosmosContainerResponse -> database.getDocClientWrapper() .queryOffers("select * from c where c.offerResourceId = '" @@ -1114,24 +923,8 @@ private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Co Offer offer = offerFeedResponse.getResults().get(0); offer.setThroughput(requestUnitsPerSecond); return database.getDocClientWrapper().replaceOffer(offer).single(); - }).map(offerResourceResponse -> offerResourceResponse.getResource().getThroughput()).doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if(!callerFunc.isPresent()) { - parentContext.set(getDatabase().getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), Signal.complete(), 200); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }); + }).map(offerResourceResponse -> offerResourceResponse.getResource().getThroughput()); + return this.getDatabase().getClient().getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono + , tracingAttributes, context, spanName); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java index 0e430db5cbce..2df8033ba2e2 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java @@ -3,8 +3,6 @@ package com.azure.cosmos; import com.azure.core.util.Context; -import com.azure.core.util.FluxUtil; -import com.azure.core.util.tracing.ProcessKind; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.Offer; @@ -25,13 +23,10 @@ import com.azure.cosmos.util.UtilBridgeInternal; import reactor.core.Exceptions; import reactor.core.publisher.Mono; -import reactor.core.publisher.Signal; import java.util.HashMap; import java.util.Map; -import java.util.Objects; import java.util.Optional; -import java.util.concurrent.atomic.AtomicReference; import static com.azure.core.util.FluxUtil.withContext; import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; @@ -206,14 +201,7 @@ public Mono createContainer( } final CosmosContainerRequestOptions requestOptions = options; - return withContext(context -> - { - Optional parentSpan = context.getData("user-span-name"); - if(parentSpan.isPresent()) { - System.out.println("here"); - } - return createContainer(containerProperties, requestOptions, context); - }).subscriberContext(reactorContext -> { + return withContext(context -> createContainer(containerProperties, requestOptions, context)).subscriberContext(reactorContext -> { Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); if (master.isPresent()) { reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); @@ -708,10 +696,6 @@ private Mono createContainerIfNotExistsInternal( CosmosAsyncContainer container, CosmosContainerRequestOptions options, Context context) { - final boolean isTracingEnabled = client.getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "createContainerIfNotExistsInternal." + containerProperties.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -720,7 +704,7 @@ private Mono createContainerIfNotExistsInternal( put(TracerProvider.DB_STATEMENT, spanName); }}; - return container.read(options).onErrorResume(exception -> { + Mono responseMono = container.read(options).onErrorResume(exception -> { final Throwable unwrappedException = Exceptions.unwrap(exception); if (unwrappedException instanceof CosmosClientException) { final CosmosClientException cosmosClientException = (CosmosClientException) unwrappedException; @@ -729,35 +713,15 @@ private Mono createContainerIfNotExistsInternal( } } return Mono.error(unwrappedException); - }).doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(client.getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - client.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - client.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); - } }); + return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, + spanName); } private Mono createContainer( CosmosContainerProperties containerProperties, CosmosContainerRequestOptions options, Context context) { - final boolean isTracingEnabled = client.getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "createContainer." + containerProperties.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -766,37 +730,15 @@ private Mono createContainer( put(TracerProvider.DB_STATEMENT, spanName); }}; - return getDocClientWrapper() + Mono responseMono = getDocClientWrapper() .createCollection(this.getLink(), ModelBridgeInternal.getV2Collection(containerProperties), ModelBridgeInternal.toRequestOptions(options)) - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(client.getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, this)).single() - .doOnSuccess(response -> { - if (isTracingEnabled) { - client.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), - response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - client.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }); + .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, this)).single(); + return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, + spanName); } private Mono read(CosmosDatabaseRequestOptions options, Context context) { - final boolean isTracingEnabled = client.getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "readDatabase." + this.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -805,32 +747,14 @@ private Mono read(CosmosDatabaseRequestOptions opti put(TracerProvider.DB_STATEMENT, spanName); }}; - return getDocClientWrapper().readDatabase(getLink(), ModelBridgeInternal.toRequestOptions(options)).doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(client.getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - client.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - client.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }).map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single(); + Mono responseMono = getDocClientWrapper().readDatabase(getLink(), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single(); + return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, + spanName); } private Mono delete(CosmosDatabaseRequestOptions options, Context context) { - final boolean isTracingEnabled = client.getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "deleteDatabase." + this.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -839,34 +763,14 @@ private Mono delete(CosmosDatabaseRequestOptions op put(TracerProvider.DB_STATEMENT, spanName); }}; - return getDocClientWrapper().deleteDatabase(getLink(), ModelBridgeInternal.toRequestOptions(options)).doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(client.getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single() - .doOnSuccess(response -> { - if (isTracingEnabled) { - client.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), - response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - client.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }); + Mono responseMono = getDocClientWrapper().deleteDatabase(getLink(), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single(); + return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, + spanName); } private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Context context) { - final boolean isTracingEnabled = this.client.getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "replaceOffer." + this.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -875,7 +779,7 @@ private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Co put(TracerProvider.DB_STATEMENT, spanName); }}; - return this.read() + Mono responseMono = this.read() .flatMap(cosmosDatabaseResponse -> this.getDocClientWrapper() .queryOffers("select * from c where c.offerResourceId = '" + cosmosDatabaseResponse.getProperties() @@ -895,34 +799,12 @@ private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Co .single(); }).map(offerResourceResponse -> offerResourceResponse .getResource() - .getThroughput())) - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(this.client.getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - this.client.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), - 200); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - this.client.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }); + .getThroughput())); + return this.client.getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono, tracingAttributes + , context, spanName); } private Mono readProvisionedThroughput(Context context) { - final boolean isTracingEnabled = this.client.getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "readProvisionedThroughput." + this.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -931,7 +813,7 @@ private Mono readProvisionedThroughput(Context context) { put(TracerProvider.DB_STATEMENT, spanName); }}; - return this.read() + Mono responseMono = this.read() .flatMap(cosmosDatabaseResponse -> getDocClientWrapper() .queryOffers("select * from c where c.offerResourceId = '" + cosmosDatabaseResponse @@ -953,26 +835,8 @@ private Mono readProvisionedThroughput(Context context) { .single(); }).map(cosmosContainerResponse1 -> cosmosContainerResponse1 .getResource() - .getThroughput())).doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(this.client.getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), - ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - this.client.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), - 200); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - this.client.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }); + .getThroughput())); + return this.client.getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono, tracingAttributes + , context, spanName); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java index 55afad42f265..00b1578e18b4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java @@ -3,8 +3,6 @@ package com.azure.cosmos; import com.azure.core.util.Context; -import com.azure.core.util.FluxUtil; -import com.azure.core.util.tracing.ProcessKind; import com.azure.cosmos.implementation.StoredProcedure; import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.implementation.Trigger; @@ -22,13 +20,10 @@ import com.azure.cosmos.util.CosmosPagedFlux; import com.azure.cosmos.util.UtilBridgeInternal; import reactor.core.publisher.Mono; -import reactor.core.publisher.Signal; import java.util.HashMap; import java.util.Map; -import java.util.Objects; import java.util.Optional; -import java.util.concurrent.atomic.AtomicReference; import static com.azure.core.util.FluxUtil.withContext; import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; @@ -71,7 +66,7 @@ public Mono createStoredProcedure(CosmosStor * In case of failure the {@link Mono} will error. * * @param properties the cosmos stored procedure properties. - * @param options the stored procedure request options. + * @param options the stored procedure request options. * @return an {@link Mono} containing the single cosmos stored procedure resource response or an error. */ public Mono createStoredProcedure( @@ -94,7 +89,6 @@ public Mono createStoredProcedure( } - /** * Reads all cosmos stored procedures in a container. *

@@ -116,10 +110,10 @@ public CosmosPagedFlux readAllStoredProcedures( this.container.getDatabase().getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper() - .readStoredProcedures(container.getLink(), options) - .map(response -> BridgeInternal.createFeedResponse( - ModelBridgeInternal.getCosmosStoredProcedurePropertiesFromV2Results(response.getResults()), - response.getResponseHeaders())); + .readStoredProcedures(container.getLink(), options) + .map(response -> BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosStoredProcedurePropertiesFromV2Results(response.getResults()), + response.getResponseHeaders())); }); } @@ -130,7 +124,7 @@ public CosmosPagedFlux readAllStoredProcedures( * The {@link CosmosPagedFlux} will contain one or several feed response pages of the obtained stored procedures. * In case of failure the {@link CosmosPagedFlux} will error. * - * @param query the the query. + * @param query the the query. * @param options the feed options. * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the obtained stored * procedures or @@ -150,7 +144,7 @@ public CosmosPagedFlux queryStoredProcedures( * In case of failure the {@link CosmosPagedFlux} will error. * * @param querySpec the SQL query specification. - * @param options the feed options. + * @param options the feed options. * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the obtained stored * procedures or * an error. @@ -165,10 +159,10 @@ public CosmosPagedFlux queryStoredProcedures( this.container.getDatabase().getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper() - .queryStoredProcedures(container.getLink(), querySpec, options) - .map(response -> BridgeInternal.createFeedResponse( - ModelBridgeInternal.getCosmosStoredProcedurePropertiesFromV2Results(response.getResults()), - response.getResponseHeaders())); + .queryStoredProcedures(container.getLink(), querySpec, options) + .map(response -> BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosStoredProcedurePropertiesFromV2Results(response.getResults()), + response.getResponseHeaders())); }); } @@ -230,10 +224,10 @@ public CosmosPagedFlux readAllUserDefinedFu this.container.getDatabase().getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper() - .readUserDefinedFunctions(container.getLink(), options) - .map(response -> BridgeInternal.createFeedResponse( - ModelBridgeInternal.getCosmosUserDefinedFunctionPropertiesFromV2Results(response.getResults()), - response.getResponseHeaders())); + .readUserDefinedFunctions(container.getLink(), options) + .map(response -> BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosUserDefinedFunctionPropertiesFromV2Results(response.getResults()), + response.getResponseHeaders())); }); } @@ -245,7 +239,7 @@ public CosmosPagedFlux readAllUserDefinedFu * functions. * In case of failure the {@link CosmosPagedFlux} will error. * - * @param query the query. + * @param query the query. * @param options the feed options. * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the obtained user defined * functions @@ -266,7 +260,7 @@ public CosmosPagedFlux queryUserDefinedFunc * In case of failure the {@link CosmosPagedFlux} will error. * * @param querySpec the SQL query specification. - * @param options the feed options. + * @param options the feed options. * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the obtained user defined * functions * or an error. @@ -281,10 +275,10 @@ public CosmosPagedFlux queryUserDefinedFunc this.container.getDatabase().getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper() - .queryUserDefinedFunctions(container.getLink(), querySpec, options) - .map(response -> BridgeInternal.createFeedResponse( - ModelBridgeInternal.getCosmosUserDefinedFunctionPropertiesFromV2Results(response.getResults()), - response.getResponseHeaders())); + .queryUserDefinedFunctions(container.getLink(), querySpec, options) + .map(response -> BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosUserDefinedFunctionPropertiesFromV2Results(response.getResults()), + response.getResponseHeaders())); }); } @@ -341,10 +335,10 @@ public CosmosPagedFlux readAllTriggers(FeedOptions opti this.container.getDatabase().getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper() - .readTriggers(container.getLink(), options) - .map(response -> BridgeInternal.createFeedResponse( - ModelBridgeInternal.getCosmosTriggerPropertiesFromV2Results(response.getResults()), - response.getResponseHeaders())); + .readTriggers(container.getLink(), options) + .map(response -> BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosTriggerPropertiesFromV2Results(response.getResults()), + response.getResponseHeaders())); }); } @@ -355,7 +349,7 @@ public CosmosPagedFlux readAllTriggers(FeedOptions opti * The {@link CosmosPagedFlux} will contain one or several feed response pages of the obtained triggers. * In case of failure the {@link CosmosPagedFlux} will error. * - * @param query the query. + * @param query the query. * @param options the feed options. * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the obtained triggers or an * error. @@ -372,7 +366,7 @@ public CosmosPagedFlux queryTriggers(String query, Feed * In case of failure the {@link CosmosPagedFlux} will error. * * @param querySpec the SQL query specification. - * @param options the feed options. + * @param options the feed options. * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the obtained triggers or an * error. */ @@ -386,10 +380,10 @@ public CosmosPagedFlux queryTriggers( this.container.getDatabase().getClient().getServiceEndpoint()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper() - .queryTriggers(container.getLink(), querySpec, options) - .map(response -> BridgeInternal.createFeedResponse( - ModelBridgeInternal.getCosmosTriggerPropertiesFromV2Results(response.getResults()), - response.getResponseHeaders())); + .queryTriggers(container.getLink(), querySpec, options) + .map(response -> BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosTriggerPropertiesFromV2Results(response.getResults()), + response.getResponseHeaders())); }); } @@ -406,10 +400,6 @@ public CosmosAsyncTrigger getTrigger(String id) { private Mono createStoredProcedure(StoredProcedure sProc, CosmosStoredProcedureRequestOptions options, Context context) { - final boolean isTracingEnabled = this.container.getDatabase().getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "createStoredProcedure." + container.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -418,38 +408,15 @@ private Mono createStoredProcedure(StoredPro put(TracerProvider.DB_STATEMENT, spanName); }}; - return database.getDocClientWrapper() - .createStoredProcedure(container.getLink(), sProc, ModelBridgeInternal.toRequestOptions(options)).doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(this.container.getDatabase().getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, this.container)) - .single().doOnSuccess(response -> { - if (isTracingEnabled) { - this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.complete(), response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.error(throwable), 0); - } - }); + Mono responseMono = database.getDocClientWrapper() + .createStoredProcedure(container.getLink(), sProc, ModelBridgeInternal.toRequestOptions(options)).map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, this.container)) + .single(); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); } private Mono createUserDefinedFunction( UserDefinedFunction udf, Context context) { - final boolean isTracingEnabled = this.container.getDatabase().getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "createUserDefinedFunction." + container.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -458,36 +425,13 @@ private Mono createUserDefinedFunction( put(TracerProvider.DB_STATEMENT, spanName); }}; - return database.getDocClientWrapper() + Mono responseMono = database.getDocClientWrapper() .createUserDefinedFunction(container.getLink(), udf, null).map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, - this.container)).single().doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(this.container.getDatabase().getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.complete(), response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.error(throwable), 0); - } - }); + this.container)).single(); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); } private Mono createTrigger(CosmosTriggerProperties properties, Context context) { - final boolean isTracingEnabled = this.container.getDatabase().getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "createTrigger." + container.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -497,30 +441,10 @@ private Mono createTrigger(CosmosTriggerProperties p }}; Trigger trigger = new Trigger(ModelBridgeInternal.toJsonFromJsonSerializable(properties)); - return database.getDocClientWrapper() + Mono responseMono = database.getDocClientWrapper() .createTrigger(container.getLink(), trigger, null) .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, this.container)) - .single() - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(this.container.getDatabase().getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.complete(), response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.error(throwable), 0); - } - }); + .single(); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java index 2b84be5b9d8a..102445d05626 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java @@ -3,8 +3,6 @@ package com.azure.cosmos; import com.azure.core.util.Context; -import com.azure.core.util.FluxUtil; -import com.azure.core.util.tracing.ProcessKind; import com.azure.cosmos.implementation.Paths; import com.azure.cosmos.implementation.StoredProcedure; import com.azure.cosmos.implementation.TracerProvider; @@ -13,13 +11,10 @@ import com.azure.cosmos.models.CosmosStoredProcedureRequestOptions; import com.azure.cosmos.models.ModelBridgeInternal; import reactor.core.publisher.Mono; -import reactor.core.publisher.Signal; import java.util.HashMap; import java.util.Map; -import java.util.Objects; import java.util.Optional; -import java.util.concurrent.atomic.AtomicReference; import static com.azure.core.util.FluxUtil.withContext; @@ -207,11 +202,7 @@ String getLink() { } private Mono read(CosmosStoredProcedureRequestOptions options, - Context context) { - final boolean isTracingEnabled = this.cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; + Context context) { String spanName = "readStoredProcedure." + cosmosContainer.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -223,40 +214,15 @@ private Mono read(CosmosStoredProcedureReque if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } - return cosmosContainer.getDatabase().getDocClientWrapper().readStoredProcedure(getLink(), + Mono responseMono = + cosmosContainer.getDatabase().getDocClientWrapper().readStoredProcedure(getLink(), ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer)).single() - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(this.cosmosContainer.getDatabase().getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), - ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - this.cosmosContainer.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.complete(), - 200); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - this.cosmosContainer.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.error(throwable), 0); - } - }); + .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer)).single(); + return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); } private Mono delete(CosmosStoredProcedureRequestOptions options, Context context) { - final boolean isTracingEnabled = this.cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "deleteStoredProcedure." + cosmosContainer.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -271,43 +237,17 @@ private Mono delete(CosmosStoredProcedureReq if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } - return cosmosContainer.getDatabase() + Mono responseMono = cosmosContainer.getDatabase() .getDocClientWrapper() .deleteStoredProcedure(getLink(), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer)) - .single() - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(this.cosmosContainer.getDatabase().getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), - ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - this.cosmosContainer.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.complete(), - 200); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - this.cosmosContainer.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.error(throwable), 0); - } - }); + .single(); + return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); } private Mono execute(Object[] procedureParams, CosmosStoredProcedureRequestOptions options, Context context) { - final boolean isTracingEnabled = this.cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "executeStoredProcedure." + cosmosContainer.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -319,43 +259,18 @@ private Mono execute(Object[] procedureParam if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } - return cosmosContainer.getDatabase() + Mono responseMono = cosmosContainer.getDatabase() .getDocClientWrapper() .executeStoredProcedure(getLink(), ModelBridgeInternal.toRequestOptions(options), procedureParams) - .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer, this.id)) - .single() - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(this.cosmosContainer.getDatabase().getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), - ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - this.cosmosContainer.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.complete(), - 200); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - this.cosmosContainer.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.error(throwable), 0); - } - }); + .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer, + this.id)) + .single(); + return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); } private Mono replace(CosmosStoredProcedureProperties storedProcedureSettings, CosmosStoredProcedureRequestOptions options, Context context) { - final boolean isTracingEnabled = this.cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "replaceStoredProcedure." + cosmosContainer.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -367,35 +282,13 @@ private Mono replace(CosmosStoredProcedurePr if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } - return cosmosContainer.getDatabase() + Mono responseMono = cosmosContainer.getDatabase() .getDocClientWrapper() .replaceStoredProcedure(new StoredProcedure(ModelBridgeInternal.toJsonFromJsonSerializable(storedProcedureSettings)), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer)) - .single() - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(this.cosmosContainer.getDatabase().getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), - ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - this.cosmosContainer.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.complete(), - 200); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - this.cosmosContainer.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.error(throwable), 0); - } - }); + .single(); + return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java index be0e6d4039eb..498d7b5d8588 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java @@ -3,8 +3,6 @@ package com.azure.cosmos; import com.azure.core.util.Context; -import com.azure.core.util.FluxUtil; -import com.azure.core.util.tracing.ProcessKind; import com.azure.cosmos.implementation.Paths; import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.implementation.Trigger; @@ -12,13 +10,10 @@ import com.azure.cosmos.models.CosmosTriggerProperties; import com.azure.cosmos.models.ModelBridgeInternal; import reactor.core.publisher.Mono; -import reactor.core.publisher.Signal; import java.util.HashMap; import java.util.Map; -import java.util.Objects; import java.util.Optional; -import java.util.concurrent.atomic.AtomicReference; import static com.azure.core.util.FluxUtil.withContext; @@ -132,10 +127,6 @@ String getLink() { } private Mono read(Context context) { - final boolean isTracingEnabled = this.container.getDatabase().getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "readTrigger." + container.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -144,39 +135,15 @@ private Mono read(Context context) { put(TracerProvider.DB_STATEMENT, spanName); }}; - return container.getDatabase() + Mono responseMono = container.getDatabase() .getDocClientWrapper() .readTrigger(getLink(), null) .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, container)) - .single() - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(this.container.getDatabase().getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.complete(), response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.error(throwable), 0); - } - }); + .single(); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); } private Mono replace(CosmosTriggerProperties triggerSettings, Context context) { - final boolean isTracingEnabled = this.container.getDatabase().getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "replaceTrigger." + container.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -185,39 +152,15 @@ private Mono replace(CosmosTriggerProperties trigger put(TracerProvider.DB_STATEMENT, spanName); }}; - return container.getDatabase() + Mono responseMono = container.getDatabase() .getDocClientWrapper() .replaceTrigger(new Trigger(ModelBridgeInternal.toJsonFromJsonSerializable(triggerSettings)), null) .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, container)) - .single() - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(this.container.getDatabase().getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.complete(), response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.error(throwable), 0); - } - }); + .single(); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); } private Mono delete(Context context) { - final boolean isTracingEnabled = this.container.getDatabase().getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "deleteTrigger." + container.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -226,31 +169,11 @@ private Mono delete(Context context) { put(TracerProvider.DB_STATEMENT, spanName); }}; - return container.getDatabase() + Mono responseMono = container.getDatabase() .getDocClientWrapper() .deleteTrigger(getLink(), null) .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, container)) - .single() - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(this.container.getDatabase().getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.complete(), response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.error(throwable), 0); - } - }); + .single(); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java index 27d517bd1d7e..3c9eae92d9ca 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java @@ -3,8 +3,6 @@ package com.azure.cosmos; import com.azure.core.util.Context; -import com.azure.core.util.FluxUtil; -import com.azure.core.util.tracing.ProcessKind; import com.azure.cosmos.implementation.Paths; import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.implementation.UserDefinedFunction; @@ -12,11 +10,9 @@ import com.azure.cosmos.models.CosmosUserDefinedFunctionProperties; import com.azure.cosmos.models.ModelBridgeInternal; import reactor.core.publisher.Mono; -import reactor.core.publisher.Signal; import java.util.HashMap; import java.util.Map; -import java.util.Objects; import java.util.Optional; import java.util.concurrent.atomic.AtomicReference; @@ -150,37 +146,14 @@ private Mono read(Context context) { put(TracerProvider.DB_STATEMENT, spanName); }}; - return container.getDatabase().getDocClientWrapper().readUserDefinedFunction(getLink(), null) - .map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, container)).single() - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(this.container.getDatabase().getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.complete(), response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.error(throwable), 0); - } - }); + Mono responseMono = + container.getDatabase().getDocClientWrapper().readUserDefinedFunction(getLink(), null) + .map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, container)).single(); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); } private Mono replace(CosmosUserDefinedFunctionProperties udfSettings, Context context) { - final boolean isTracingEnabled = this.container.getDatabase().getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "replaceUDF." + container.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -189,39 +162,15 @@ private Mono replace(CosmosUserDefinedFu put(TracerProvider.DB_STATEMENT, spanName); }}; - return container.getDatabase() + Mono responseMono = container.getDatabase() .getDocClientWrapper() .replaceUserDefinedFunction(new UserDefinedFunction(ModelBridgeInternal.toJsonFromJsonSerializable(udfSettings)), null) .map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, container)) - .single() - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(this.container.getDatabase().getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.complete(), response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.error(throwable), 0); - } - }); + .single(); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); } private Mono delete(Context context) { - final boolean isTracingEnabled = this.container.getDatabase().getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "deleteUDF." + container.getId(); Map tracingAttributes = new HashMap() {{ put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); @@ -230,31 +179,11 @@ private Mono delete(Context context) { put(TracerProvider.DB_STATEMENT, spanName); }}; - return container.getDatabase() + Mono responseMono = container.getDatabase() .getDocClientWrapper() .deleteUserDefinedFunction(this.getLink(), null) .map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, container)) - .single() - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(this.container.getDatabase().getClient().getTracerProvider().startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.complete(), response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - this.container.getDatabase().getClient().getTracerProvider().endSpan(parentContext.get(), - Signal.error(throwable), 0); - } - }); + .single(); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java index f3381af9504f..25bcb7b394bc 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java @@ -3,21 +3,24 @@ package com.azure.cosmos.implementation; import com.azure.core.util.Context; +import com.azure.core.util.FluxUtil; import com.azure.core.util.logging.ClientLogger; import com.azure.core.util.tracing.ProcessKind; import com.azure.core.util.tracing.Tracer; import com.azure.cosmos.CosmosClientException; -import com.azure.cosmos.models.CosmosAsyncDatabaseResponse; +import com.azure.cosmos.models.CosmosAsyncItemResponse; import com.azure.cosmos.models.CosmosResponse; +import reactor.core.publisher.Mono; import reactor.core.publisher.Signal; import java.io.PrintWriter; import java.io.StringWriter; import java.util.ArrayList; -import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; public class TracerProvider { private final ClientLogger logger = new ClientLogger(TracerProvider.class); @@ -47,7 +50,7 @@ public boolean isEnabled() { /** * For each tracer plugged into the SDK a new tracing span is created. - * + *

* The {@code context} will be checked for containing information about a parent span. If a parent span is found the * new span will be added as a child, otherwise the span will be created and added to the context and any downstream * start calls will use the created span as the parent. @@ -69,8 +72,9 @@ public Context startSpan(String methodName, Context context, ProcessKind process /** * Given a context containing the current tracing span the span is marked completed with status info from * {@link Signal}. For each tracer plugged into the SDK the current tracing span is marked as completed. + * * @param context Additional metadata that is passed through the call stack. - * @param signal The signal indicates the status and contains the metadata we need to end the tracing span. + * @param signal The signal indicates the status and contains the metadata we need to end the tracing span. */ public void endSpan(Context context, Signal signal, int statusCode) { Objects.requireNonNull(context, "'context' cannot be null."); @@ -100,14 +104,104 @@ public void endSpan(Context context, Signal signal, int statusCo } } + public Mono traceEnabledCosmosResponsePublisher(Mono resultPublisher, + Map tracingAttributes, + Context context, + String spanName) { + final boolean isTracingEnabled = this.isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + return resultPublisher + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(this.startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + this.endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + this.endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }); + } + + public Mono traceEnabledNonCosmosResponsePublisher(Mono resultPublisher, + Map tracingAttributes, + Context context, + String spanName) { + final boolean isTracingEnabled = this.isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + return resultPublisher + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(this.startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + this.endSpan(parentContext.get(), Signal.complete(), 200); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + this.endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }); + } + + public Mono> traceEnabledCosmosItemResponsePublisher(Mono> resultPublisher, + Map tracingAttributes, + Context context, + String spanName) { + final boolean isTracingEnabled = this.isEnabled(); + final AtomicReference parentContext = isTracingEnabled + ? new AtomicReference<>(Context.NONE) + : null; + return resultPublisher + .doOnSubscribe(ignoredValue -> { + if (isTracingEnabled) { + reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); + Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); + Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); + if (!callerFunc.isPresent()) { + parentContext.set(this.startSpan(spanName, + context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + } + } + }).doOnSuccess(response -> { + if (isTracingEnabled) { + this.endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + } + }).doOnError(throwable -> { + if (isTracingEnabled) { + this.endSpan(parentContext.get(), Signal.error(throwable), 0); + } + }); + } + private void end(int statusCode, Throwable throwable, Context context) { for (Tracer tracer : tracers) { - if(throwable != null) { - tracer.setAttribute(TracerProvider.ERROR_MSG, throwable.getMessage(), context); - tracer.setAttribute(TracerProvider.ERROR_TYPE, throwable.getClass().getName(), context); + if (throwable != null) { + tracer.setAttribute(TracerProvider.ERROR_MSG, throwable.getMessage(), context); + tracer.setAttribute(TracerProvider.ERROR_TYPE, throwable.getClass().getName(), context); StringWriter errorStack = new StringWriter(); throwable.printStackTrace(new PrintWriter(errorStack)); - tracer.setAttribute(TracerProvider.ERROR_STACK, errorStack.toString(), context); + tracer.setAttribute(TracerProvider.ERROR_STACK, errorStack.toString(), context); } tracer.end(statusCode, throwable, context); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/module-info.java b/sdk/cosmos/azure-cosmos/src/main/java/module-info.java index cba344b67ef1..6a926894171c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/module-info.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/module-info.java @@ -48,4 +48,5 @@ opens com.azure.cosmos.implementation.routing to com.fasterxml.jackson.databind; uses com.azure.cosmos.implementation.guava25.base.PatternCompiler; + uses com.azure.core.util.tracing.Tracer; } diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java new file mode 100644 index 000000000000..963d05c49ce9 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java @@ -0,0 +1,265 @@ +package com.azure.cosmos; + +import com.azure.core.util.Context; +import com.azure.core.util.tracing.ProcessKind; +import com.azure.core.util.tracing.Tracer; +import com.azure.cosmos.implementation.CosmosItemProperties; +import com.azure.cosmos.implementation.LifeCycleUtils; +import com.azure.cosmos.implementation.TestConfigurations; +import com.azure.cosmos.implementation.TracerProvider; +import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils; +import com.azure.cosmos.models.CosmosAsyncContainerResponse; +import com.azure.cosmos.models.CosmosAsyncDatabaseResponse; +import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.CosmosStoredProcedureProperties; +import com.azure.cosmos.models.CosmosTriggerProperties; +import com.azure.cosmos.models.CosmosUserDefinedFunctionProperties; +import com.azure.cosmos.models.FeedOptions; +import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.models.TriggerOperation; +import com.azure.cosmos.models.TriggerType; +import com.azure.cosmos.rx.TestSuiteBase; +import org.mockito.Matchers; +import org.mockito.Mockito; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; +import reactor.core.publisher.Signal; + +import java.util.ServiceLoader; +import java.util.UUID; + +public class CosmosTracerTest extends TestSuiteBase { + private static final String ITEM_ID = "tracerDoc"; + CosmosAsyncClient client; + CosmosAsyncDatabase cosmosAsyncDatabase; + CosmosAsyncContainer cosmosAsyncContainer; + + @BeforeClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT) + public void beforeClass() throws Exception { + client = new CosmosClientBuilder() + .endpoint(TestConfigurations.HOST) + .key(TestConfigurations.MASTER_KEY) + .buildAsyncClient(); + cosmosAsyncDatabase = getSharedCosmosDatabase(client); + cosmosAsyncContainer = getSharedMultiPartitionCosmosContainer(client); + + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void cosmosAsyncClient() { + TracerProvider tracer = Mockito.spy(new TracerProvider(ServiceLoader.load(Tracer.class))); + ReflectionUtils.setTracerProvider(client, tracer); + + CosmosAsyncDatabaseResponse databaseResponse = + client.createDatabaseIfNotExists(cosmosAsyncDatabase.getId()).block(); + Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + Mockito.verify(tracer, Mockito.times(2)).endSpan(Matchers.any(Context.class), Matchers.any(Signal.class), + Matchers.anyInt()); //endSpan invocation is 1 greater than startSpan, because createDatabaseIfNotExists + // uses public read api + + client.readAllDatabases(new FeedOptions()).byPage().single().block(); + Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + Mockito.verify(tracer, Mockito.atLeast(4)).endSpan(Matchers.any(Context.class), Matchers.any(Signal.class), + Matchers.anyInt()); //We used atLeast in endSpan because for query we are calling endSpan on every doOnEach + + String query = "select * from c where c.id = '" + cosmosAsyncDatabase.getId() + "'"; + client.queryDatabases(query, new FeedOptions()).byPage().single().block(); + Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + Mockito.verify(tracer, Mockito.atLeast(6)).endSpan(Matchers.any(Context.class), Matchers.any(Signal.class), + Matchers.anyInt()); //We used atLeast in endSpan because for query we are calling endSpan on every doOnEach + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void cosmosAsyncDatabase() { + TracerProvider tracer = Mockito.spy(new TracerProvider(ServiceLoader.load(Tracer.class))); + ReflectionUtils.setTracerProvider(client, tracer); + + CosmosAsyncContainerResponse containerResponse = + cosmosAsyncDatabase.createContainerIfNotExists(cosmosAsyncContainer.getId(), + "/pk", 5000).block(); + Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + try { + cosmosAsyncDatabase.readProvisionedThroughput().block(); + } catch (CosmosClientException ex) { + //do nothing + } + + Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosAsyncDatabase.readAllUsers().byPage().single().block(); + Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosAsyncDatabase.readAllContainers().byPage().single().block(); + Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void cosmosAsyncContainer() { + TracerProvider tracer = Mockito.spy(new TracerProvider(ServiceLoader.load(Tracer.class))); + ReflectionUtils.setTracerProvider(client, tracer); + + cosmosAsyncContainer.read().block(); + Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + try { + cosmosAsyncContainer.readProvisionedThroughput().block(); + } catch (CosmosClientException ex) { + //do nothing + } + Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + CosmosItemProperties properties = new CosmosItemProperties(); + properties.setId(ITEM_ID); + cosmosAsyncContainer.createItem(properties).block(); + Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosAsyncContainer.upsertItem(properties, + new CosmosItemRequestOptions()).block(); + Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosAsyncContainer.readItem(ITEM_ID, PartitionKey.NONE, + CosmosItemProperties.class).block(); + Mockito.verify(tracer, Mockito.times(5)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosAsyncContainer.deleteItem(ITEM_ID, PartitionKey.NONE).block(); + Mockito.verify(tracer, Mockito.times(6)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosAsyncContainer.readAllItems(new FeedOptions(), CosmosItemRequestOptions.class).byPage().single().block(); + Mockito.verify(tracer, Mockito.times(7)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + String query = "select * from c where c.id = '" + ITEM_ID + "'"; + cosmosAsyncContainer.queryItems(query, new FeedOptions(), CosmosItemRequestOptions.class).byPage().single().block(); + Mockito.verify(tracer, Mockito.times(8)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void cosmosAsyncScripts() { + TracerProvider tracer = Mockito.spy(new TracerProvider(ServiceLoader.load(Tracer.class))); + ReflectionUtils.setTracerProvider(client, tracer); + + cosmosAsyncContainer.getScripts().readAllStoredProcedures(new FeedOptions()).byPage().single().block(); + Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosAsyncContainer.getScripts().readAllTriggers(new FeedOptions()).byPage().single().block(); + Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosAsyncContainer.getScripts().readAllUserDefinedFunctions(new FeedOptions()).byPage().single().block(); + Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + CosmosUserDefinedFunctionProperties cosmosUserDefinedFunctionProperties = + getCosmosUserDefinedFunctionProperties(); + CosmosUserDefinedFunctionProperties resultUdf = + cosmosAsyncContainer.getScripts().createUserDefinedFunction(cosmosUserDefinedFunctionProperties).block().getProperties(); + Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosAsyncContainer.getScripts().getUserDefinedFunction(cosmosUserDefinedFunctionProperties.getId()).read().block(); + Mockito.verify(tracer, Mockito.times(5)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosUserDefinedFunctionProperties.setBody("function() {var x = 15;}"); + cosmosAsyncContainer.getScripts().getUserDefinedFunction(resultUdf.getId()).replace(resultUdf).block(); + Mockito.verify(tracer, Mockito.times(6)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosAsyncContainer.getScripts().readAllUserDefinedFunctions(new FeedOptions()).byPage().single().block(); + Mockito.verify(tracer, Mockito.times(7)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosAsyncContainer.getScripts().getUserDefinedFunction(cosmosUserDefinedFunctionProperties.getId()).delete().block(); + Mockito.verify(tracer, Mockito.times(8)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + CosmosTriggerProperties cosmosTriggerProperties = getCosmosTriggerProperties(); + CosmosTriggerProperties resultTrigger = + cosmosAsyncContainer.getScripts().createTrigger(cosmosTriggerProperties).block().getProperties(); + Mockito.verify(tracer, Mockito.times(9)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosAsyncContainer.getScripts().getTrigger(cosmosTriggerProperties.getId()).read().block(); + Mockito.verify(tracer, Mockito.times(10)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosAsyncContainer.getScripts().getTrigger(cosmosTriggerProperties.getId()).replace(resultTrigger).block(); + Mockito.verify(tracer, Mockito.times(11)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosAsyncContainer.getScripts().readAllTriggers(new FeedOptions()).byPage().single().block(); + Mockito.verify(tracer, Mockito.times(12)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosAsyncContainer.getScripts().getTrigger(cosmosTriggerProperties.getId()).delete().block(); + Mockito.verify(tracer, Mockito.times(13)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + CosmosStoredProcedureProperties procedureProperties = getCosmosStoredProcedureProperties(); + CosmosStoredProcedureProperties resultSproc = + cosmosAsyncContainer.getScripts().createStoredProcedure(procedureProperties).block().getProperties(); + Mockito.verify(tracer, Mockito.times(14)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).read().block(); + Mockito.verify(tracer, Mockito.times(15)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).replace(resultSproc).block(); + Mockito.verify(tracer, Mockito.times(16)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosAsyncContainer.getScripts().readAllStoredProcedures(new FeedOptions()).byPage().single().block(); + Mockito.verify(tracer, Mockito.times(17)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + + cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).delete().block(); + Mockito.verify(tracer, Mockito.times(18)).startSpan(Matchers.anyString(), Matchers.any(Context.class), + Matchers.any(ProcessKind.class)); + } + + @AfterClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT) + public void afterClass() throws Exception { + LifeCycleUtils.closeQuietly(client); + } + + private static CosmosUserDefinedFunctionProperties getCosmosUserDefinedFunctionProperties() { + CosmosUserDefinedFunctionProperties udf = new CosmosUserDefinedFunctionProperties(); + udf.setId(UUID.randomUUID().toString()); + udf.setBody("function() {var x = 10;}"); + return udf; + } + + private static CosmosTriggerProperties getCosmosTriggerProperties() { + CosmosTriggerProperties trigger = new CosmosTriggerProperties(); + trigger.setId(UUID.randomUUID().toString()); + trigger.setBody("function() {var x = 10;}"); + trigger.setTriggerOperation(TriggerOperation.CREATE); + trigger.setTriggerType(TriggerType.PRE); + return trigger; + } + + private static CosmosStoredProcedureProperties getCosmosStoredProcedureProperties() { + CosmosStoredProcedureProperties storedProcedureDef = new CosmosStoredProcedureProperties(); + storedProcedureDef.setId(UUID.randomUUID().toString()); + storedProcedureDef.setBody("function() {var x = 10;}"); + return storedProcedureDef; + } +} diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/ReflectionUtils.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/ReflectionUtils.java index 9dbdc312ffcd..5a9063584e37 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/ReflectionUtils.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/directconnectivity/ReflectionUtils.java @@ -9,6 +9,7 @@ import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.GlobalEndpointManager; import com.azure.cosmos.implementation.RxDocumentClientImpl; +import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.implementation.http.HttpClient; import org.apache.commons.lang3.reflect.FieldUtils; @@ -104,4 +105,8 @@ public static GatewayServiceConfigurationReader getServiceConfigurationReader(Rx public static void setBackgroundRefreshLocationTimeIntervalInMS(GlobalEndpointManager globalEndPointManager, int millSec){ set(globalEndPointManager, millSec, "backgroundRefreshLocationTimeIntervalInMS"); } + + public static void setTracerProvider(CosmosAsyncClient cosmosAsyncClient, TracerProvider tracerProvider){ + set(cosmosAsyncClient, tracerProvider, "tracerProvider"); + } } From eabbec95e6c61b20fa9365824e0944c5c5dcaf1a Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Fri, 24 Apr 2020 16:05:35 -0400 Subject: [PATCH 07/40] refactoring code with PR ready version --- sdk/cosmos/azure-cosmos/pom.xml | 2 +- .../com/azure/cosmos/CosmosAsyncClient.java | 48 +---- .../azure/cosmos/CosmosAsyncContainer.java | 179 ++++-------------- .../com/azure/cosmos/CosmosAsyncDatabase.java | 113 +++-------- .../com/azure/cosmos/CosmosAsyncScripts.java | 71 ++----- .../cosmos/CosmosAsyncStoredProcedure.java | 69 ++----- .../com/azure/cosmos/CosmosAsyncTrigger.java | 52 +---- .../com/azure/cosmos/CosmosAsyncUser.java | 6 +- .../CosmosAsyncUserDefinedFunction.java | 53 +----- .../CosmosPagedFluxOptions.java | 9 +- .../cosmos/implementation/TracerProvider.java | 121 ++++++------ .../azure/cosmos/util/CosmosPagedFlux.java | 46 +---- .../src/main/java/module-info.java | 4 +- .../com/azure/cosmos/CosmosTracerTest.java | 67 ++++--- 14 files changed, 226 insertions(+), 614 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/pom.xml b/sdk/cosmos/azure-cosmos/pom.xml index 9c8f93ca2ced..445feb8ede6c 100644 --- a/sdk/cosmos/azure-cosmos/pom.xml +++ b/sdk/cosmos/azure-cosmos/pom.xml @@ -119,7 +119,7 @@ Licensed under the MIT License. com.azure azure-core - 1.4.0-beta.1 + 1.4.0 io.netty diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index ccb4a72386ec..3ec556efe5b8 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -24,10 +24,8 @@ import reactor.core.publisher.Mono; import java.io.Closeable; -import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.Optional; import static com.azure.core.util.FluxUtil.withContext; import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; @@ -217,13 +215,7 @@ public Mono createDatabase(CosmosDatabaseProperties Database wrappedDatabase = new Database(); wrappedDatabase.setId(databaseSettings.getId()); final CosmosDatabaseRequestOptions requestOptions = options; - return withContext(context -> createDatabase(wrappedDatabase, requestOptions, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> createDatabase(wrappedDatabase, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -279,13 +271,7 @@ public Mono createDatabase(CosmosDatabaseProperties Database wrappedDatabase = new Database(); wrappedDatabase.setId(databaseSettings.getId()); final CosmosDatabaseRequestOptions requestOptions = options; - return withContext(context -> createDatabase(wrappedDatabase, requestOptions, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> createDatabase(wrappedDatabase, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -337,7 +323,7 @@ public Mono createDatabase(String id, int throughpu public CosmosPagedFlux readAllDatabases(FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { String spanName = "readAllDatabases"; - pagedFluxOptions.setTracerInformation(this.tracerProvider, spanName, this.serviceEndpoint); + pagedFluxOptions.setTracerInformation(this.tracerProvider, spanName, this.serviceEndpoint, null); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().readDatabases(options) .map(response -> @@ -390,7 +376,7 @@ public CosmosPagedFlux queryDatabases(String query, Fe public CosmosPagedFlux queryDatabases(SqlQuerySpec querySpec, FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { String spanName = "queryDatabases." + querySpec.getQueryText(); - pagedFluxOptions.setTracerInformation(this.tracerProvider, spanName, this.serviceEndpoint); + pagedFluxOptions.setTracerInformation(this.tracerProvider, spanName, this.serviceEndpoint, null); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().queryDatabases(querySpec, options) .map(response -> BridgeInternal.createFeedResponse( @@ -422,24 +408,13 @@ TracerProvider getTracerProvider(){ } private Mono createDatabaseIfNotExistsInternal(CosmosAsyncDatabase database) { - return withContext(context -> createDatabaseIfNotExistsInternal(database, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> createDatabaseIfNotExistsInternal(database, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } private Mono createDatabaseIfNotExistsInternal(CosmosAsyncDatabase database, Context context) { String spanName = "createDatabaseIfNotExistsInternal." + database.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, database.getId()); - put(TracerProvider.DB_URL, serviceEndpoint); - put(TracerProvider.DB_STATEMENT, spanName); - }}; + Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), this.serviceEndpoint, spanName); Mono responseMono = database.read().onErrorResume(exception -> { final Throwable unwrappedException = Exceptions.unwrap(exception); if (unwrappedException instanceof CosmosClientException) { @@ -454,14 +429,11 @@ private Mono createDatabaseIfNotExistsInternal(Cosm return tracerProvider.traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); } - private Mono createDatabase(Database database, CosmosDatabaseRequestOptions options, Context context){ + private Mono createDatabase(Database database, CosmosDatabaseRequestOptions options, + Context context) { String spanName = "createDatabase." + database.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, database.getId()); - put(TracerProvider.DB_URL, serviceEndpoint); - put(TracerProvider.DB_STATEMENT, spanName); - }}; + Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), + this.serviceEndpoint, spanName); Mono responseMono = asyncDocumentClient.createDatabase(database, ModelBridgeInternal.toRequestOptions(options)) .map(databaseResourceResponse -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(databaseResourceResponse, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 929857fbfeb0..29cd1edbc909 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -28,7 +28,6 @@ import java.util.HashMap; import java.util.Map; -import java.util.Optional; import java.util.stream.Collectors; import static com.azure.core.util.FluxUtil.withContext; @@ -91,13 +90,7 @@ public Mono read(CosmosContainerRequestOptions opt } final CosmosContainerRequestOptions requestOptions = options; - return withContext(context -> read(requestOptions, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> read(requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -117,24 +110,13 @@ public Mono delete(CosmosContainerRequestOptions o } final CosmosContainerRequestOptions requestOptions = options; - return withContext(context -> delete(requestOptions, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> delete(requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } private Mono delete(CosmosContainerRequestOptions options, Context context) { String spanName = "deleteContainer." + this.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, getId()); - put(TracerProvider.DB_URL, database.getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), + database.getClient().getServiceEndpoint(), spanName); Mono responseMono = database.getDocClientWrapper().deleteCollection(getLink(), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); @@ -194,26 +176,15 @@ public Mono replace( } final CosmosContainerRequestOptions requestOptions = options; - return withContext(context -> replace(containerProperties, requestOptions, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> replace(containerProperties, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } private Mono replace(CosmosContainerProperties containerProperties, CosmosContainerRequestOptions options, Context context) { String spanName = "replaceContainer." + this.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, getId()); - put(TracerProvider.DB_URL, database.getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), + database.getClient().getServiceEndpoint(), spanName); Mono responseMono = database.getDocClientWrapper() .replaceCollection(ModelBridgeInternal.getV2Collection(containerProperties), ModelBridgeInternal.toRequestOptions(options)) @@ -278,24 +249,13 @@ public Mono> createItem(T item, CosmosItemRequest } final CosmosItemRequestOptions requestOptions = options; - return withContext(context -> createItem(item, requestOptions, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> createItem(item, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } private Mono> createItem(T item, CosmosItemRequestOptions options, Context context) { String spanName = "createItem." + this.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, getId()); - put(TracerProvider.DB_URL, database.getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), + database.getClient().getServiceEndpoint(), spanName); @SuppressWarnings("unchecked") Class itemType = (Class) item.getClass(); RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); @@ -342,13 +302,7 @@ public Mono> upsertItem(T item, CosmosItemRequest } final CosmosItemRequestOptions requestOptions = options; - return withContext(context -> upsertItem(item, requestOptions, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> upsertItem(item, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -384,7 +338,7 @@ public CosmosPagedFlux readAllItems(FeedOptions options, Class classTy return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { String spanName = "readAllItems." + this.getId(); pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), spanName, - this.getDatabase().getClient().getServiceEndpoint()); + this.getDatabase().getClient().getServiceEndpoint(), database.getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDatabase().getDocClientWrapper().readDocuments(getLink(), options).map( response -> prepareFeedResponse(response, classType)); @@ -466,7 +420,7 @@ private CosmosPagedFlux queryItemsInternal( return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { String spanName = "queryItems." + this.getId() + "." + sqlQuerySpec.getQueryText(); pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), spanName, - this.getDatabase().getClient().getServiceEndpoint()); + this.getDatabase().getClient().getServiceEndpoint(), database.getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, feedOptions); return getDatabase().getDocClientWrapper() .queryDocuments(CosmosAsyncContainer.this.getLink(), sqlQuerySpec, feedOptions) @@ -522,13 +476,7 @@ public Mono> readItem( ModelBridgeInternal.setPartitionKey(options, partitionKey); RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); - return withContext(context -> readItem(itemId, partitionKey, requestOptions, itemType, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> readItem(itemId, partitionKey, requestOptions, itemType, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -573,13 +521,7 @@ public Mono> replaceItem( @SuppressWarnings("unchecked") Class itemType = (Class) item.getClass(); final CosmosItemRequestOptions requestOptions = options; - return withContext(context -> replaceItem(itemType, itemId, doc, requestOptions,context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> replaceItem(itemType, itemId, doc, requestOptions,context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -617,13 +559,7 @@ public Mono> deleteItem( } ModelBridgeInternal.setPartitionKey(options, partitionKey); RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); - return withContext(context -> deleteItem(itemId, requestOptions, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> deleteItem(itemId, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } private String getItemLink(String itemId) { @@ -659,7 +595,7 @@ public CosmosPagedFlux readAllConflicts(FeedOptions op return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { String spanName = "readAllConflicts." + this.getId(); pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), spanName, - this.getDatabase().getClient().getServiceEndpoint()); + this.getDatabase().getClient().getServiceEndpoint(), database.getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper().readConflicts(getLink(), options) .map(response -> BridgeInternal.createFeedResponse( @@ -691,7 +627,7 @@ public CosmosPagedFlux queryConflicts(String query, Fe return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { String spanName = "queryConflicts." + this.getId() + "." + query; pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), spanName, - this.getDatabase().getClient().getServiceEndpoint()); + this.getDatabase().getClient().getServiceEndpoint(), database.getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper().queryConflicts(getLink(), query, options) .map(response -> BridgeInternal.createFeedResponse( @@ -716,13 +652,7 @@ public CosmosAsyncConflict getConflict(String id) { * @return a {@link Mono} containing throughput or an error. */ public Mono readProvisionedThroughput() { - return withContext(context -> readProvisionedThroughput(context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> readProvisionedThroughput(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -734,13 +664,7 @@ public Mono readProvisionedThroughput() { * @return a {@link Mono} containing throughput or an error. */ public Mono replaceProvisionedThroughput(int requestUnitsPerSecond) { - return withContext(context -> replaceProvisionedThroughput(requestUnitsPerSecond, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> replaceProvisionedThroughput(requestUnitsPerSecond, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -769,13 +693,8 @@ private Mono> deleteItem( RequestOptions requestOptions, Context context) { String spanName = "deleteItem." + this.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, getId()); - put(TracerProvider.DB_URL, database.getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), + database.getClient().getServiceEndpoint(), spanName); Mono> responseMono = this.getDatabase() .getDocClientWrapper() .deleteDocument(getItemLink(itemId), requestOptions) @@ -793,13 +712,8 @@ private Mono> replaceItem( CosmosItemRequestOptions options, Context context) { String spanName = "replaceItem." + this.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, getId()); - put(TracerProvider.DB_URL, database.getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), + database.getClient().getServiceEndpoint(), spanName); Mono> responseMono = this.getDatabase() .getDocClientWrapper() .replaceDocument(getItemLink(itemId), doc, ModelBridgeInternal.toRequestOptions(options)) @@ -812,13 +726,8 @@ private Mono> replaceItem( private Mono> upsertItem(T item, CosmosItemRequestOptions options, Context context) { String spanName = "upsertItem." + this.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, getId()); - put(TracerProvider.DB_URL, database.getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), + database.getClient().getServiceEndpoint(), spanName); @SuppressWarnings("unchecked") Class itemType = (Class) item.getClass(); Mono> responseMono = this.getDatabase().getDocClientWrapper() @@ -836,13 +745,8 @@ private Mono> readItem( RequestOptions requestOptions, Class itemType, Context context) { String spanName = "readItem." + this.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, getId()); - put(TracerProvider.DB_URL, database.getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), + database.getClient().getServiceEndpoint(), spanName); Mono> responseMono = this.getDatabase().getDocClientWrapper() .readDocument(getItemLink(itemId), requestOptions) .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) @@ -853,13 +757,8 @@ private Mono> readItem( private Mono read(CosmosContainerRequestOptions options, Context context) { String spanName = "readContainer." + this.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, getId()); - put(TracerProvider.DB_URL, database.getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), + database.getClient().getServiceEndpoint(), spanName); Mono responseMono = database.getDocClientWrapper().readCollection(getLink(), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); @@ -870,13 +769,8 @@ private Mono read(CosmosContainerRequestOptions op private Mono readProvisionedThroughput(Context context) { String spanName = "readProvisionedThroughput." + this.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, getId()); - put(TracerProvider.DB_URL, getDatabase().getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), + database.getClient().getServiceEndpoint(), spanName); Mono responseMono = this.read() .flatMap(cosmosContainerResponse -> database.getDocClientWrapper() @@ -900,13 +794,8 @@ private Mono readProvisionedThroughput(Context context) { private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Context context) { String spanName = "replaceProvisionedThroughput." + this.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, getId()); - put(TracerProvider.DB_URL, getDatabase().getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), + database.getClient().getServiceEndpoint(), spanName); Mono responseMono = this.read() .flatMap(cosmosContainerResponse -> database.getDocClientWrapper() diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java index 2df8033ba2e2..bcfb8415b06d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java @@ -26,7 +26,6 @@ import java.util.HashMap; import java.util.Map; -import java.util.Optional; import static com.azure.core.util.FluxUtil.withContext; import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; @@ -84,17 +83,9 @@ public Mono read(CosmosDatabaseRequestOptions optio options = new CosmosDatabaseRequestOptions(); } final CosmosDatabaseRequestOptions requestOptions = options; - return withContext(context -> read(requestOptions, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> read(requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } - - /** * Deletes a database. *

@@ -124,13 +115,7 @@ public Mono delete(CosmosDatabaseRequestOptions opt } final CosmosDatabaseRequestOptions requestOptions = options; - return withContext(context -> delete(requestOptions, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> delete(requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /* CosmosAsyncContainer operations */ @@ -201,13 +186,7 @@ public Mono createContainer( } final CosmosContainerRequestOptions requestOptions = options; - return withContext(context -> createContainer(containerProperties, requestOptions, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> createContainer(containerProperties, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } @@ -357,13 +336,7 @@ public Mono createContainerIfNotExists( private Mono createContainerIfNotExistsInternal( CosmosContainerProperties containerProperties, CosmosAsyncContainer container, CosmosContainerRequestOptions options) { - return withContext(context -> createContainerIfNotExistsInternal(containerProperties, container, options, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> createContainerIfNotExistsInternal(containerProperties, container, options, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -381,7 +354,7 @@ public CosmosPagedFlux readAllContainers(FeedOptions return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { String spanName = "readAllContainers." + this.getId(); pagedFluxOptions.setTracerInformation(this.getClient().getTracerProvider(), spanName, - this.getClient().getServiceEndpoint()); + this.getClient().getServiceEndpoint(), getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().readCollections(getLink(), options) .map(response -> BridgeInternal.createFeedResponse( @@ -466,7 +439,7 @@ public CosmosPagedFlux queryContainers(SqlQuerySpec q return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { String spanName = "queryContainers." + this.getId() + "." + querySpec.getQueryText(); pagedFluxOptions.setTracerInformation(this.getClient().getTracerProvider(), spanName, - this.getClient().getServiceEndpoint()); + this.getClient().getServiceEndpoint(), getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().queryCollections(getLink(), querySpec, options) .map(response -> BridgeInternal.createFeedResponse( @@ -548,7 +521,7 @@ public CosmosPagedFlux readAllUsers(FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { String spanName = "readAllUsers." + this.getId(); pagedFluxOptions.setTracerInformation(this.getClient().getTracerProvider(), spanName, - this.getClient().getServiceEndpoint()); + this.getClient().getServiceEndpoint(), getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().readUsers(getLink(), options) .map(response -> BridgeInternal.createFeedResponse( @@ -619,7 +592,7 @@ public CosmosPagedFlux queryUsers(SqlQuerySpec querySpec, return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { String spanName = "queryUsers." + this.getId() + "." + querySpec.getQueryText(); pagedFluxOptions.setTracerInformation(this.getClient().getTracerProvider(), spanName, - this.getClient().getServiceEndpoint()); + this.getClient().getServiceEndpoint(), getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().queryUsers(getLink(), querySpec, options) .map(response -> BridgeInternal.createFeedResponseWithQueryMetrics( @@ -644,13 +617,7 @@ public CosmosAsyncUser getUser(String id) { * @return a {@link Mono} containing throughput or an error. */ public Mono readProvisionedThroughput() { - return withContext(context -> readProvisionedThroughput(context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> readProvisionedThroughput(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -662,13 +629,7 @@ public Mono readProvisionedThroughput() { * @return a {@link Mono} containing throughput or an error. */ public Mono replaceProvisionedThroughput(int requestUnitsPerSecond) { - return withContext(context -> replaceProvisionedThroughput(requestUnitsPerSecond, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> replaceProvisionedThroughput(requestUnitsPerSecond, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } CosmosAsyncClient getClient() { @@ -697,13 +658,8 @@ private Mono createContainerIfNotExistsInternal( CosmosContainerRequestOptions options, Context context) { String spanName = "createContainerIfNotExistsInternal." + containerProperties.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, getId()); - put(TracerProvider.DB_URL, client.getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(getId(), + getClient().getServiceEndpoint(), spanName); Mono responseMono = container.read(options).onErrorResume(exception -> { final Throwable unwrappedException = Exceptions.unwrap(exception); if (unwrappedException instanceof CosmosClientException) { @@ -723,13 +679,8 @@ private Mono createContainer( CosmosContainerRequestOptions options, Context context) { String spanName = "createContainer." + containerProperties.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, getId()); - put(TracerProvider.DB_URL, client.getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(getId(), + getClient().getServiceEndpoint(), spanName); Mono responseMono = getDocClientWrapper() .createCollection(this.getLink(), ModelBridgeInternal.getV2Collection(containerProperties), ModelBridgeInternal.toRequestOptions(options)) @@ -740,13 +691,8 @@ private Mono createContainer( private Mono read(CosmosDatabaseRequestOptions options, Context context) { String spanName = "readDatabase." + this.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, getId()); - put(TracerProvider.DB_URL, client.getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(getId(), + getClient().getServiceEndpoint(), spanName); Mono responseMono = getDocClientWrapper().readDatabase(getLink(), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single(); @@ -756,13 +702,8 @@ private Mono read(CosmosDatabaseRequestOptions opti private Mono delete(CosmosDatabaseRequestOptions options, Context context) { String spanName = "deleteDatabase." + this.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, getId()); - put(TracerProvider.DB_URL, client.getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(getId(), + getClient().getServiceEndpoint(), spanName); Mono responseMono = getDocClientWrapper().deleteDatabase(getLink(), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single(); @@ -772,13 +713,8 @@ private Mono delete(CosmosDatabaseRequestOptions op private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Context context) { String spanName = "replaceOffer." + this.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, getId()); - put(TracerProvider.DB_URL, client.getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(getId(), + getClient().getServiceEndpoint(), spanName); Mono responseMono = this.read() .flatMap(cosmosDatabaseResponse -> this.getDocClientWrapper() .queryOffers("select * from c where c.offerResourceId = '" @@ -806,13 +742,8 @@ private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Co private Mono readProvisionedThroughput(Context context) { String spanName = "readProvisionedThroughput." + this.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, getId()); - put(TracerProvider.DB_URL, client.getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(getId(), + getClient().getServiceEndpoint(), spanName); Mono responseMono = this.read() .flatMap(cosmosDatabaseResponse -> getDocClientWrapper() .queryOffers("select * from c where c.offerResourceId = '" diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java index 00b1578e18b4..c4618f20016f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java @@ -21,9 +21,7 @@ import com.azure.cosmos.util.UtilBridgeInternal; import reactor.core.publisher.Mono; -import java.util.HashMap; import java.util.Map; -import java.util.Optional; import static com.azure.core.util.FluxUtil.withContext; import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; @@ -79,13 +77,7 @@ public Mono createStoredProcedure( sProc.setId(properties.getId()); sProc.setBody(properties.getBody()); final CosmosStoredProcedureRequestOptions requestOptions = options; - return withContext(context -> createStoredProcedure(sProc, requestOptions, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> createStoredProcedure(sProc, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } @@ -107,7 +99,8 @@ public CosmosPagedFlux readAllStoredProcedures( String spanName = "readAllStoredProcedures." + this.container.getId(); pagedFluxOptions.setTracerInformation(this.container.getDatabase().getClient().getTracerProvider(), spanName, - this.container.getDatabase().getClient().getServiceEndpoint()); + this.container.getDatabase().getClient().getServiceEndpoint(), + this.container.getDatabase().getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper() .readStoredProcedures(container.getLink(), options) @@ -156,7 +149,8 @@ public CosmosPagedFlux queryStoredProcedures( String spanName = "queryStoredProcedures." + this.container.getId() + "." + querySpec.getQueryText(); pagedFluxOptions.setTracerInformation(this.container.getDatabase().getClient().getTracerProvider(), spanName, - this.container.getDatabase().getClient().getServiceEndpoint()); + this.container.getDatabase().getClient().getServiceEndpoint(), + this.container.getDatabase().getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper() .queryStoredProcedures(container.getLink(), querySpec, options) @@ -195,13 +189,7 @@ public Mono createUserDefinedFunction( udf.setId(properties.getId()); udf.setBody(properties.getBody()); - return withContext(context -> createUserDefinedFunction(udf, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> createUserDefinedFunction(udf, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -221,7 +209,8 @@ public CosmosPagedFlux readAllUserDefinedFu String spanName = "readAllUserDefinedFunctions." + this.container.getId(); pagedFluxOptions.setTracerInformation(this.container.getDatabase().getClient().getTracerProvider(), spanName, - this.container.getDatabase().getClient().getServiceEndpoint()); + this.container.getDatabase().getClient().getServiceEndpoint(), + this.container.getDatabase().getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper() .readUserDefinedFunctions(container.getLink(), options) @@ -272,7 +261,8 @@ public CosmosPagedFlux queryUserDefinedFunc String spanName = "queryUserDefinedFunctions." + this.container.getId() + "." + querySpec.getQueryText(); pagedFluxOptions.setTracerInformation(this.container.getDatabase().getClient().getTracerProvider(), spanName, - this.container.getDatabase().getClient().getServiceEndpoint()); + this.container.getDatabase().getClient().getServiceEndpoint(), + this.container.getDatabase().getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper() .queryUserDefinedFunctions(container.getLink(), querySpec, options) @@ -305,13 +295,7 @@ public CosmosAsyncUserDefinedFunction getUserDefinedFunction(String id) { * @return an {@link Mono} containing the single resource response with the created trigger or an error. */ public Mono createTrigger(CosmosTriggerProperties properties) { - return withContext(context -> createTrigger(properties, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> createTrigger(properties, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -332,7 +316,8 @@ public CosmosPagedFlux readAllTriggers(FeedOptions opti String spanName = "readAllTriggers." + this.container.getId(); pagedFluxOptions.setTracerInformation(this.container.getDatabase().getClient().getTracerProvider(), spanName, - this.container.getDatabase().getClient().getServiceEndpoint()); + this.container.getDatabase().getClient().getServiceEndpoint(), + this.container.getDatabase().getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper() .readTriggers(container.getLink(), options) @@ -377,7 +362,8 @@ public CosmosPagedFlux queryTriggers( String spanName = "queryTriggers." + this.container.getId(); pagedFluxOptions.setTracerInformation(this.container.getDatabase().getClient().getTracerProvider(), spanName, - this.container.getDatabase().getClient().getServiceEndpoint()); + this.container.getDatabase().getClient().getServiceEndpoint(), + this.container.getDatabase().getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper() .queryTriggers(container.getLink(), querySpec, options) @@ -401,13 +387,8 @@ private Mono createStoredProcedure(StoredPro CosmosStoredProcedureRequestOptions options, Context context) { String spanName = "createStoredProcedure." + container.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, database.getId()); - put(TracerProvider.DB_URL, container.getDatabase().getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), + database.getClient().getServiceEndpoint(), spanName); Mono responseMono = database.getDocClientWrapper() .createStoredProcedure(container.getLink(), sProc, ModelBridgeInternal.toRequestOptions(options)).map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, this.container)) .single(); @@ -418,13 +399,8 @@ private Mono createUserDefinedFunction( UserDefinedFunction udf, Context context) { String spanName = "createUserDefinedFunction." + container.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, database.getId()); - put(TracerProvider.DB_URL, container.getDatabase().getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), + database.getClient().getServiceEndpoint(), spanName); Mono responseMono = database.getDocClientWrapper() .createUserDefinedFunction(container.getLink(), udf, null).map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, this.container)).single(); @@ -433,13 +409,8 @@ private Mono createUserDefinedFunction( private Mono createTrigger(CosmosTriggerProperties properties, Context context) { String spanName = "createTrigger." + container.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, database.getId()); - put(TracerProvider.DB_URL, container.getDatabase().getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), + database.getClient().getServiceEndpoint(), spanName); Trigger trigger = new Trigger(ModelBridgeInternal.toJsonFromJsonSerializable(properties)); Mono responseMono = database.getDocClientWrapper() .createTrigger(container.getLink(), trigger, null) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java index 102445d05626..9a1a95fed2fc 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java @@ -14,7 +14,6 @@ import java.util.HashMap; import java.util.Map; -import java.util.Optional; import static com.azure.core.util.FluxUtil.withContext; @@ -78,13 +77,7 @@ public Mono read() { * @return an {@link Mono} containing the single resource response with the read stored procedure or an error. */ public Mono read(CosmosStoredProcedureRequestOptions options) { - return withContext(context -> read(options, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> read(options, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -113,13 +106,7 @@ public Mono delete() { * @return an {@link Mono} containing the single resource response for the deleted stored procedure or an error. */ public Mono delete(CosmosStoredProcedureRequestOptions options) { - return withContext(context -> delete(options, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> delete(options, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -136,13 +123,7 @@ public Mono delete(CosmosStoredProcedureRequ */ public Mono execute(Object[] procedureParams, CosmosStoredProcedureRequestOptions options) { - return withContext(context -> execute(procedureParams, options, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> execute(procedureParams, options, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -174,13 +155,7 @@ public Mono replace(CosmosStoredProcedurePro */ public Mono replace(CosmosStoredProcedureProperties storedProcedureSettings, CosmosStoredProcedureRequestOptions options) { - return withContext(context -> replace(storedProcedureSettings, options, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> replace(storedProcedureSettings, options, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } String getURIPathSegment() { @@ -204,13 +179,8 @@ String getLink() { private Mono read(CosmosStoredProcedureRequestOptions options, Context context) { String spanName = "readStoredProcedure." + cosmosContainer.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, cosmosContainer.getDatabase().getId()); - put(TracerProvider.DB_URL, cosmosContainer.getDatabase().getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(cosmosContainer.getDatabase().getId(), + cosmosContainer.getDatabase().getClient().getServiceEndpoint(), spanName); if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } @@ -224,13 +194,8 @@ private Mono read(CosmosStoredProcedureReque private Mono delete(CosmosStoredProcedureRequestOptions options, Context context) { String spanName = "deleteStoredProcedure." + cosmosContainer.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, cosmosContainer.getDatabase().getId()); - put(TracerProvider.DB_URL, cosmosContainer.getDatabase().getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(cosmosContainer.getDatabase().getId(), + cosmosContainer.getDatabase().getClient().getServiceEndpoint(), spanName); if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } @@ -249,13 +214,8 @@ private Mono execute(Object[] procedureParam CosmosStoredProcedureRequestOptions options, Context context) { String spanName = "executeStoredProcedure." + cosmosContainer.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, cosmosContainer.getDatabase().getId()); - put(TracerProvider.DB_URL, cosmosContainer.getDatabase().getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(cosmosContainer.getDatabase().getId(), + cosmosContainer.getDatabase().getClient().getServiceEndpoint(), spanName); if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } @@ -272,13 +232,8 @@ private Mono replace(CosmosStoredProcedurePr CosmosStoredProcedureRequestOptions options, Context context) { String spanName = "replaceStoredProcedure." + cosmosContainer.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, cosmosContainer.getDatabase().getId()); - put(TracerProvider.DB_URL, cosmosContainer.getDatabase().getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(cosmosContainer.getDatabase().getId(), + cosmosContainer.getDatabase().getClient().getServiceEndpoint(), spanName); if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java index 498d7b5d8588..5ac774388ee4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java @@ -11,9 +11,7 @@ import com.azure.cosmos.models.ModelBridgeInternal; import reactor.core.publisher.Mono; -import java.util.HashMap; import java.util.Map; -import java.util.Optional; import static com.azure.core.util.FluxUtil.withContext; @@ -60,13 +58,7 @@ CosmosAsyncTrigger setId(String id) { * @return an {@link Mono} containing the single resource response for the read cosmos trigger or an error. */ public Mono read() { - return withContext(context -> read(context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> read(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -80,13 +72,7 @@ public Mono read() { * @return an {@link Mono} containing the single resource response with the replaced cosmos trigger or an error. */ public Mono replace(CosmosTriggerProperties triggerSettings) { - return withContext(context -> replace(triggerSettings, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> replace(triggerSettings, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -99,13 +85,7 @@ public Mono replace(CosmosTriggerProperties triggerS * @return an {@link Mono} containing the single resource response for the deleted cosmos trigger or an error. */ public Mono delete() { - return withContext(context -> delete(context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> delete(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } String getURIPathSegment() { @@ -128,12 +108,8 @@ String getLink() { private Mono read(Context context) { String spanName = "readTrigger." + container.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, container.getDatabase().getId()); - put(TracerProvider.DB_URL, container.getDatabase().getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; + Map tracingAttributes = TracerProvider.createTracingMap(container.getDatabase().getId(), + container.getDatabase().getClient().getServiceEndpoint(), spanName); Mono responseMono = container.getDatabase() .getDocClientWrapper() @@ -145,13 +121,8 @@ private Mono read(Context context) { private Mono replace(CosmosTriggerProperties triggerSettings, Context context) { String spanName = "replaceTrigger." + container.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, container.getDatabase().getId()); - put(TracerProvider.DB_URL, container.getDatabase().getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(container.getDatabase().getId(), + container.getDatabase().getClient().getServiceEndpoint(), spanName); Mono responseMono = container.getDatabase() .getDocClientWrapper() .replaceTrigger(new Trigger(ModelBridgeInternal.toJsonFromJsonSerializable(triggerSettings)), null) @@ -162,13 +133,8 @@ private Mono replace(CosmosTriggerProperties trigger private Mono delete(Context context) { String spanName = "deleteTrigger." + container.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, container.getDatabase().getId()); - put(TracerProvider.DB_URL, container.getDatabase().getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(container.getDatabase().getId(), + container.getDatabase().getClient().getServiceEndpoint(), spanName); Mono responseMono = container.getDatabase() .getDocClientWrapper() .deleteTrigger(getLink(), null) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java index a231fa5d328f..528eb57cab2b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java @@ -149,7 +149,8 @@ public CosmosPagedFlux readAllPermissions(FeedOption String spanName = "readAllPermissions." + this.getId(); pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), spanName, - this.getDatabase().getClient().getServiceEndpoint()); + this.getDatabase().getClient().getServiceEndpoint(), + this.getDatabase().getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDatabase().getDocClientWrapper() .readPermissions(getLink(), options) @@ -191,7 +192,8 @@ public CosmosPagedFlux queryPermissions(String query String spanName = "queryPermissions." + this.getId() + "." + query; pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), spanName, - this.getDatabase().getClient().getServiceEndpoint()); + this.getDatabase().getClient().getServiceEndpoint(), + this.getDatabase().getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDatabase().getDocClientWrapper() .queryPermissions(getLink(), query, options) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java index 3c9eae92d9ca..82a7a0802869 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java @@ -11,9 +11,7 @@ import com.azure.cosmos.models.ModelBridgeInternal; import reactor.core.publisher.Mono; -import java.util.HashMap; import java.util.Map; -import java.util.Optional; import java.util.concurrent.atomic.AtomicReference; import static com.azure.core.util.FluxUtil.withContext; @@ -63,13 +61,7 @@ CosmosAsyncUserDefinedFunction setId(String id) { * @return an {@link Mono} containing the single resource response for the read user defined function or an error. */ public Mono read() { - return withContext(context -> read(context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> read(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -85,13 +77,7 @@ public Mono read() { * or an error. */ public Mono replace(CosmosUserDefinedFunctionProperties udfSettings) { - return withContext(context -> replace(udfSettings, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> replace(udfSettings, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -106,13 +92,7 @@ public Mono replace(CosmosUserDefinedFun * an error. */ public Mono delete() { - return withContext(context -> delete(context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return withContext(context -> delete(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } String getURIPathSegment() { @@ -139,13 +119,8 @@ private Mono read(Context context) { ? new AtomicReference<>(Context.NONE) : null; String spanName = "readUDF." + container.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, container.getDatabase().getId()); - put(TracerProvider.DB_URL, container.getDatabase().getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(container.getDatabase().getId(), + container.getDatabase().getClient().getServiceEndpoint(), spanName); Mono responseMono = container.getDatabase().getDocClientWrapper().readUserDefinedFunction(getLink(), null) .map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, container)).single(); @@ -155,13 +130,8 @@ private Mono read(Context context) { private Mono replace(CosmosUserDefinedFunctionProperties udfSettings, Context context) { String spanName = "replaceUDF." + container.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, container.getDatabase().getId()); - put(TracerProvider.DB_URL, container.getDatabase().getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(container.getDatabase().getId(), + container.getDatabase().getClient().getServiceEndpoint(), spanName); Mono responseMono = container.getDatabase() .getDocClientWrapper() .replaceUserDefinedFunction(new UserDefinedFunction(ModelBridgeInternal.toJsonFromJsonSerializable(udfSettings)), null) @@ -172,13 +142,8 @@ private Mono replace(CosmosUserDefinedFu private Mono delete(Context context) { String spanName = "deleteUDF." + container.getId(); - Map tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_INSTANCE, container.getDatabase().getId()); - put(TracerProvider.DB_URL, container.getDatabase().getClient().getServiceEndpoint()); - put(TracerProvider.DB_STATEMENT, spanName); - }}; - + Map tracingAttributes = TracerProvider.createTracingMap(container.getDatabase().getId(), + container.getDatabase().getClient().getServiceEndpoint(), spanName); Mono responseMono = container.getDatabase() .getDocClientWrapper() .deleteUserDefinedFunction(this.getLink(), null) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosPagedFluxOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosPagedFluxOptions.java index eb6340bba9ba..1f0a3904db47 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosPagedFluxOptions.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosPagedFluxOptions.java @@ -5,7 +5,6 @@ import com.azure.cosmos.util.CosmosPagedFlux; -import java.util.HashMap; import java.util.Map; /** @@ -89,12 +88,8 @@ public String getTracerSpanName() { return tracerSpanName; } - public void setTracerInformation(TracerProvider tracerProvider, String tracerSpanName, String serviceUrl) { - this.tracingAttributes = new HashMap() {{ - put(TracerProvider.DB_TYPE, TracerProvider.DB_TYPE_VALUE); - put(TracerProvider.DB_URL, serviceUrl); - put(TracerProvider.DB_STATEMENT, tracerSpanName); - }}; + public void setTracerInformation(TracerProvider tracerProvider, String tracerSpanName, String serviceUrl, String databaseId) { + this.tracingAttributes = TracerProvider.createTracingMap(databaseId, serviceUrl, tracerSpanName); this.tracerSpanName = tracerSpanName; this.tracerProvider = tracerProvider; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java index 25bcb7b394bc..691b1b522c53 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java @@ -5,7 +5,6 @@ import com.azure.core.util.Context; import com.azure.core.util.FluxUtil; import com.azure.core.util.logging.ClientLogger; -import com.azure.core.util.tracing.ProcessKind; import com.azure.core.util.tracing.Tracer; import com.azure.cosmos.CosmosClientException; import com.azure.cosmos.models.CosmosAsyncItemResponse; @@ -16,11 +15,14 @@ import java.io.PrintWriter; import java.io.StringWriter; import java.util.ArrayList; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; -import java.util.Optional; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; + +import static com.azure.core.util.tracing.Tracer.AZ_TRACING_NAMESPACE_KEY; public class TracerProvider { private final ClientLogger logger = new ClientLogger(TracerProvider.class); @@ -33,11 +35,35 @@ public class TracerProvider { public static final String ERROR_MSG = "error.msg"; public static final String ERROR_TYPE = "error.type"; public static final String ERROR_STACK = "error.stack"; - public static final String MASTER_CALL = "masterCall"; - public static final String NESTED_CALL = "nestedCall"; + public static final String COSMOS_CALL_DEPTH = "cosmosCallDepth"; + public static final String MICROSOFT_DOCOMENTDB = "Microsoft.DocumentDB"; public static final Object ATTRIBUTE_MAP = "span-attributes"; + public static Function callDepthAttributeFunc = ( + reactor.util.context.Context reactorContext) -> { + CallDepth callDepth = reactorContext.getOrDefault(COSMOS_CALL_DEPTH, CallDepth.ZERO); + if (callDepth.equals(CallDepth.ZERO)) { + return reactorContext.put(COSMOS_CALL_DEPTH, CallDepth.ONE); + } else if (callDepth.equals(CallDepth.ONE)) { + return reactorContext.put(COSMOS_CALL_DEPTH, CallDepth.TWO_OR_MORE); + } else { + return reactorContext; + } + }; + + public static Map createTracingMap(String databaseId, String serviceEndpoint, String spanName) { + return new HashMap() {{ + if (databaseId != null) { + put(TracerProvider.DB_INSTANCE, databaseId); + } + + put(AZ_TRACING_NAMESPACE_KEY, MICROSOFT_DOCOMENTDB); + put(TracerProvider.DB_TYPE, DB_TYPE_VALUE); + put(TracerProvider.DB_URL, serviceEndpoint); + put(TracerProvider.DB_STATEMENT, spanName); + }}; + } public TracerProvider(Iterable tracers) { Objects.requireNonNull(tracers, "'tracers' cannot be null."); @@ -58,12 +84,14 @@ public boolean isEnabled() { * @param context Additional metadata that is passed through the call stack. * @return An updated context object. */ - public Context startSpan(String methodName, Context context, ProcessKind processKind) { + public Context startSpan(String methodName, Context context, Map attributeMap) { Context local = Objects.requireNonNull(context, "'context' cannot be null."); - Objects.requireNonNull(processKind, "'processKind' cannot be null."); - for (Tracer tracer : tracers) { - local = tracer.start(methodName, local, processKind); + local = tracer.start(methodName, local); // start the span and return the started span + Context lambdaContext = local; + attributeMap.forEach((key, value) -> { + tracer.setAttribute(key, value, lambdaContext); // set attrs on the span + }); } return local; @@ -108,66 +136,30 @@ public Mono traceEnabledCosmosResponsePublisher(Mo Map tracingAttributes, Context context, String spanName) { - final boolean isTracingEnabled = this.isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; - return resultPublisher - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(this.startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - this.endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - this.endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }); + return traceEnabledPublisher(resultPublisher, tracingAttributes, context, spanName, + (T response) -> response.getStatusCode()); } public Mono traceEnabledNonCosmosResponsePublisher(Mono resultPublisher, Map tracingAttributes, Context context, String spanName) { - final boolean isTracingEnabled = this.isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; - return resultPublisher - .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { - parentContext.set(this.startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); - } - } - }).doOnSuccess(response -> { - if (isTracingEnabled) { - this.endSpan(parentContext.get(), Signal.complete(), 200); - } - }).doOnError(throwable -> { - if (isTracingEnabled) { - this.endSpan(parentContext.get(), Signal.error(throwable), 0); - } - }); + return traceEnabledPublisher(resultPublisher, tracingAttributes, context, spanName, (T response) -> 200); } public Mono> traceEnabledCosmosItemResponsePublisher(Mono> resultPublisher, Map tracingAttributes, Context context, String spanName) { + return traceEnabledPublisher(resultPublisher, tracingAttributes, context, spanName, + (CosmosAsyncItemResponse response) -> response.getStatusCode()); + } + + public Mono traceEnabledPublisher(Mono resultPublisher, + Map tracingAttributes, + Context context, + String spanName, + Function statusCodeFunc) { final boolean isTracingEnabled = this.isEnabled(); final AtomicReference parentContext = isTracingEnabled ? new AtomicReference<>(Context.NONE) @@ -175,17 +167,16 @@ public Mono> traceEnabledCosmosItemResponsePublis return resultPublisher .doOnSubscribe(ignoredValue -> { if (isTracingEnabled) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { + CallDepth callDepth = FluxUtil.toReactorContext(context).getOrDefault(COSMOS_CALL_DEPTH, + CallDepth.ZERO); + if (!callDepth.equals(CallDepth.TWO_OR_MORE)) { parentContext.set(this.startSpan(spanName, - context.addData(TracerProvider.ATTRIBUTE_MAP, tracingAttributes), ProcessKind.DATABASE)); + context, tracingAttributes)); } } }).doOnSuccess(response -> { if (isTracingEnabled) { - this.endSpan(parentContext.get(), Signal.complete(), response.getStatusCode()); + this.endSpan(parentContext.get(), Signal.complete(), statusCodeFunc.apply(response)); } }).doOnError(throwable -> { if (isTracingEnabled) { @@ -206,4 +197,10 @@ private void end(int statusCode, Throwable throwable, Context context) { tracer.end(statusCode, throwable, context); } } + + public enum CallDepth { + ZERO, + ONE, + TWO_OR_MORE, + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java index 1527d08267b9..388d99b61f8a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java @@ -3,22 +3,17 @@ package com.azure.cosmos.util; -import com.azure.core.http.HttpHeader; import com.azure.core.util.Context; import com.azure.core.util.FluxUtil; import com.azure.core.util.IterableStream; import com.azure.core.util.paging.ContinuablePagedFlux; -import com.azure.core.util.tracing.ProcessKind; import com.azure.cosmos.implementation.CosmosPagedFluxOptions; import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.models.FeedResponse; -import io.netty.handler.codec.http.HttpConstants; import reactor.core.CoreSubscriber; import reactor.core.publisher.Flux; import reactor.core.publisher.Signal; -import java.util.Objects; -import java.util.Optional; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; @@ -47,13 +42,7 @@ public final class CosmosPagedFlux extends ContinuablePagedFlux> byPage() { CosmosPagedFluxOptions cosmosPagedFluxOptions = new CosmosPagedFluxOptions(); - return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } @Override @@ -61,13 +50,7 @@ public Flux> byPage(String continuationToken) { CosmosPagedFluxOptions cosmosPagedFluxOptions = new CosmosPagedFluxOptions(); cosmosPagedFluxOptions.setRequestContinuation(continuationToken); - return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } @Override @@ -75,13 +58,7 @@ public Flux> byPage(int preferredPageSize) { CosmosPagedFluxOptions cosmosPagedFluxOptions = new CosmosPagedFluxOptions(); cosmosPagedFluxOptions.setMaxItemCount(preferredPageSize); - return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } @Override @@ -90,13 +67,7 @@ public Flux> byPage(String continuationToken, int preferredPageS cosmosPagedFluxOptions.setRequestContinuation(continuationToken); cosmosPagedFluxOptions.setMaxItemCount(preferredPageSize); - return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)).subscriberContext(reactorContext -> { - Optional master = reactorContext.getOrEmpty(TracerProvider.MASTER_CALL); - if (master.isPresent()) { - reactorContext = reactorContext.put(TracerProvider.NESTED_CALL, true); - } - return reactorContext.put(TracerProvider.MASTER_CALL, true); - }); + return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -123,12 +94,11 @@ private Flux> byPage(CosmosPagedFluxOptions pagedFluxOptions, Co return this.optionsFluxFunction.apply(pagedFluxOptions).doOnSubscribe(ignoredValue -> { if (pagedFluxOptions.getTracerProvider().isEnabled()) { - reactor.util.context.Context reactorContext = FluxUtil.toReactorContext(context); - Objects.requireNonNull(reactorContext.hasKey(TracerProvider.MASTER_CALL)); - Optional callerFunc = reactorContext.getOrEmpty(TracerProvider.NESTED_CALL); - if (!callerFunc.isPresent()) { + TracerProvider.CallDepth callDepth = FluxUtil.toReactorContext(context).getOrDefault(TracerProvider.COSMOS_CALL_DEPTH, + TracerProvider.CallDepth.ZERO); + if (!callDepth.equals(TracerProvider.CallDepth.TWO_OR_MORE)) { parentContext.set(pagedFluxOptions.getTracerProvider().startSpan(pagedFluxOptions.getTracerSpanName(), - context.addData(TracerProvider.ATTRIBUTE_MAP, pagedFluxOptions.getTracingAttributes()), ProcessKind.DATABASE)); + context, pagedFluxOptions.getTracingAttributes())); } } }).doOnEach(responseSignal -> { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/module-info.java b/sdk/cosmos/azure-cosmos/src/main/java/module-info.java index 5500ddcc3e81..b29c0d76586a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/module-info.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/module-info.java @@ -3,7 +3,8 @@ module com.azure.cosmos { - requires transitive azure.core; + requires transitive com.azure.core; + requires com.fasterxml.jackson.datatype.jsr310; requires io.netty.transport; requires io.netty.handler; @@ -26,6 +27,7 @@ requires reactor.core; requires org.reactivestreams; requires com.fasterxml.jackson.databind; + requires org.slf4j; // public API surface area exports com.azure.cosmos; diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java index 963d05c49ce9..7ac4460a5881 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java @@ -1,7 +1,6 @@ package com.azure.cosmos; import com.azure.core.util.Context; -import com.azure.core.util.tracing.ProcessKind; import com.azure.core.util.tracing.Tracer; import com.azure.cosmos.implementation.CosmosItemProperties; import com.azure.cosmos.implementation.LifeCycleUtils; @@ -54,21 +53,21 @@ public void cosmosAsyncClient() { CosmosAsyncDatabaseResponse databaseResponse = client.createDatabaseIfNotExists(cosmosAsyncDatabase.getId()).block(); Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); Mockito.verify(tracer, Mockito.times(2)).endSpan(Matchers.any(Context.class), Matchers.any(Signal.class), Matchers.anyInt()); //endSpan invocation is 1 greater than startSpan, because createDatabaseIfNotExists // uses public read api client.readAllDatabases(new FeedOptions()).byPage().single().block(); Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); Mockito.verify(tracer, Mockito.atLeast(4)).endSpan(Matchers.any(Context.class), Matchers.any(Signal.class), Matchers.anyInt()); //We used atLeast in endSpan because for query we are calling endSpan on every doOnEach String query = "select * from c where c.id = '" + cosmosAsyncDatabase.getId() + "'"; client.queryDatabases(query, new FeedOptions()).byPage().single().block(); Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); Mockito.verify(tracer, Mockito.atLeast(6)).endSpan(Matchers.any(Context.class), Matchers.any(Signal.class), Matchers.anyInt()); //We used atLeast in endSpan because for query we are calling endSpan on every doOnEach } @@ -82,7 +81,7 @@ public void cosmosAsyncDatabase() { cosmosAsyncDatabase.createContainerIfNotExists(cosmosAsyncContainer.getId(), "/pk", 5000).block(); Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); try { cosmosAsyncDatabase.readProvisionedThroughput().block(); @@ -91,15 +90,15 @@ public void cosmosAsyncDatabase() { } Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); cosmosAsyncDatabase.readAllUsers().byPage().single().block(); Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); cosmosAsyncDatabase.readAllContainers().byPage().single().block(); Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) @@ -109,7 +108,7 @@ public void cosmosAsyncContainer() { cosmosAsyncContainer.read().block(); Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); try { cosmosAsyncContainer.readProvisionedThroughput().block(); @@ -117,36 +116,36 @@ public void cosmosAsyncContainer() { //do nothing } Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); CosmosItemProperties properties = new CosmosItemProperties(); properties.setId(ITEM_ID); cosmosAsyncContainer.createItem(properties).block(); Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); cosmosAsyncContainer.upsertItem(properties, new CosmosItemRequestOptions()).block(); Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); cosmosAsyncContainer.readItem(ITEM_ID, PartitionKey.NONE, CosmosItemProperties.class).block(); Mockito.verify(tracer, Mockito.times(5)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); cosmosAsyncContainer.deleteItem(ITEM_ID, PartitionKey.NONE).block(); Mockito.verify(tracer, Mockito.times(6)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); cosmosAsyncContainer.readAllItems(new FeedOptions(), CosmosItemRequestOptions.class).byPage().single().block(); Mockito.verify(tracer, Mockito.times(7)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); String query = "select * from c where c.id = '" + ITEM_ID + "'"; cosmosAsyncContainer.queryItems(query, new FeedOptions(), CosmosItemRequestOptions.class).byPage().single().block(); Mockito.verify(tracer, Mockito.times(8)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) @@ -156,83 +155,81 @@ public void cosmosAsyncScripts() { cosmosAsyncContainer.getScripts().readAllStoredProcedures(new FeedOptions()).byPage().single().block(); Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); cosmosAsyncContainer.getScripts().readAllTriggers(new FeedOptions()).byPage().single().block(); Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); cosmosAsyncContainer.getScripts().readAllUserDefinedFunctions(new FeedOptions()).byPage().single().block(); Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); CosmosUserDefinedFunctionProperties cosmosUserDefinedFunctionProperties = getCosmosUserDefinedFunctionProperties(); CosmosUserDefinedFunctionProperties resultUdf = cosmosAsyncContainer.getScripts().createUserDefinedFunction(cosmosUserDefinedFunctionProperties).block().getProperties(); Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); cosmosAsyncContainer.getScripts().getUserDefinedFunction(cosmosUserDefinedFunctionProperties.getId()).read().block(); Mockito.verify(tracer, Mockito.times(5)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); cosmosUserDefinedFunctionProperties.setBody("function() {var x = 15;}"); cosmosAsyncContainer.getScripts().getUserDefinedFunction(resultUdf.getId()).replace(resultUdf).block(); Mockito.verify(tracer, Mockito.times(6)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); cosmosAsyncContainer.getScripts().readAllUserDefinedFunctions(new FeedOptions()).byPage().single().block(); Mockito.verify(tracer, Mockito.times(7)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); cosmosAsyncContainer.getScripts().getUserDefinedFunction(cosmosUserDefinedFunctionProperties.getId()).delete().block(); Mockito.verify(tracer, Mockito.times(8)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); CosmosTriggerProperties cosmosTriggerProperties = getCosmosTriggerProperties(); CosmosTriggerProperties resultTrigger = cosmosAsyncContainer.getScripts().createTrigger(cosmosTriggerProperties).block().getProperties(); Mockito.verify(tracer, Mockito.times(9)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); cosmosAsyncContainer.getScripts().getTrigger(cosmosTriggerProperties.getId()).read().block(); Mockito.verify(tracer, Mockito.times(10)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); cosmosAsyncContainer.getScripts().getTrigger(cosmosTriggerProperties.getId()).replace(resultTrigger).block(); Mockito.verify(tracer, Mockito.times(11)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); cosmosAsyncContainer.getScripts().readAllTriggers(new FeedOptions()).byPage().single().block(); Mockito.verify(tracer, Mockito.times(12)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); cosmosAsyncContainer.getScripts().getTrigger(cosmosTriggerProperties.getId()).delete().block(); Mockito.verify(tracer, Mockito.times(13)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); CosmosStoredProcedureProperties procedureProperties = getCosmosStoredProcedureProperties(); CosmosStoredProcedureProperties resultSproc = cosmosAsyncContainer.getScripts().createStoredProcedure(procedureProperties).block().getProperties(); Mockito.verify(tracer, Mockito.times(14)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).read().block(); Mockito.verify(tracer, Mockito.times(15)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).replace(resultSproc).block(); Mockito.verify(tracer, Mockito.times(16)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); cosmosAsyncContainer.getScripts().readAllStoredProcedures(new FeedOptions()).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(17)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).delete().block(); Mockito.verify(tracer, Mockito.times(18)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.any(ProcessKind.class)); + Matchers.anyMap()); } @AfterClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT) From 36dde959516c9db168cca34023a105bbb19aa5da Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Fri, 24 Apr 2020 17:06:40 -0400 Subject: [PATCH 08/40] ending span on doOnComplete --- .../src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java index 388d99b61f8a..0ce63cca60de 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java @@ -101,7 +101,7 @@ private Flux> byPage(CosmosPagedFluxOptions pagedFluxOptions, Co context, pagedFluxOptions.getTracingAttributes())); } } - }).doOnEach(responseSignal -> { + }).doOnComplete(() -> { if (pagedFluxOptions.getTracerProvider().isEnabled()) { pagedFluxOptions.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), 200); } From 36743ca506ec2fca6e25fe1c80062abcaba04688 Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Mon, 4 May 2020 11:02:31 -0400 Subject: [PATCH 09/40] resolving error --- .../CosmosAsyncUserDefinedFunction.java | 4 - .../cosmos/implementation/TracerProvider.java | 27 +++--- .../azure/cosmos/util/CosmosPagedFlux.java | 1 + .../com/azure/cosmos/CosmosTracerTest.java | 84 +++++++++---------- 4 files changed, 55 insertions(+), 61 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java index 82a7a0802869..0d3ed2d6575a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java @@ -114,10 +114,6 @@ String getLink() { } private Mono read(Context context) { - final boolean isTracingEnabled = this.container.getDatabase().getClient().getTracerProvider().isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; String spanName = "readUDF." + container.getId(); Map tracingAttributes = TracerProvider.createTracingMap(container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint(), spanName); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java index 691b1b522c53..376099c2d697 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java @@ -4,11 +4,11 @@ import com.azure.core.util.Context; import com.azure.core.util.FluxUtil; -import com.azure.core.util.logging.ClientLogger; import com.azure.core.util.tracing.Tracer; import com.azure.cosmos.CosmosClientException; import com.azure.cosmos.models.CosmosAsyncItemResponse; import com.azure.cosmos.models.CosmosResponse; +import com.azure.cosmos.models.Resource; import reactor.core.publisher.Mono; import reactor.core.publisher.Signal; @@ -25,7 +25,6 @@ import static com.azure.core.util.tracing.Tracer.AZ_TRACING_NAMESPACE_KEY; public class TracerProvider { - private final ClientLogger logger = new ClientLogger(TracerProvider.class); private final List tracers = new ArrayList<>(); public final static String DB_TYPE_VALUE = "cosmosdb"; public final static String DB_TYPE = "db.type"; @@ -38,11 +37,10 @@ public class TracerProvider { public static final String COSMOS_CALL_DEPTH = "cosmosCallDepth"; public static final String MICROSOFT_DOCOMENTDB = "Microsoft.DocumentDB"; - public static final Object ATTRIBUTE_MAP = "span-attributes"; - - public static Function callDepthAttributeFunc = ( + public final static Function callDepthAttributeFunc = ( reactor.util.context.Context reactorContext) -> { CallDepth callDepth = reactorContext.getOrDefault(COSMOS_CALL_DEPTH, CallDepth.ZERO); + assert callDepth != null; if (callDepth.equals(CallDepth.ZERO)) { return reactorContext.put(COSMOS_CALL_DEPTH, CallDepth.ONE); } else if (callDepth.equals(CallDepth.ONE)) { @@ -53,7 +51,7 @@ public class TracerProvider { }; public static Map createTracingMap(String databaseId, String serviceEndpoint, String spanName) { - return new HashMap() {{ + return new HashMap() {{ if (databaseId != null) { put(TracerProvider.DB_INSTANCE, databaseId); } @@ -67,7 +65,7 @@ public static Map createTracingMap(String databaseId, String ser public TracerProvider(Iterable tracers) { Objects.requireNonNull(tracers, "'tracers' cannot be null."); - tracers.forEach(e -> this.tracers.add(e)); + tracers.forEach(this.tracers::add); } public boolean isEnabled() { @@ -104,7 +102,7 @@ public Context startSpan(String methodName, Context context, Map signal, int statusCode) { + public > void endSpan(Context context, Signal signal, int statusCode) { Objects.requireNonNull(context, "'context' cannot be null."); Objects.requireNonNull(signal, "'signal' cannot be null."); @@ -132,12 +130,12 @@ public void endSpan(Context context, Signal signal, int statusCo } } - public Mono traceEnabledCosmosResponsePublisher(Mono resultPublisher, - Map tracingAttributes, - Context context, - String spanName) { + public > Mono traceEnabledCosmosResponsePublisher(Mono resultPublisher, + Map tracingAttributes, + Context context, + String spanName) { return traceEnabledPublisher(resultPublisher, tracingAttributes, context, spanName, - (T response) -> response.getStatusCode()); + CosmosResponse::getStatusCode); } public Mono traceEnabledNonCosmosResponsePublisher(Mono resultPublisher, @@ -152,7 +150,7 @@ public Mono> traceEnabledCosmosItemResponsePublis Context context, String spanName) { return traceEnabledPublisher(resultPublisher, tracingAttributes, context, spanName, - (CosmosAsyncItemResponse response) -> response.getStatusCode()); + CosmosAsyncItemResponse::getStatusCode); } public Mono traceEnabledPublisher(Mono resultPublisher, @@ -169,6 +167,7 @@ public Mono traceEnabledPublisher(Mono resultPublisher, if (isTracingEnabled) { CallDepth callDepth = FluxUtil.toReactorContext(context).getOrDefault(COSMOS_CALL_DEPTH, CallDepth.ZERO); + assert callDepth != null; if (!callDepth.equals(CallDepth.TWO_OR_MORE)) { parentContext.set(this.startSpan(spanName, context, tracingAttributes)); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java index 0ce63cca60de..998cb6d7c42c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java @@ -96,6 +96,7 @@ private Flux> byPage(CosmosPagedFluxOptions pagedFluxOptions, Co if (pagedFluxOptions.getTracerProvider().isEnabled()) { TracerProvider.CallDepth callDepth = FluxUtil.toReactorContext(context).getOrDefault(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.CallDepth.ZERO); + assert callDepth != null; if (!callDepth.equals(TracerProvider.CallDepth.TWO_OR_MORE)) { parentContext.set(pagedFluxOptions.getTracerProvider().startSpan(pagedFluxOptions.getTracerSpanName(), context, pagedFluxOptions.getTracingAttributes())); diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java index 7ac4460a5881..3bff6207d233 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java @@ -7,14 +7,14 @@ import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils; -import com.azure.cosmos.models.CosmosAsyncContainerResponse; -import com.azure.cosmos.models.CosmosAsyncDatabaseResponse; import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.CosmosResponse; import com.azure.cosmos.models.CosmosStoredProcedureProperties; import com.azure.cosmos.models.CosmosTriggerProperties; import com.azure.cosmos.models.CosmosUserDefinedFunctionProperties; import com.azure.cosmos.models.FeedOptions; import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.models.Resource; import com.azure.cosmos.models.TriggerOperation; import com.azure.cosmos.models.TriggerType; import com.azure.cosmos.rx.TestSuiteBase; @@ -35,7 +35,7 @@ public class CosmosTracerTest extends TestSuiteBase { CosmosAsyncContainer cosmosAsyncContainer; @BeforeClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT) - public void beforeClass() throws Exception { + public void beforeClass() { client = new CosmosClientBuilder() .endpoint(TestConfigurations.HOST) .key(TestConfigurations.MASTER_KEY) @@ -50,25 +50,24 @@ public void cosmosAsyncClient() { TracerProvider tracer = Mockito.spy(new TracerProvider(ServiceLoader.load(Tracer.class))); ReflectionUtils.setTracerProvider(client, tracer); - CosmosAsyncDatabaseResponse databaseResponse = - client.createDatabaseIfNotExists(cosmosAsyncDatabase.getId()).block(); + client.createDatabaseIfNotExists(cosmosAsyncDatabase.getId()).block(); Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); - Mockito.verify(tracer, Mockito.times(2)).endSpan(Matchers.any(Context.class), Matchers.any(Signal.class), + Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(2)).endSpan(Matchers.any(Context.class), Matchers.>>any(), Matchers.anyInt()); //endSpan invocation is 1 greater than startSpan, because createDatabaseIfNotExists // uses public read api client.readAllDatabases(new FeedOptions()).byPage().single().block(); Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); - Mockito.verify(tracer, Mockito.atLeast(4)).endSpan(Matchers.any(Context.class), Matchers.any(Signal.class), + Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.atLeast(3)).endSpan(Matchers.any(Context.class), Matchers.>>any(), Matchers.anyInt()); //We used atLeast in endSpan because for query we are calling endSpan on every doOnEach String query = "select * from c where c.id = '" + cosmosAsyncDatabase.getId() + "'"; client.queryDatabases(query, new FeedOptions()).byPage().single().block(); Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); - Mockito.verify(tracer, Mockito.atLeast(6)).endSpan(Matchers.any(Context.class), Matchers.any(Signal.class), + Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.atLeast(4)).endSpan(Matchers.any(Context.class), Matchers.>>any(), Matchers.anyInt()); //We used atLeast in endSpan because for query we are calling endSpan on every doOnEach } @@ -77,11 +76,10 @@ public void cosmosAsyncDatabase() { TracerProvider tracer = Mockito.spy(new TracerProvider(ServiceLoader.load(Tracer.class))); ReflectionUtils.setTracerProvider(client, tracer); - CosmosAsyncContainerResponse containerResponse = - cosmosAsyncDatabase.createContainerIfNotExists(cosmosAsyncContainer.getId(), + cosmosAsyncDatabase.createContainerIfNotExists(cosmosAsyncContainer.getId(), "/pk", 5000).block(); Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); try { cosmosAsyncDatabase.readProvisionedThroughput().block(); @@ -90,15 +88,15 @@ public void cosmosAsyncDatabase() { } Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); cosmosAsyncDatabase.readAllUsers().byPage().single().block(); Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); cosmosAsyncDatabase.readAllContainers().byPage().single().block(); Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) @@ -108,7 +106,7 @@ public void cosmosAsyncContainer() { cosmosAsyncContainer.read().block(); Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); try { cosmosAsyncContainer.readProvisionedThroughput().block(); @@ -116,36 +114,36 @@ public void cosmosAsyncContainer() { //do nothing } Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); CosmosItemProperties properties = new CosmosItemProperties(); properties.setId(ITEM_ID); cosmosAsyncContainer.createItem(properties).block(); Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); cosmosAsyncContainer.upsertItem(properties, new CosmosItemRequestOptions()).block(); Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); cosmosAsyncContainer.readItem(ITEM_ID, PartitionKey.NONE, CosmosItemProperties.class).block(); Mockito.verify(tracer, Mockito.times(5)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); cosmosAsyncContainer.deleteItem(ITEM_ID, PartitionKey.NONE).block(); Mockito.verify(tracer, Mockito.times(6)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); cosmosAsyncContainer.readAllItems(new FeedOptions(), CosmosItemRequestOptions.class).byPage().single().block(); Mockito.verify(tracer, Mockito.times(7)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); String query = "select * from c where c.id = '" + ITEM_ID + "'"; cosmosAsyncContainer.queryItems(query, new FeedOptions(), CosmosItemRequestOptions.class).byPage().single().block(); Mockito.verify(tracer, Mockito.times(8)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) @@ -155,85 +153,85 @@ public void cosmosAsyncScripts() { cosmosAsyncContainer.getScripts().readAllStoredProcedures(new FeedOptions()).byPage().single().block(); Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); cosmosAsyncContainer.getScripts().readAllTriggers(new FeedOptions()).byPage().single().block(); Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); cosmosAsyncContainer.getScripts().readAllUserDefinedFunctions(new FeedOptions()).byPage().single().block(); Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); CosmosUserDefinedFunctionProperties cosmosUserDefinedFunctionProperties = getCosmosUserDefinedFunctionProperties(); CosmosUserDefinedFunctionProperties resultUdf = cosmosAsyncContainer.getScripts().createUserDefinedFunction(cosmosUserDefinedFunctionProperties).block().getProperties(); Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); cosmosAsyncContainer.getScripts().getUserDefinedFunction(cosmosUserDefinedFunctionProperties.getId()).read().block(); Mockito.verify(tracer, Mockito.times(5)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); cosmosUserDefinedFunctionProperties.setBody("function() {var x = 15;}"); cosmosAsyncContainer.getScripts().getUserDefinedFunction(resultUdf.getId()).replace(resultUdf).block(); Mockito.verify(tracer, Mockito.times(6)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); cosmosAsyncContainer.getScripts().readAllUserDefinedFunctions(new FeedOptions()).byPage().single().block(); Mockito.verify(tracer, Mockito.times(7)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); cosmosAsyncContainer.getScripts().getUserDefinedFunction(cosmosUserDefinedFunctionProperties.getId()).delete().block(); Mockito.verify(tracer, Mockito.times(8)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); CosmosTriggerProperties cosmosTriggerProperties = getCosmosTriggerProperties(); CosmosTriggerProperties resultTrigger = cosmosAsyncContainer.getScripts().createTrigger(cosmosTriggerProperties).block().getProperties(); Mockito.verify(tracer, Mockito.times(9)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); cosmosAsyncContainer.getScripts().getTrigger(cosmosTriggerProperties.getId()).read().block(); Mockito.verify(tracer, Mockito.times(10)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); cosmosAsyncContainer.getScripts().getTrigger(cosmosTriggerProperties.getId()).replace(resultTrigger).block(); Mockito.verify(tracer, Mockito.times(11)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); cosmosAsyncContainer.getScripts().readAllTriggers(new FeedOptions()).byPage().single().block(); Mockito.verify(tracer, Mockito.times(12)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); cosmosAsyncContainer.getScripts().getTrigger(cosmosTriggerProperties.getId()).delete().block(); Mockito.verify(tracer, Mockito.times(13)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); CosmosStoredProcedureProperties procedureProperties = getCosmosStoredProcedureProperties(); CosmosStoredProcedureProperties resultSproc = cosmosAsyncContainer.getScripts().createStoredProcedure(procedureProperties).block().getProperties(); Mockito.verify(tracer, Mockito.times(14)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).read().block(); Mockito.verify(tracer, Mockito.times(15)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).replace(resultSproc).block(); Mockito.verify(tracer, Mockito.times(16)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); cosmosAsyncContainer.getScripts().readAllStoredProcedures(new FeedOptions()).byPage().single().block(); cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).delete().block(); Mockito.verify(tracer, Mockito.times(18)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMap()); + Matchers.anyMapOf(String.class, String.class)); } @AfterClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT) - public void afterClass() throws Exception { + public void afterClass() { LifeCycleUtils.closeQuietly(client); } From 8ca24ec529a3e6ea073cde068f3e46b2173cc6e5 Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Tue, 5 May 2020 16:06:26 -0400 Subject: [PATCH 10/40] perf improvement --- .../com/azure/cosmos/CosmosAsyncClient.java | 52 +++--- .../azure/cosmos/CosmosAsyncContainer.java | 97 +++++------ .../com/azure/cosmos/CosmosAsyncDatabase.java | 156 ++++++++++-------- .../azure/cosmos/CosmosAsyncPermission.java | 75 +++++++-- .../com/azure/cosmos/CosmosAsyncScripts.java | 131 +++++++++------ .../cosmos/CosmosAsyncStoredProcedure.java | 40 ++--- .../com/azure/cosmos/CosmosAsyncTrigger.java | 24 +-- .../com/azure/cosmos/CosmosAsyncUser.java | 92 ++++++++--- .../CosmosAsyncUserDefinedFunction.java | 24 +-- .../CosmosPagedFluxOptions.java | 32 ++-- .../cosmos/implementation/TracerProvider.java | 78 ++++----- .../azure/cosmos/util/CosmosPagedFlux.java | 20 +-- .../com/azure/cosmos/CosmosTracerTest.java | 130 +++++++-------- 13 files changed, 545 insertions(+), 406 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index 3ec556efe5b8..3b1ff821d391 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -25,7 +25,6 @@ import java.io.Closeable; import java.util.List; -import java.util.Map; import static com.azure.core.util.FluxUtil.withContext; import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; @@ -215,7 +214,8 @@ public Mono createDatabase(CosmosDatabaseProperties Database wrappedDatabase = new Database(); wrappedDatabase.setId(databaseSettings.getId()); final CosmosDatabaseRequestOptions requestOptions = options; - return withContext(context -> createDatabase(wrappedDatabase, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> createDatabase(wrappedDatabase, requestOptions + , context)), getTracerProvider()); } /** @@ -271,7 +271,8 @@ public Mono createDatabase(CosmosDatabaseProperties Database wrappedDatabase = new Database(); wrappedDatabase.setId(databaseSettings.getId()); final CosmosDatabaseRequestOptions requestOptions = options; - return withContext(context -> createDatabase(wrappedDatabase, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> createDatabase(wrappedDatabase, requestOptions + , context)), getTracerProvider()); } /** @@ -359,7 +360,7 @@ public CosmosPagedFlux readAllDatabases() { * @return a {@link CosmosPagedFlux} containing one or several feed response pages of read databases or an error. */ public CosmosPagedFlux queryDatabases(String query, FeedOptions options) { - return queryDatabases(new SqlQuerySpec(query), options); + return queryDatabasesInternal(false, new SqlQuerySpec(query), options); } /** @@ -374,15 +375,7 @@ public CosmosPagedFlux queryDatabases(String query, Fe * @return a {@link CosmosPagedFlux} containing one or several feed response pages of read databases or an error. */ public CosmosPagedFlux queryDatabases(SqlQuerySpec querySpec, FeedOptions options) { - return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "queryDatabases." + querySpec.getQueryText(); - 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())); - }); + return queryDatabasesInternal(true, querySpec, options); } /** @@ -407,14 +400,32 @@ TracerProvider getTracerProvider(){ return this.tracerProvider; } + private CosmosPagedFlux queryDatabasesInternal(boolean isParameterised, SqlQuerySpec querySpec, FeedOptions options){ + return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName; + if (isParameterised) { + spanName = "queryDatabases." + querySpec.getQueryText(); + } else { + 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())); + }); + } + private Mono createDatabaseIfNotExistsInternal(CosmosAsyncDatabase database) { - return withContext(context -> createDatabaseIfNotExistsInternal(database, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> createDatabaseIfNotExistsInternal(database, + context)), getTracerProvider()); } private Mono createDatabaseIfNotExistsInternal(CosmosAsyncDatabase database, Context context) { String spanName = "createDatabaseIfNotExistsInternal." + database.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), this.serviceEndpoint, spanName); Mono responseMono = database.read().onErrorResume(exception -> { final Throwable unwrappedException = Exceptions.unwrap(exception); if (unwrappedException instanceof CosmosClientException) { @@ -426,19 +437,20 @@ private Mono createDatabaseIfNotExistsInternal(Cosm } return Mono.error(unwrappedException); }); - return tracerProvider.traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); + return tracerProvider.traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), + this.serviceEndpoint); } private Mono createDatabase(Database database, CosmosDatabaseRequestOptions options, Context context) { String spanName = "createDatabase." + database.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), - this.serviceEndpoint, spanName); - Mono responseMono = asyncDocumentClient.createDatabase(database, ModelBridgeInternal.toRequestOptions(options)) + Mono responseMono = asyncDocumentClient.createDatabase(database, + ModelBridgeInternal.toRequestOptions(options)) .map(databaseResourceResponse -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(databaseResourceResponse, this)) .single(); - return tracerProvider.traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); + return tracerProvider.traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), + this.serviceEndpoint); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 29cd1edbc909..2c9df7932ae9 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -26,8 +26,6 @@ import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; -import java.util.HashMap; -import java.util.Map; import java.util.stream.Collectors; import static com.azure.core.util.FluxUtil.withContext; @@ -90,7 +88,7 @@ public Mono read(CosmosContainerRequestOptions opt } final CosmosContainerRequestOptions requestOptions = options; - return withContext(context -> read(requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> read(requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc), database.getClient().getTracerProvider()); } /** @@ -110,18 +108,17 @@ public Mono delete(CosmosContainerRequestOptions o } final CosmosContainerRequestOptions requestOptions = options; - return withContext(context -> delete(requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> delete(requestOptions, context)), + database.getClient().getTracerProvider()); } private Mono delete(CosmosContainerRequestOptions options, Context context) { String spanName = "deleteContainer." + this.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), - database.getClient().getServiceEndpoint(), spanName); Mono responseMono = database.getDocClientWrapper().deleteCollection(getLink(), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - tracingAttributes, context, spanName); + context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } /** @@ -176,21 +173,20 @@ public Mono replace( } final CosmosContainerRequestOptions requestOptions = options; - return withContext(context -> replace(containerProperties, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> replace(containerProperties, requestOptions, + context)), database.getClient().getTracerProvider()); } private Mono replace(CosmosContainerProperties containerProperties, CosmosContainerRequestOptions options, Context context) { String spanName = "replaceContainer." + this.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), - database.getClient().getServiceEndpoint(), spanName); Mono responseMono = database.getDocClientWrapper() .replaceCollection(ModelBridgeInternal.getV2Collection(containerProperties), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - tracingAttributes, context, spanName); + context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } /* CosmosAsyncItem operations */ @@ -249,13 +245,12 @@ public Mono> createItem(T item, CosmosItemRequest } final CosmosItemRequestOptions requestOptions = options; - return withContext(context -> createItem(item, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> createItem(item, requestOptions, context)), + database.getClient().getTracerProvider()); } private Mono> createItem(T item, CosmosItemRequestOptions options, Context context) { String spanName = "createItem." + this.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), - database.getClient().getServiceEndpoint(), spanName); @SuppressWarnings("unchecked") Class itemType = (Class) item.getClass(); RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); @@ -266,7 +261,8 @@ private Mono> createItem(T item, CosmosItemReques true) .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) .single(); - return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, tracingAttributes, context,spanName ); + return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, context + , spanName, database.getId(), database.getClient().getServiceEndpoint()); } /** @@ -302,7 +298,8 @@ public Mono> upsertItem(T item, CosmosItemRequest } final CosmosItemRequestOptions requestOptions = options; - return withContext(context -> upsertItem(item, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> upsertItem(item, requestOptions, context)), + database.getClient().getTracerProvider()); } /** @@ -359,7 +356,7 @@ public CosmosPagedFlux readAllItems(FeedOptions options, Class classTy * error. */ public CosmosPagedFlux queryItems(String query, Class classType) { - return queryItems(new SqlQuerySpec(query), classType); + return queryItemsInternal(false, new SqlQuerySpec(query), new FeedOptions(), classType); } /** @@ -377,7 +374,7 @@ public CosmosPagedFlux queryItems(String query, Class classType) { * error. */ public CosmosPagedFlux queryItems(String query, FeedOptions options, Class classType) { - return queryItems(new SqlQuerySpec(query), options, classType); + return queryItemsInternal(false, new SqlQuerySpec(query), options, classType); } /** @@ -394,7 +391,7 @@ public CosmosPagedFlux queryItems(String query, FeedOptions options, Clas * error. */ public CosmosPagedFlux queryItems(SqlQuerySpec querySpec, Class classType) { - return queryItems(querySpec, new FeedOptions(), classType); + return queryItemsInternal(true, querySpec, new FeedOptions(), classType); } /** @@ -412,13 +409,19 @@ public CosmosPagedFlux queryItems(SqlQuerySpec querySpec, Class classT * error. */ public CosmosPagedFlux queryItems(SqlQuerySpec querySpec, FeedOptions options, Class classType) { - return queryItemsInternal(querySpec, options, classType); + return queryItemsInternal(true, querySpec, options, classType); } private CosmosPagedFlux queryItemsInternal( - SqlQuerySpec sqlQuerySpec, FeedOptions feedOptions, Class classType) { + boolean isParameterised, SqlQuerySpec sqlQuerySpec, FeedOptions feedOptions, Class classType) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "queryItems." + this.getId() + "." + sqlQuerySpec.getQueryText(); + String spanName; + if (isParameterised) { + spanName = "queryItems." + this.getId() + "." + sqlQuerySpec.getQueryText(); + } else { + spanName = "queryItems." + this.getId(); + } + pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), spanName, this.getDatabase().getClient().getServiceEndpoint(), database.getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, feedOptions); @@ -476,7 +479,8 @@ public Mono> readItem( ModelBridgeInternal.setPartitionKey(options, partitionKey); RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); - return withContext(context -> readItem(itemId, partitionKey, requestOptions, itemType, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> readItem(itemId, partitionKey, + requestOptions, itemType, context)), database.getClient().getTracerProvider()); } /** @@ -521,7 +525,7 @@ public Mono> replaceItem( @SuppressWarnings("unchecked") Class itemType = (Class) item.getClass(); final CosmosItemRequestOptions requestOptions = options; - return withContext(context -> replaceItem(itemType, itemId, doc, requestOptions,context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> replaceItem(itemType, itemId, doc, requestOptions,context)), database.getClient().getTracerProvider()); } /** @@ -559,7 +563,7 @@ public Mono> deleteItem( } ModelBridgeInternal.setPartitionKey(options, partitionKey); RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); - return withContext(context -> deleteItem(itemId, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> deleteItem(itemId, requestOptions, context)), database.getClient().getTracerProvider()); } private String getItemLink(String itemId) { @@ -625,7 +629,7 @@ public CosmosPagedFlux queryConflicts(String query) { */ public CosmosPagedFlux queryConflicts(String query, FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "queryConflicts." + this.getId() + "." + query; + String spanName = "queryConflicts." + this.getId(); pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), spanName, this.getDatabase().getClient().getServiceEndpoint(), database.getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); @@ -652,7 +656,7 @@ public CosmosAsyncConflict getConflict(String id) { * @return a {@link Mono} containing throughput or an error. */ public Mono readProvisionedThroughput() { - return withContext(context -> readProvisionedThroughput(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> readProvisionedThroughput(context)), database.getClient().getTracerProvider()); } /** @@ -664,7 +668,7 @@ public Mono readProvisionedThroughput() { * @return a {@link Mono} containing throughput or an error. */ public Mono replaceProvisionedThroughput(int requestUnitsPerSecond) { - return withContext(context -> replaceProvisionedThroughput(requestUnitsPerSecond, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> replaceProvisionedThroughput(requestUnitsPerSecond, context)), database.getClient().getTracerProvider()); } /** @@ -693,16 +697,13 @@ private Mono> deleteItem( RequestOptions requestOptions, Context context) { String spanName = "deleteItem." + this.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), - database.getClient().getServiceEndpoint(), spanName); - Mono> responseMono = this.getDatabase() + Mono> responseMono = this.getDatabase() .getDocClientWrapper() .deleteDocument(getItemLink(itemId), requestOptions) .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponseWithObjectType(response)) .single(); return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, - tracingAttributes, - context, spanName); + context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } private Mono> replaceItem( @@ -712,22 +713,17 @@ private Mono> replaceItem( CosmosItemRequestOptions options, Context context) { String spanName = "replaceItem." + this.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), - database.getClient().getServiceEndpoint(), spanName); Mono> responseMono = this.getDatabase() .getDocClientWrapper() .replaceDocument(getItemLink(itemId), doc, ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) .single(); return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, - tracingAttributes, - context, spanName); + context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } private Mono> upsertItem(T item, CosmosItemRequestOptions options, Context context) { String spanName = "upsertItem." + this.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), - database.getClient().getServiceEndpoint(), spanName); @SuppressWarnings("unchecked") Class itemType = (Class) item.getClass(); Mono> responseMono = this.getDatabase().getDocClientWrapper() @@ -736,8 +732,8 @@ private Mono> upsertItem(T item, CosmosItemReques true) .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) .single(); - return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, tracingAttributes, - context, spanName); + return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, + context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } private Mono> readItem( @@ -745,32 +741,25 @@ private Mono> readItem( RequestOptions requestOptions, Class itemType, Context context) { String spanName = "readItem." + this.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), - database.getClient().getServiceEndpoint(), spanName); Mono> responseMono = this.getDatabase().getDocClientWrapper() .readDocument(getItemLink(itemId), requestOptions) .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) .single(); - return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, tracingAttributes, - context, spanName); + return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, + context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } private Mono read(CosmosContainerRequestOptions options, Context context) { String spanName = "readContainer." + this.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), - database.getClient().getServiceEndpoint(), spanName); Mono responseMono = database.getDocClientWrapper().readCollection(getLink(), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - tracingAttributes, - context, spanName); + context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } private Mono readProvisionedThroughput(Context context) { String spanName = "readProvisionedThroughput." + this.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), - database.getClient().getServiceEndpoint(), spanName); Mono responseMono = this.read() .flatMap(cosmosContainerResponse -> database.getDocClientWrapper() @@ -789,13 +778,11 @@ private Mono readProvisionedThroughput(Context context) { .single(); }).map(cosmosOfferResponse -> cosmosOfferResponse.getResource().getThroughput()); return this.getDatabase().getClient().getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono - , tracingAttributes, context, spanName); + , context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Context context) { String spanName = "replaceProvisionedThroughput." + this.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), - database.getClient().getServiceEndpoint(), spanName); Mono responseMono = this.read() .flatMap(cosmosContainerResponse -> database.getDocClientWrapper() @@ -814,6 +801,6 @@ private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Co return database.getDocClientWrapper().replaceOffer(offer).single(); }).map(offerResourceResponse -> offerResourceResponse.getResource().getThroughput()); return this.getDatabase().getClient().getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono - , tracingAttributes, context, spanName); + , context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java index bcfb8415b06d..68d9c5a2562e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java @@ -24,9 +24,6 @@ import reactor.core.Exceptions; import reactor.core.publisher.Mono; -import java.util.HashMap; -import java.util.Map; - import static com.azure.core.util.FluxUtil.withContext; import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; @@ -82,8 +79,8 @@ public Mono read(CosmosDatabaseRequestOptions optio if (options == null) { options = new CosmosDatabaseRequestOptions(); } - final CosmosDatabaseRequestOptions requestOptions = options; - return withContext(context -> read(requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + final CosmosDatabaseRequestOptions requestOptions = options; + return TracerProvider.cosmosWithContext(withContext(context -> read(requestOptions, context)), client.getTracerProvider()); } /** @@ -115,7 +112,8 @@ public Mono delete(CosmosDatabaseRequestOptions opt } final CosmosDatabaseRequestOptions requestOptions = options; - return withContext(context -> delete(requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> delete(requestOptions, context)), + client.getTracerProvider()); } /* CosmosAsyncContainer operations */ @@ -185,8 +183,9 @@ public Mono createContainer( options = new CosmosContainerRequestOptions(); } - final CosmosContainerRequestOptions requestOptions = options; - return withContext(context -> createContainer(containerProperties, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + final CosmosContainerRequestOptions requestOptions = options; + return TracerProvider.cosmosWithContext(withContext(context -> createContainer(containerProperties, + requestOptions, context)), client.getTracerProvider()); } @@ -336,7 +335,7 @@ public Mono createContainerIfNotExists( private Mono createContainerIfNotExistsInternal( CosmosContainerProperties containerProperties, CosmosAsyncContainer container, CosmosContainerRequestOptions options) { - return withContext(context -> createContainerIfNotExistsInternal(containerProperties, container, options, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> createContainerIfNotExistsInternal(containerProperties, container, options, context)), client.getTracerProvider()); } /** @@ -389,7 +388,7 @@ public CosmosPagedFlux readAllContainers() { * obtained containers or an error. */ public CosmosPagedFlux queryContainers(String query) { - return queryContainers(new SqlQuerySpec(query)); + return queryContainersInternal(false, new SqlQuerySpec(query), new FeedOptions()); } /** @@ -405,7 +404,7 @@ public CosmosPagedFlux queryContainers(String query) * obtained containers or an error. */ public CosmosPagedFlux queryContainers(String query, FeedOptions options) { - return queryContainers(new SqlQuerySpec(query), options); + return queryContainersInternal(false, new SqlQuerySpec(query), options); } /** @@ -420,7 +419,7 @@ public CosmosPagedFlux queryContainers(String query, * obtained containers or an error. */ public CosmosPagedFlux queryContainers(SqlQuerySpec querySpec) { - return queryContainers(querySpec, new FeedOptions()); + return queryContainersInternal(true, querySpec, new FeedOptions()); } /** @@ -435,17 +434,9 @@ public CosmosPagedFlux queryContainers(SqlQuerySpec q * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the * obtained containers or an error. */ - public CosmosPagedFlux queryContainers(SqlQuerySpec querySpec, FeedOptions options) { - return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "queryContainers." + this.getId() + "." + querySpec.getQueryText(); - pagedFluxOptions.setTracerInformation(this.getClient().getTracerProvider(), spanName, - this.getClient().getServiceEndpoint(), getId()); - setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); - return getDocClientWrapper().queryCollections(getLink(), querySpec, options) - .map(response -> BridgeInternal.createFeedResponse( - ModelBridgeInternal.getCosmosContainerPropertiesFromV2Results(response.getResults()), - response.getResponseHeaders())); - }); + public CosmosPagedFlux queryContainers(SqlQuerySpec querySpec + , FeedOptions options) { + return queryContainersInternal(true, querySpec, options); } /** @@ -471,8 +462,8 @@ public CosmosAsyncContainer getContainer(String id) { * created cosmos user or an error. */ public Mono createUser(CosmosUserProperties userProperties) { - return getDocClientWrapper().createUser(this.getLink(), ModelBridgeInternal.getV2User(userProperties), null) - .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, this)).single(); + return TracerProvider.cosmosWithContext(withContext(context -> createUserInternal(userProperties, context)), + client.getTracerProvider()); } @@ -488,8 +479,8 @@ public Mono createUser(CosmosUserProperties userPropert * upserted user or an error. */ public Mono upsertUser(CosmosUserProperties userProperties) { - return getDocClientWrapper().upsertUser(this.getLink(), ModelBridgeInternal.getV2User(userProperties), null) - .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, this)).single(); + return TracerProvider.cosmosWithContext(withContext(context -> upsertUserInternal(userProperties, context)), + client.getTracerProvider()); } /** @@ -558,7 +549,7 @@ public CosmosPagedFlux queryUsers(String query) { * obtained users or an error. */ public CosmosPagedFlux queryUsers(String query, FeedOptions options) { - return queryUsers(new SqlQuerySpec(query), options); + return queryUsersInternal(false, new SqlQuerySpec(query), options); } /** @@ -573,7 +564,7 @@ public CosmosPagedFlux queryUsers(String query, FeedOption * obtained users or an error. */ public CosmosPagedFlux queryUsers(SqlQuerySpec querySpec) { - return queryUsers(querySpec, new FeedOptions()); + return queryUsersInternal(true, querySpec, new FeedOptions()); } /** @@ -589,16 +580,7 @@ public CosmosPagedFlux queryUsers(SqlQuerySpec querySpec) * obtained users or an error. */ public CosmosPagedFlux queryUsers(SqlQuerySpec querySpec, FeedOptions options) { - return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "queryUsers." + this.getId() + "." + querySpec.getQueryText(); - pagedFluxOptions.setTracerInformation(this.getClient().getTracerProvider(), spanName, - this.getClient().getServiceEndpoint(), getId()); - setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); - return getDocClientWrapper().queryUsers(getLink(), querySpec, options) - .map(response -> BridgeInternal.createFeedResponseWithQueryMetrics( - ModelBridgeInternal.getCosmosUserPropertiesFromV2Results(response.getResults()), response.getResponseHeaders(), - ModelBridgeInternal.queryMetrics(response))); - }); + return queryUsersInternal(true, querySpec, options); } /** @@ -617,7 +599,8 @@ public CosmosAsyncUser getUser(String id) { * @return a {@link Mono} containing throughput or an error. */ public Mono readProvisionedThroughput() { - return withContext(context -> readProvisionedThroughput(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> readProvisionedThroughput(context)), + client.getTracerProvider()); } /** @@ -629,7 +612,7 @@ public Mono readProvisionedThroughput() { * @return a {@link Mono} containing throughput or an error. */ public Mono replaceProvisionedThroughput(int requestUnitsPerSecond) { - return withContext(context -> replaceProvisionedThroughput(requestUnitsPerSecond, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> replaceProvisionedThroughput(requestUnitsPerSecond, context)), client.getTracerProvider()); } CosmosAsyncClient getClient() { @@ -652,14 +635,51 @@ String getLink() { return this.link; } + public CosmosPagedFlux queryContainersInternal(boolean isParameterised, SqlQuerySpec querySpec + , FeedOptions options) { + return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName; + if (isParameterised) { + spanName = "queryContainers." + this.getId() + "." + querySpec.getQueryText(); + } else { + spanName = "queryContainers." + this.getId(); + } + + pagedFluxOptions.setTracerInformation(this.getClient().getTracerProvider(), spanName, + this.getClient().getServiceEndpoint(), getId()); + setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); + return getDocClientWrapper().queryCollections(getLink(), querySpec, options) + .map(response -> BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosContainerPropertiesFromV2Results(response.getResults()), + response.getResponseHeaders())); + }); + } + + private CosmosPagedFlux queryUsersInternal(boolean isParameterised, SqlQuerySpec querySpec, FeedOptions options) { + return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName; + if (isParameterised) { + spanName = "queryUsers." + this.getId() + "." + querySpec.getQueryText(); + } else { + spanName = "queryUsers." + this.getId(); + } + + pagedFluxOptions.setTracerInformation(this.getClient().getTracerProvider(), spanName, + this.getClient().getServiceEndpoint(), getId()); + setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); + return getDocClientWrapper().queryUsers(getLink(), querySpec, options) + .map(response -> BridgeInternal.createFeedResponseWithQueryMetrics( + ModelBridgeInternal.getCosmosUserPropertiesFromV2Results(response.getResults()), response.getResponseHeaders(), + ModelBridgeInternal.queryMetrics(response))); + }); + } + private Mono createContainerIfNotExistsInternal( CosmosContainerProperties containerProperties, CosmosAsyncContainer container, CosmosContainerRequestOptions options, Context context) { String spanName = "createContainerIfNotExistsInternal." + containerProperties.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(getId(), - getClient().getServiceEndpoint(), spanName); Mono responseMono = container.read(options).onErrorResume(exception -> { final Throwable unwrappedException = Exceptions.unwrap(exception); if (unwrappedException instanceof CosmosClientException) { @@ -670,8 +690,8 @@ private Mono createContainerIfNotExistsInternal( } return Mono.error(unwrappedException); }); - return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, - spanName); + return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, getId(), getClient().getServiceEndpoint()); } private Mono createContainer( @@ -679,42 +699,34 @@ private Mono createContainer( CosmosContainerRequestOptions options, Context context) { String spanName = "createContainer." + containerProperties.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(getId(), - getClient().getServiceEndpoint(), spanName); Mono responseMono = getDocClientWrapper() .createCollection(this.getLink(), ModelBridgeInternal.getV2Collection(containerProperties), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, this)).single(); - return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, - spanName); + return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, getId(), getClient().getServiceEndpoint()); } private Mono read(CosmosDatabaseRequestOptions options, Context context) { String spanName = "readDatabase." + this.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(getId(), - getClient().getServiceEndpoint(), spanName); Mono responseMono = getDocClientWrapper().readDatabase(getLink(), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single(); - return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, - spanName); + return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, getId(), getClient().getServiceEndpoint()); } private Mono delete(CosmosDatabaseRequestOptions options, Context context) { String spanName = "deleteDatabase." + this.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(getId(), - getClient().getServiceEndpoint(), spanName); Mono responseMono = getDocClientWrapper().deleteDatabase(getLink(), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single(); - return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, - spanName); + return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, getId(), getClient().getServiceEndpoint()); } private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Context context) { String spanName = "replaceOffer." + this.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(getId(), - getClient().getServiceEndpoint(), spanName); Mono responseMono = this.read() .flatMap(cosmosDatabaseResponse -> this.getDocClientWrapper() .queryOffers("select * from c where c.offerResourceId = '" @@ -736,14 +748,12 @@ private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Co }).map(offerResourceResponse -> offerResourceResponse .getResource() .getThroughput())); - return this.client.getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono, tracingAttributes - , context, spanName); + return this.client.getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono, context, spanName + , getId(), getClient().getServiceEndpoint()); } private Mono readProvisionedThroughput(Context context) { String spanName = "readProvisionedThroughput." + this.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(getId(), - getClient().getServiceEndpoint(), spanName); Mono responseMono = this.read() .flatMap(cosmosDatabaseResponse -> getDocClientWrapper() .queryOffers("select * from c where c.offerResourceId = '" @@ -767,7 +777,23 @@ private Mono readProvisionedThroughput(Context context) { }).map(cosmosContainerResponse1 -> cosmosContainerResponse1 .getResource() .getThroughput())); - return this.client.getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono, tracingAttributes - , context, spanName); + return this.client.getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono + , context, spanName, getId(), getClient().getServiceEndpoint()); + } + + private Mono createUserInternal(CosmosUserProperties userProperties, Context context) { + String spanName = "createUser." + this.getId(); + Mono responseMono = getDocClientWrapper().createUser(this.getLink(), ModelBridgeInternal.getV2User(userProperties), null) + .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, this)).single(); + return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, getId(), getClient().getServiceEndpoint()); + } + + private Mono upsertUserInternal(CosmosUserProperties userProperties, Context context) { + String spanName = "upsertUser." + this.getId(); + Mono responseMono = getDocClientWrapper().upsertUser(this.getLink(), ModelBridgeInternal.getV2User(userProperties), null) + .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, this)).single(); + return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, getId(), getClient().getServiceEndpoint()); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java index f4798285b5aa..5075c2ffc293 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java @@ -2,15 +2,19 @@ // 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.TracerProvider; import com.azure.cosmos.models.CosmosAsyncPermissionResponse; import com.azure.cosmos.models.CosmosPermissionProperties; import com.azure.cosmos.models.CosmosPermissionRequestOptions; import com.azure.cosmos.models.ModelBridgeInternal; import reactor.core.publisher.Mono; +import static com.azure.core.util.FluxUtil.withContext; + /** - * Has methods to operate on a per-User Permission to access a specific resource + * Has methods to operate on a per-User Permission to access a specific resource */ public class CosmosAsyncPermission { @@ -56,11 +60,10 @@ public Mono read(CosmosPermissionRequestOptions o if (options == null) { options = new CosmosPermissionRequestOptions(); } - return cosmosUser.getDatabase() - .getDocClientWrapper() - .readPermission(getLink(), ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncPermissionResponse(response, cosmosUser)) - .single(); + + final CosmosPermissionRequestOptions requestOptions = options; + return TracerProvider.cosmosWithContext(withContext(context -> readInternal(requestOptions, context)), + cosmosUser.getDatabase().getClient().getTracerProvider()); } /** @@ -79,12 +82,11 @@ public Mono replace(CosmosPermissionProperties pe if (options == null) { options = new CosmosPermissionRequestOptions(); } - return cosmosUser.getDatabase() - .getDocClientWrapper() - .replacePermission(ModelBridgeInternal.getV2Permissions(permissionSettings), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncPermissionResponse(response, cosmosUser)) - .single(); + + final CosmosPermissionRequestOptions requestOptions = options; + return TracerProvider.cosmosWithContext(withContext(context -> replaceInternal(permissionSettings, + requestOptions, context)), + cosmosUser.getDatabase().getClient().getTracerProvider()); } /** @@ -101,11 +103,10 @@ public Mono delete(CosmosPermissionRequestOptions if (options == null) { options = new CosmosPermissionRequestOptions(); } - return cosmosUser.getDatabase() - .getDocClientWrapper() - .deletePermission(getLink(), ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncPermissionResponse(response, cosmosUser)) - .single(); + + final CosmosPermissionRequestOptions requestOptions = options; + return TracerProvider.cosmosWithContext(withContext(context -> deleteInternal(requestOptions, context)), + cosmosUser.getDatabase().getClient().getTracerProvider()); } String getURIPathSegment() { @@ -125,4 +126,44 @@ String getLink() { builder.append(id()); return builder.toString(); } + + private Mono readInternal(CosmosPermissionRequestOptions options, Context context) { + + String spanName = "readPermission." + cosmosUser.getId(); + Mono responseMono = cosmosUser.getDatabase() + .getDocClientWrapper() + .readPermission(getLink(), ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncPermissionResponse(response, cosmosUser)) + .single(); + return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, cosmosUser.getDatabase().getId(), cosmosUser.getDatabase().getClient().getServiceEndpoint()); + } + + private Mono replaceInternal(CosmosPermissionProperties permissionSettings, + CosmosPermissionRequestOptions options, + Context context) { + + String spanName = "replacePermission." + cosmosUser.getId(); + Mono responseMono = cosmosUser.getDatabase() + .getDocClientWrapper() + .replacePermission(ModelBridgeInternal.getV2Permissions(permissionSettings), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncPermissionResponse(response, cosmosUser)) + .single(); + return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, cosmosUser.getDatabase().getId(), cosmosUser.getDatabase().getClient().getServiceEndpoint()); + } + + private Mono deleteInternal(CosmosPermissionRequestOptions options, + Context context) { + + String spanName = "deletePermission." + cosmosUser.getId(); + Mono responseMono = cosmosUser.getDatabase() + .getDocClientWrapper() + .deletePermission(getLink(), ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncPermissionResponse(response, cosmosUser)) + .single(); + return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, cosmosUser.getDatabase().getId(), cosmosUser.getDatabase().getClient().getServiceEndpoint()); + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java index c4618f20016f..e9fd16f1141c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java @@ -21,8 +21,6 @@ import com.azure.cosmos.util.UtilBridgeInternal; import reactor.core.publisher.Mono; -import java.util.Map; - import static com.azure.core.util.FluxUtil.withContext; import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; @@ -126,7 +124,7 @@ public CosmosPagedFlux readAllStoredProcedures( public CosmosPagedFlux queryStoredProcedures( String query, FeedOptions options) { - return queryStoredProcedures(new SqlQuerySpec(query), options); + return queryStoredProceduresInternal(false, new SqlQuerySpec(query), options); } /** @@ -145,19 +143,7 @@ public CosmosPagedFlux queryStoredProcedures( public CosmosPagedFlux queryStoredProcedures( SqlQuerySpec querySpec, FeedOptions options) { - return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "queryStoredProcedures." + this.container.getId() + "." + querySpec.getQueryText(); - pagedFluxOptions.setTracerInformation(this.container.getDatabase().getClient().getTracerProvider(), - spanName, - this.container.getDatabase().getClient().getServiceEndpoint(), - this.container.getDatabase().getId()); - setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); - return database.getDocClientWrapper() - .queryStoredProcedures(container.getLink(), querySpec, options) - .map(response -> BridgeInternal.createFeedResponse( - ModelBridgeInternal.getCosmosStoredProcedurePropertiesFromV2Results(response.getResults()), - response.getResponseHeaders())); - }); + return queryStoredProceduresInternal(true, querySpec, options); } /** @@ -257,19 +243,7 @@ public CosmosPagedFlux queryUserDefinedFunc public CosmosPagedFlux queryUserDefinedFunctions( SqlQuerySpec querySpec, FeedOptions options) { - return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "queryUserDefinedFunctions." + this.container.getId() + "." + querySpec.getQueryText(); - pagedFluxOptions.setTracerInformation(this.container.getDatabase().getClient().getTracerProvider(), - spanName, - this.container.getDatabase().getClient().getServiceEndpoint(), - this.container.getDatabase().getId()); - setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); - return database.getDocClientWrapper() - .queryUserDefinedFunctions(container.getLink(), querySpec, options) - .map(response -> BridgeInternal.createFeedResponse( - ModelBridgeInternal.getCosmosUserDefinedFunctionPropertiesFromV2Results(response.getResults()), - response.getResponseHeaders())); - }); + return queryUserDefinedFunctionsInternal(true, querySpec, options); } /** @@ -340,7 +314,7 @@ public CosmosPagedFlux readAllTriggers(FeedOptions opti * error. */ public CosmosPagedFlux queryTriggers(String query, FeedOptions options) { - return queryTriggers(new SqlQuerySpec(query), options); + return queryTriggersInternal(false, new SqlQuerySpec(query), options); } /** @@ -356,66 +330,123 @@ public CosmosPagedFlux queryTriggers(String query, Feed * error. */ public CosmosPagedFlux queryTriggers( + SqlQuerySpec querySpec, + FeedOptions options) { + return queryTriggersInternal(true, querySpec, options); + } + + /** + * Gets a CosmosAsyncTrigger object without making a service call + * + * @param id id of the cosmos trigger + * @return a cosmos trigger + */ + public CosmosAsyncTrigger getTrigger(String id) { + return new CosmosAsyncTrigger(id, this.container); + } + + private CosmosPagedFlux queryStoredProceduresInternal( + boolean isParameterised, SqlQuerySpec querySpec, FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "queryTriggers." + this.container.getId(); + String spanName; + if (isParameterised) { + spanName = "queryStoredProcedures." + this.container.getId() + "." + querySpec.getQueryText(); + } else { + spanName = "queryStoredProcedures." + this.container.getId(); + } + pagedFluxOptions.setTracerInformation(this.container.getDatabase().getClient().getTracerProvider(), spanName, this.container.getDatabase().getClient().getServiceEndpoint(), this.container.getDatabase().getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper() - .queryTriggers(container.getLink(), querySpec, options) + .queryStoredProcedures(container.getLink(), querySpec, options) .map(response -> BridgeInternal.createFeedResponse( - ModelBridgeInternal.getCosmosTriggerPropertiesFromV2Results(response.getResults()), + ModelBridgeInternal.getCosmosStoredProcedurePropertiesFromV2Results(response.getResults()), response.getResponseHeaders())); }); } - /** - * Gets a CosmosAsyncTrigger object without making a service call - * - * @param id id of the cosmos trigger - * @return a cosmos trigger - */ - public CosmosAsyncTrigger getTrigger(String id) { - return new CosmosAsyncTrigger(id, this.container); + private CosmosPagedFlux queryUserDefinedFunctionsInternal( + boolean isParameterised, + SqlQuerySpec querySpec, + FeedOptions options) { + return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName; + if (isParameterised) { + spanName = "queryUserDefinedFunctions." + this.container.getId() + "." + querySpec.getQueryText(); + } else { + spanName = "queryUserDefinedFunctions." + this.container.getId(); + } + + pagedFluxOptions.setTracerInformation(this.container.getDatabase().getClient().getTracerProvider(), + spanName, + this.container.getDatabase().getClient().getServiceEndpoint(), + this.container.getDatabase().getId()); + setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); + return database.getDocClientWrapper() + .queryUserDefinedFunctions(container.getLink(), querySpec, options) + .map(response -> BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosUserDefinedFunctionPropertiesFromV2Results(response.getResults()), + response.getResponseHeaders())); + }); + } + + private CosmosPagedFlux queryTriggersInternal( + boolean isParameterised, + SqlQuerySpec querySpec, + FeedOptions options) { + return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + String spanName; + if (isParameterised) { + spanName = "queryTriggers." + this.container.getId() + "." + querySpec.getQueryText(); + } else { + spanName = "queryTriggers." + this.container.getId(); + } + + pagedFluxOptions.setTracerInformation(this.container.getDatabase().getClient().getTracerProvider(), + spanName, + this.container.getDatabase().getClient().getServiceEndpoint(), + this.container.getDatabase().getId()); + setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); + return database.getDocClientWrapper() + .queryTriggers(container.getLink(), querySpec, options) + .map(response -> BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosTriggerPropertiesFromV2Results(response.getResults()), + response.getResponseHeaders())); + }); } private Mono createStoredProcedure(StoredProcedure sProc, CosmosStoredProcedureRequestOptions options, Context context) { String spanName = "createStoredProcedure." + container.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), - database.getClient().getServiceEndpoint(), spanName); Mono responseMono = database.getDocClientWrapper() .createStoredProcedure(container.getLink(), sProc, ModelBridgeInternal.toRequestOptions(options)).map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, this.container)) .single(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } private Mono createUserDefinedFunction( UserDefinedFunction udf, Context context) { String spanName = "createUserDefinedFunction." + container.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), - database.getClient().getServiceEndpoint(), spanName); Mono responseMono = database.getDocClientWrapper() .createUserDefinedFunction(container.getLink(), udf, null).map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, this.container)).single(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } private Mono createTrigger(CosmosTriggerProperties properties, Context context) { String spanName = "createTrigger." + container.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(database.getId(), - database.getClient().getServiceEndpoint(), spanName); Trigger trigger = new Trigger(ModelBridgeInternal.toJsonFromJsonSerializable(properties)); Mono responseMono = database.getDocClientWrapper() .createTrigger(container.getLink(), trigger, null) .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, this.container)) .single(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java index 9a1a95fed2fc..39e064790fa2 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java @@ -12,9 +12,6 @@ import com.azure.cosmos.models.ModelBridgeInternal; import reactor.core.publisher.Mono; -import java.util.HashMap; -import java.util.Map; - import static com.azure.core.util.FluxUtil.withContext; /** @@ -77,7 +74,8 @@ public Mono read() { * @return an {@link Mono} containing the single resource response with the read stored procedure or an error. */ public Mono read(CosmosStoredProcedureRequestOptions options) { - return withContext(context -> read(options, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> read(options, context)), + cosmosContainer.getDatabase().getClient().getTracerProvider()); } /** @@ -106,7 +104,8 @@ public Mono delete() { * @return an {@link Mono} containing the single resource response for the deleted stored procedure or an error. */ public Mono delete(CosmosStoredProcedureRequestOptions options) { - return withContext(context -> delete(options, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> delete(options, context)), + cosmosContainer.getDatabase().getClient().getTracerProvider()); } /** @@ -123,7 +122,8 @@ public Mono delete(CosmosStoredProcedureRequ */ public Mono execute(Object[] procedureParams, CosmosStoredProcedureRequestOptions options) { - return withContext(context -> execute(procedureParams, options, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> execute(procedureParams, options, context)), + cosmosContainer.getDatabase().getClient().getTracerProvider()); } /** @@ -155,7 +155,8 @@ public Mono replace(CosmosStoredProcedurePro */ public Mono replace(CosmosStoredProcedureProperties storedProcedureSettings, CosmosStoredProcedureRequestOptions options) { - return withContext(context -> replace(storedProcedureSettings, options, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> replace(storedProcedureSettings, options, + context)), cosmosContainer.getDatabase().getClient().getTracerProvider()); } String getURIPathSegment() { @@ -179,43 +180,36 @@ String getLink() { private Mono read(CosmosStoredProcedureRequestOptions options, Context context) { String spanName = "readStoredProcedure." + cosmosContainer.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(cosmosContainer.getDatabase().getId(), - cosmosContainer.getDatabase().getClient().getServiceEndpoint(), spanName); if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } Mono responseMono = cosmosContainer.getDatabase().getDocClientWrapper().readStoredProcedure(getLink(), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer)).single(); - return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, + cosmosContainer)).single(); + return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); } private Mono delete(CosmosStoredProcedureRequestOptions options, Context context) { String spanName = "deleteStoredProcedure." + cosmosContainer.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(cosmosContainer.getDatabase().getId(), - cosmosContainer.getDatabase().getClient().getServiceEndpoint(), spanName); - if (options == null) { - options = new CosmosStoredProcedureRequestOptions(); - } if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } + Mono responseMono = cosmosContainer.getDatabase() .getDocClientWrapper() .deleteStoredProcedure(getLink(), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer)) .single(); - return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); + return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); } private Mono execute(Object[] procedureParams, CosmosStoredProcedureRequestOptions options, Context context) { String spanName = "executeStoredProcedure." + cosmosContainer.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(cosmosContainer.getDatabase().getId(), - cosmosContainer.getDatabase().getClient().getServiceEndpoint(), spanName); if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } @@ -225,15 +219,13 @@ private Mono execute(Object[] procedureParam .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer, this.id)) .single(); - return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); + return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); } private Mono replace(CosmosStoredProcedureProperties storedProcedureSettings, CosmosStoredProcedureRequestOptions options, Context context) { String spanName = "replaceStoredProcedure." + cosmosContainer.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(cosmosContainer.getDatabase().getId(), - cosmosContainer.getDatabase().getClient().getServiceEndpoint(), spanName); if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } @@ -243,7 +235,7 @@ private Mono replace(CosmosStoredProcedurePr ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer)) .single(); - return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); + return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java index 5ac774388ee4..e0efc005db0d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java @@ -11,8 +11,6 @@ import com.azure.cosmos.models.ModelBridgeInternal; import reactor.core.publisher.Mono; -import java.util.Map; - import static com.azure.core.util.FluxUtil.withContext; /** @@ -58,7 +56,8 @@ CosmosAsyncTrigger setId(String id) { * @return an {@link Mono} containing the single resource response for the read cosmos trigger or an error. */ public Mono read() { - return withContext(context -> read(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> read(context)), + container.getDatabase().getClient().getTracerProvider()); } /** @@ -72,7 +71,8 @@ public Mono read() { * @return an {@link Mono} containing the single resource response with the replaced cosmos trigger or an error. */ public Mono replace(CosmosTriggerProperties triggerSettings) { - return withContext(context -> replace(triggerSettings, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> replace(triggerSettings, context)), + container.getDatabase().getClient().getTracerProvider()); } /** @@ -85,7 +85,8 @@ public Mono replace(CosmosTriggerProperties triggerS * @return an {@link Mono} containing the single resource response for the deleted cosmos trigger or an error. */ public Mono delete() { - return withContext(context -> delete(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> delete(context)), + container.getDatabase().getClient().getTracerProvider()); } String getURIPathSegment() { @@ -108,38 +109,31 @@ String getLink() { private Mono read(Context context) { String spanName = "readTrigger." + container.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(container.getDatabase().getId(), - container.getDatabase().getClient().getServiceEndpoint(), spanName); - Mono responseMono = container.getDatabase() .getDocClientWrapper() .readTrigger(getLink(), null) .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, container)) .single(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); } private Mono replace(CosmosTriggerProperties triggerSettings, Context context) { String spanName = "replaceTrigger." + container.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(container.getDatabase().getId(), - container.getDatabase().getClient().getServiceEndpoint(), spanName); Mono responseMono = container.getDatabase() .getDocClientWrapper() .replaceTrigger(new Trigger(ModelBridgeInternal.toJsonFromJsonSerializable(triggerSettings)), null) .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, container)) .single(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); } private Mono delete(Context context) { String spanName = "deleteTrigger." + container.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(container.getDatabase().getId(), - container.getDatabase().getClient().getServiceEndpoint(), spanName); Mono responseMono = container.getDatabase() .getDocClientWrapper() .deleteTrigger(getLink(), null) .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, container)) .single(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java index 528eb57cab2b..f51e3e61af9c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java @@ -3,7 +3,10 @@ package com.azure.cosmos; +import com.azure.core.util.Context; import com.azure.cosmos.implementation.Paths; +import com.azure.cosmos.implementation.Permission; +import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.models.CosmosAsyncPermissionResponse; import com.azure.cosmos.models.CosmosAsyncUserResponse; import com.azure.cosmos.models.CosmosPermissionProperties; @@ -11,11 +14,11 @@ import com.azure.cosmos.models.CosmosUserProperties; import com.azure.cosmos.models.FeedOptions; import com.azure.cosmos.models.ModelBridgeInternal; -import com.azure.cosmos.implementation.Permission; import com.azure.cosmos.util.CosmosPagedFlux; import com.azure.cosmos.util.UtilBridgeInternal; import reactor.core.publisher.Mono; +import static com.azure.core.util.FluxUtil.withContext; import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; /** @@ -56,9 +59,8 @@ CosmosAsyncUser setId(String id) { * @return a {@link Mono} containing the single resource response with the read user or an error. */ public Mono read() { - return this.database.getDocClientWrapper() - .readUser(getLink(), null) - .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, database)).single(); + return TracerProvider.cosmosWithContext(withContext(context -> readInternal(context)), + database.getClient().getTracerProvider()); } /** @@ -68,9 +70,8 @@ public Mono read() { * @return a {@link Mono} containing the single resource response with the replaced user or an error. */ public Mono replace(CosmosUserProperties userSettings) { - return this.database.getDocClientWrapper() - .replaceUser(ModelBridgeInternal.getV2User(userSettings), null) - .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, database)).single(); + return TracerProvider.cosmosWithContext(withContext(context -> replaceInternal(userSettings, context)), + database.getClient().getTracerProvider()); } /** @@ -79,9 +80,8 @@ public Mono replace(CosmosUserProperties userSettings) * @return a {@link Mono} containing the single resource response with the deleted user or an error. */ public Mono delete() { - return this.database.getDocClientWrapper() - .deleteUser(getLink(), null) - .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, database)).single(); + return TracerProvider.cosmosWithContext(withContext(context -> deleteInternal(context)), + database.getClient().getTracerProvider()); } /** @@ -101,11 +101,11 @@ public Mono createPermission( if (options == null) { options = new CosmosPermissionRequestOptions(); } + Permission permission = ModelBridgeInternal.getV2Permissions(permissionSettings); - return database.getDocClientWrapper() - .createPermission(getLink(), permission, ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncPermissionResponse(response, this)) - .single(); + final CosmosPermissionRequestOptions requesOptions = options; + return TracerProvider.cosmosWithContext(withContext(context -> createPermissionInternal(permission, requesOptions, context)), + database.getClient().getTracerProvider()); } /** @@ -126,10 +126,11 @@ public Mono upsertPermission( if (options == null) { options = new CosmosPermissionRequestOptions(); } - return database.getDocClientWrapper() - .upsertPermission(getLink(), permission, ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncPermissionResponse(response, this)) - .single(); + + final CosmosPermissionRequestOptions requestOptions = options; + return TracerProvider.cosmosWithContext(withContext(context -> upsertPermissionInternal(permission, + requestOptions, context)), + database.getClient().getTracerProvider()); } @@ -189,7 +190,7 @@ public CosmosPagedFlux queryPermissions(String query */ public CosmosPagedFlux queryPermissions(String query, FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "queryPermissions." + this.getId() + "." + query; + String spanName = "queryPermissions." + this.getId(); pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), spanName, this.getDatabase().getClient().getServiceEndpoint(), @@ -239,4 +240,57 @@ String getLink() { public CosmosAsyncDatabase getDatabase() { return database; } + + private Mono readInternal(Context context) { + String spanName = "readUser." + getId(); + Mono responseMono = this.database.getDocClientWrapper() + .readUser(getLink(), null) + .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, database)).single(); + return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, database.getId(), database.getClient().getServiceEndpoint()); + } + + private Mono replaceInternal(CosmosUserProperties userSettings, Context context) { + String spanName = "replaceUser." + getId(); + Mono responseMono = this.database.getDocClientWrapper() + .replaceUser(ModelBridgeInternal.getV2User(userSettings), null) + .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, database)).single(); + return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, database.getId(), database.getClient().getServiceEndpoint()); + } + + private Mono deleteInternal(Context context) { + String spanName = "deleteUser." + getId(); + Mono responseMono = this.database.getDocClientWrapper() + .deleteUser(getLink(), null) + .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, database)).single(); + return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, database.getId(), database.getClient().getServiceEndpoint()); + } + + private Mono createPermissionInternal( + Permission permission, + CosmosPermissionRequestOptions options, + Context context) { + String spanName = "createPermission." + getId(); + Mono responseMono = database.getDocClientWrapper() + .createPermission(getLink(), permission, ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncPermissionResponse(response, this)) + .single(); + return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, database.getId(), database.getClient().getServiceEndpoint()); + } + + private Mono upsertPermissionInternal( + Permission permission, + CosmosPermissionRequestOptions options, + Context context) { + String spanName = "upsertPermission." + getId(); + Mono responseMono = database.getDocClientWrapper() + .upsertPermission(getLink(), permission, ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncPermissionResponse(response, this)) + .single(); + return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, database.getId(), database.getClient().getServiceEndpoint()); + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java index 0d3ed2d6575a..6137dfb09468 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java @@ -11,9 +11,6 @@ import com.azure.cosmos.models.ModelBridgeInternal; import reactor.core.publisher.Mono; -import java.util.Map; -import java.util.concurrent.atomic.AtomicReference; - import static com.azure.core.util.FluxUtil.withContext; /** @@ -61,7 +58,8 @@ CosmosAsyncUserDefinedFunction setId(String id) { * @return an {@link Mono} containing the single resource response for the read user defined function or an error. */ public Mono read() { - return withContext(context -> read(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> read(context)), + container.getDatabase().getClient().getTracerProvider()); } /** @@ -77,7 +75,8 @@ public Mono read() { * or an error. */ public Mono replace(CosmosUserDefinedFunctionProperties udfSettings) { - return withContext(context -> replace(udfSettings, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> replace(udfSettings, context)), + container.getDatabase().getClient().getTracerProvider()); } /** @@ -92,7 +91,8 @@ public Mono replace(CosmosUserDefinedFun * an error. */ public Mono delete() { - return withContext(context -> delete(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> delete(context)), + container.getDatabase().getClient().getTracerProvider()); } String getURIPathSegment() { @@ -115,36 +115,30 @@ String getLink() { private Mono read(Context context) { String spanName = "readUDF." + container.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(container.getDatabase().getId(), - container.getDatabase().getClient().getServiceEndpoint(), spanName); Mono responseMono = container.getDatabase().getDocClientWrapper().readUserDefinedFunction(getLink(), null) .map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, container)).single(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); } private Mono replace(CosmosUserDefinedFunctionProperties udfSettings, Context context) { String spanName = "replaceUDF." + container.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(container.getDatabase().getId(), - container.getDatabase().getClient().getServiceEndpoint(), spanName); Mono responseMono = container.getDatabase() .getDocClientWrapper() .replaceUserDefinedFunction(new UserDefinedFunction(ModelBridgeInternal.toJsonFromJsonSerializable(udfSettings)), null) .map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, container)) .single(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); } private Mono delete(Context context) { String spanName = "deleteUDF." + container.getId(); - Map tracingAttributes = TracerProvider.createTracingMap(container.getDatabase().getId(), - container.getDatabase().getClient().getServiceEndpoint(), spanName); Mono responseMono = container.getDatabase() .getDocClientWrapper() .deleteUserDefinedFunction(this.getLink(), null) .map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, container)) .single(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, tracingAttributes, context, spanName); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosPagedFluxOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosPagedFluxOptions.java index 1f0a3904db47..28b4054180dc 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosPagedFluxOptions.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosPagedFluxOptions.java @@ -15,9 +15,10 @@ public class CosmosPagedFluxOptions { private String requestContinuation; private Integer maxItemCount; - private Map tracingAttributes; private TracerProvider tracerProvider; private String tracerSpanName; + private String databaseId; + private String serviceEndpoint; public CosmosPagedFluxOptions() {} @@ -64,14 +65,6 @@ public CosmosPagedFluxOptions setMaxItemCount(Integer maxItemCount) { return this; } - /** - * Gets the tracing attributes - * @return tracingAttributes - */ - public Map getTracingAttributes() { - return this.tracingAttributes; - } - /** * Gets the tracer provider * @return tracerProvider @@ -88,8 +81,25 @@ public String getTracerSpanName() { return tracerSpanName; } - public void setTracerInformation(TracerProvider tracerProvider, String tracerSpanName, String serviceUrl, String databaseId) { - this.tracingAttributes = TracerProvider.createTracingMap(databaseId, serviceUrl, tracerSpanName); + /** + * Gets the databaseId + * @return databaseId + */ + public String getDatabaseId() { + return databaseId; + } + + /** + * Gets the service end point + * @return serviceEndpoint + */ + public String getServiceEndpoint() { + return serviceEndpoint; + } + + public void setTracerInformation(TracerProvider tracerProvider, String tracerSpanName, String serviceEndpoint, String databaseId) { + this.databaseId = databaseId; + this.serviceEndpoint = serviceEndpoint; this.tracerSpanName = tracerSpanName; this.tracerProvider = tracerProvider; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java index 376099c2d697..1733bef533ba 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java @@ -15,9 +15,7 @@ import java.io.PrintWriter; import java.io.StringWriter; import java.util.ArrayList; -import java.util.HashMap; import java.util.List; -import java.util.Map; import java.util.Objects; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; @@ -37,7 +35,7 @@ public class TracerProvider { public static final String COSMOS_CALL_DEPTH = "cosmosCallDepth"; public static final String MICROSOFT_DOCOMENTDB = "Microsoft.DocumentDB"; - public final static Function callDepthAttributeFunc = ( + private final static Function callDepthAttributeFunc = ( reactor.util.context.Context reactorContext) -> { CallDepth callDepth = reactorContext.getOrDefault(COSMOS_CALL_DEPTH, CallDepth.ZERO); assert callDepth != null; @@ -50,17 +48,12 @@ public class TracerProvider { } }; - public static Map createTracingMap(String databaseId, String serviceEndpoint, String spanName) { - return new HashMap() {{ - if (databaseId != null) { - put(TracerProvider.DB_INSTANCE, databaseId); - } - - put(AZ_TRACING_NAMESPACE_KEY, MICROSOFT_DOCOMENTDB); - put(TracerProvider.DB_TYPE, DB_TYPE_VALUE); - put(TracerProvider.DB_URL, serviceEndpoint); - put(TracerProvider.DB_STATEMENT, spanName); - }}; + public static Mono cosmosWithContext(Mono mono, TracerProvider tracerProvider) { + if(tracerProvider.isEnabled()) { + return mono.subscriberContext(TracerProvider.callDepthAttributeFunc); + } else { + return mono; + } } public TracerProvider(Iterable tracers) { @@ -82,14 +75,18 @@ public boolean isEnabled() { * @param context Additional metadata that is passed through the call stack. * @return An updated context object. */ - public Context startSpan(String methodName, Context context, Map attributeMap) { + public Context startSpan(String methodName, String databaseId, String endpoint, Context context) { Context local = Objects.requireNonNull(context, "'context' cannot be null."); for (Tracer tracer : tracers) { local = tracer.start(methodName, local); // start the span and return the started span - Context lambdaContext = local; - attributeMap.forEach((key, value) -> { - tracer.setAttribute(key, value, lambdaContext); // set attrs on the span - }); + if (databaseId != null) { + tracer.setAttribute(TracerProvider.DB_INSTANCE, databaseId, local); + } + + tracer.setAttribute(AZ_TRACING_NAMESPACE_KEY, MICROSOFT_DOCOMENTDB, local); + tracer.setAttribute(TracerProvider.DB_TYPE, DB_TYPE_VALUE, local); + tracer.setAttribute(TracerProvider.DB_URL, endpoint, local); + tracer.setAttribute(TracerProvider.DB_STATEMENT, methodName, local); } return local; @@ -131,54 +128,57 @@ public > void endSpan(Context conte } public > Mono traceEnabledCosmosResponsePublisher(Mono resultPublisher, - Map tracingAttributes, Context context, - String spanName) { - return traceEnabledPublisher(resultPublisher, tracingAttributes, context, spanName, + String spanName, + String databaseId, + String endpoint) { + return traceEnabledPublisher(resultPublisher, context, spanName,databaseId, endpoint, CosmosResponse::getStatusCode); } public Mono traceEnabledNonCosmosResponsePublisher(Mono resultPublisher, - Map tracingAttributes, Context context, - String spanName) { - return traceEnabledPublisher(resultPublisher, tracingAttributes, context, spanName, (T response) -> 200); + String spanName, + String databaseId, + String endpoint) { + return traceEnabledPublisher(resultPublisher, context, spanName, databaseId, endpoint, (T response) -> 200); } public Mono> traceEnabledCosmosItemResponsePublisher(Mono> resultPublisher, - Map tracingAttributes, Context context, - String spanName) { - return traceEnabledPublisher(resultPublisher, tracingAttributes, context, spanName, + String spanName, + String databaseId, + String endpoint) { + return traceEnabledPublisher(resultPublisher, context, spanName,databaseId, endpoint, CosmosAsyncItemResponse::getStatusCode); } public Mono traceEnabledPublisher(Mono resultPublisher, - Map tracingAttributes, Context context, String spanName, + String databaseId, + String endpoint, Function statusCodeFunc) { final boolean isTracingEnabled = this.isEnabled(); final AtomicReference parentContext = isTracingEnabled ? new AtomicReference<>(Context.NONE) : null; + CallDepth callDepth = FluxUtil.toReactorContext(context).getOrDefault(COSMOS_CALL_DEPTH, + CallDepth.ZERO); + assert callDepth != null; + final boolean isNestedCall = callDepth.equals(CallDepth.TWO_OR_MORE); return resultPublisher .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled) { - CallDepth callDepth = FluxUtil.toReactorContext(context).getOrDefault(COSMOS_CALL_DEPTH, - CallDepth.ZERO); - assert callDepth != null; - if (!callDepth.equals(CallDepth.TWO_OR_MORE)) { - parentContext.set(this.startSpan(spanName, - context, tracingAttributes)); - } + if (isTracingEnabled && !isNestedCall) { + parentContext.set(this.startSpan(spanName, databaseId, endpoint, + context)); } }).doOnSuccess(response -> { - if (isTracingEnabled) { + if (isTracingEnabled && !isNestedCall) { this.endSpan(parentContext.get(), Signal.complete(), statusCodeFunc.apply(response)); } }).doOnError(throwable -> { - if (isTracingEnabled) { + if (isTracingEnabled && !isNestedCall) { this.endSpan(parentContext.get(), Signal.error(throwable), 0); } }); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java index 998cb6d7c42c..f1a8e61f6042 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java @@ -17,6 +17,8 @@ import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; +import static com.azure.cosmos.implementation.TracerProvider.COSMOS_CALL_DEPTH; + /** * Cosmos implementation of {@link ContinuablePagedFlux}. *

@@ -91,19 +93,17 @@ public void subscribe(CoreSubscriber coreSubscriber) { private Flux> byPage(CosmosPagedFluxOptions pagedFluxOptions, Context context) { final AtomicReference parentContext = new AtomicReference<>(Context.NONE); - + TracerProvider.CallDepth callDepth = FluxUtil.toReactorContext(context).getOrDefault(COSMOS_CALL_DEPTH, + TracerProvider.CallDepth.ZERO); + assert callDepth != null; + final boolean isNestedCall = callDepth.equals(TracerProvider.CallDepth.TWO_OR_MORE); return this.optionsFluxFunction.apply(pagedFluxOptions).doOnSubscribe(ignoredValue -> { - if (pagedFluxOptions.getTracerProvider().isEnabled()) { - TracerProvider.CallDepth callDepth = FluxUtil.toReactorContext(context).getOrDefault(TracerProvider.COSMOS_CALL_DEPTH, - TracerProvider.CallDepth.ZERO); - assert callDepth != null; - if (!callDepth.equals(TracerProvider.CallDepth.TWO_OR_MORE)) { - parentContext.set(pagedFluxOptions.getTracerProvider().startSpan(pagedFluxOptions.getTracerSpanName(), - context, pagedFluxOptions.getTracingAttributes())); - } + if ( pagedFluxOptions.getTracerProvider().isEnabled() && !isNestedCall) { + parentContext.set(pagedFluxOptions.getTracerProvider().startSpan(pagedFluxOptions.getTracerSpanName(), pagedFluxOptions.getDatabaseId(), pagedFluxOptions.getServiceEndpoint(), + context)); } }).doOnComplete(() -> { - if (pagedFluxOptions.getTracerProvider().isEnabled()) { + if ( pagedFluxOptions.getTracerProvider().isEnabled() && !isNestedCall) { pagedFluxOptions.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), 200); } }).doOnError(throwable -> { diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java index 3bff6207d233..f53375fcb012 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java @@ -51,22 +51,22 @@ public void cosmosAsyncClient() { ReflectionUtils.setTracerProvider(client, tracer); client.createDatabaseIfNotExists(cosmosAsyncDatabase.getId()).block(); - Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); Mockito.verify(tracer, Mockito.times(2)).endSpan(Matchers.any(Context.class), Matchers.>>any(), Matchers.anyInt()); //endSpan invocation is 1 greater than startSpan, because createDatabaseIfNotExists // uses public read api client.readAllDatabases(new FeedOptions()).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); Mockito.verify(tracer, Mockito.atLeast(3)).endSpan(Matchers.any(Context.class), Matchers.>>any(), Matchers.anyInt()); //We used atLeast in endSpan because for query we are calling endSpan on every doOnEach String query = "select * from c where c.id = '" + cosmosAsyncDatabase.getId() + "'"; client.queryDatabases(query, new FeedOptions()).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); Mockito.verify(tracer, Mockito.atLeast(4)).endSpan(Matchers.any(Context.class), Matchers.>>any(), Matchers.anyInt()); //We used atLeast in endSpan because for query we are calling endSpan on every doOnEach } @@ -78,8 +78,8 @@ public void cosmosAsyncDatabase() { cosmosAsyncDatabase.createContainerIfNotExists(cosmosAsyncContainer.getId(), "/pk", 5000).block(); - Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); try { cosmosAsyncDatabase.readProvisionedThroughput().block(); @@ -87,16 +87,16 @@ public void cosmosAsyncDatabase() { //do nothing } - Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncDatabase.readAllUsers().byPage().single().block(); - Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncDatabase.readAllContainers().byPage().single().block(); - Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) @@ -105,45 +105,44 @@ public void cosmosAsyncContainer() { ReflectionUtils.setTracerProvider(client, tracer); cosmosAsyncContainer.read().block(); - Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); try { cosmosAsyncContainer.readProvisionedThroughput().block(); } catch (CosmosClientException ex) { //do nothing } - Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); CosmosItemProperties properties = new CosmosItemProperties(); properties.setId(ITEM_ID); cosmosAsyncContainer.createItem(properties).block(); - Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.upsertItem(properties, new CosmosItemRequestOptions()).block(); - Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.readItem(ITEM_ID, PartitionKey.NONE, CosmosItemProperties.class).block(); - Mockito.verify(tracer, Mockito.times(5)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(5)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.deleteItem(ITEM_ID, PartitionKey.NONE).block(); - Mockito.verify(tracer, Mockito.times(6)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); - + Mockito.verify(tracer, Mockito.times(6)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.readAllItems(new FeedOptions(), CosmosItemRequestOptions.class).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(7)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(7)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); String query = "select * from c where c.id = '" + ITEM_ID + "'"; cosmosAsyncContainer.queryItems(query, new FeedOptions(), CosmosItemRequestOptions.class).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(8)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(8)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) @@ -152,82 +151,81 @@ public void cosmosAsyncScripts() { ReflectionUtils.setTracerProvider(client, tracer); cosmosAsyncContainer.getScripts().readAllStoredProcedures(new FeedOptions()).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.getScripts().readAllTriggers(new FeedOptions()).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.getScripts().readAllUserDefinedFunctions(new FeedOptions()).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); CosmosUserDefinedFunctionProperties cosmosUserDefinedFunctionProperties = getCosmosUserDefinedFunctionProperties(); CosmosUserDefinedFunctionProperties resultUdf = cosmosAsyncContainer.getScripts().createUserDefinedFunction(cosmosUserDefinedFunctionProperties).block().getProperties(); - Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.getScripts().getUserDefinedFunction(cosmosUserDefinedFunctionProperties.getId()).read().block(); - Mockito.verify(tracer, Mockito.times(5)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(5)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosUserDefinedFunctionProperties.setBody("function() {var x = 15;}"); cosmosAsyncContainer.getScripts().getUserDefinedFunction(resultUdf.getId()).replace(resultUdf).block(); - Mockito.verify(tracer, Mockito.times(6)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(6)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.getScripts().readAllUserDefinedFunctions(new FeedOptions()).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(7)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(7)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.getScripts().getUserDefinedFunction(cosmosUserDefinedFunctionProperties.getId()).delete().block(); - Mockito.verify(tracer, Mockito.times(8)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(8)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); CosmosTriggerProperties cosmosTriggerProperties = getCosmosTriggerProperties(); CosmosTriggerProperties resultTrigger = cosmosAsyncContainer.getScripts().createTrigger(cosmosTriggerProperties).block().getProperties(); - Mockito.verify(tracer, Mockito.times(9)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(9)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.getScripts().getTrigger(cosmosTriggerProperties.getId()).read().block(); - Mockito.verify(tracer, Mockito.times(10)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(10)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.getScripts().getTrigger(cosmosTriggerProperties.getId()).replace(resultTrigger).block(); - Mockito.verify(tracer, Mockito.times(11)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); - + Mockito.verify(tracer, Mockito.times(11)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.getScripts().readAllTriggers(new FeedOptions()).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(12)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(12)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.getScripts().getTrigger(cosmosTriggerProperties.getId()).delete().block(); - Mockito.verify(tracer, Mockito.times(13)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(13)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); CosmosStoredProcedureProperties procedureProperties = getCosmosStoredProcedureProperties(); CosmosStoredProcedureProperties resultSproc = cosmosAsyncContainer.getScripts().createStoredProcedure(procedureProperties).block().getProperties(); - Mockito.verify(tracer, Mockito.times(14)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(14)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).read().block(); - Mockito.verify(tracer, Mockito.times(15)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(15)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).replace(resultSproc).block(); - Mockito.verify(tracer, Mockito.times(16)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(16)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.getScripts().readAllStoredProcedures(new FeedOptions()).byPage().single().block(); cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).delete().block(); - Mockito.verify(tracer, Mockito.times(18)).startSpan(Matchers.anyString(), Matchers.any(Context.class), - Matchers.anyMapOf(String.class, String.class)); + Mockito.verify(tracer, Mockito.times(18)).startSpan(Matchers.anyString(), Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); } @AfterClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT) From 1f9f1397b3ef30c383800bd32e6ebc1986c2567b Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Wed, 6 May 2020 12:37:16 -0400 Subject: [PATCH 11/40] fixing build error --- .../main/java/com/azure/cosmos/CosmosAsyncContainer.java | 2 +- .../main/java/com/azure/cosmos/CosmosAsyncScripts.java | 9 ++++++--- .../com/azure/cosmos/implementation/TracerProvider.java | 2 +- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 2c9df7932ae9..842cc54c8df3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -88,7 +88,7 @@ public Mono read(CosmosContainerRequestOptions opt } final CosmosContainerRequestOptions requestOptions = options; - return TracerProvider.cosmosWithContext(withContext(context -> read(requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc), database.getClient().getTracerProvider()); + return TracerProvider.cosmosWithContext(withContext(context -> read(requestOptions, context)), database.getClient().getTracerProvider()); } /** diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java index e9fd16f1141c..06f47bc15c76 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java @@ -75,7 +75,8 @@ public Mono createStoredProcedure( sProc.setId(properties.getId()); sProc.setBody(properties.getBody()); final CosmosStoredProcedureRequestOptions requestOptions = options; - return withContext(context -> createStoredProcedure(sProc, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> createStoredProcedure(sProc, requestOptions, context)), + container.getDatabase().getClient().getTracerProvider()); } @@ -175,7 +176,8 @@ public Mono createUserDefinedFunction( udf.setId(properties.getId()); udf.setBody(properties.getBody()); - return withContext(context -> createUserDefinedFunction(udf, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> createUserDefinedFunction(udf, context)), + container.getDatabase().getClient().getTracerProvider()); } /** @@ -269,7 +271,8 @@ public CosmosAsyncUserDefinedFunction getUserDefinedFunction(String id) { * @return an {@link Mono} containing the single resource response with the created trigger or an error. */ public Mono createTrigger(CosmosTriggerProperties properties) { - return withContext(context -> createTrigger(properties, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); + return TracerProvider.cosmosWithContext(withContext(context -> createTrigger(properties, context)), + container.getDatabase().getClient().getTracerProvider()); } /** diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java index 1733bef533ba..078c297daf29 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java @@ -35,7 +35,7 @@ public class TracerProvider { public static final String COSMOS_CALL_DEPTH = "cosmosCallDepth"; public static final String MICROSOFT_DOCOMENTDB = "Microsoft.DocumentDB"; - private final static Function callDepthAttributeFunc = ( + public final static Function callDepthAttributeFunc = ( reactor.util.context.Context reactorContext) -> { CallDepth callDepth = reactorContext.getOrDefault(COSMOS_CALL_DEPTH, CallDepth.ZERO); assert callDepth != null; From d14a0bb5ee30f455dd35b545f564ec6b0ee5beb7 Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Mon, 11 May 2020 10:29:15 -0400 Subject: [PATCH 12/40] resolving merge conflict --- .../com/azure/cosmos/CosmosAsyncClient.java | 33 +++++++ .../com/azure/cosmos/CosmosAsyncDatabase.java | 96 +++++++++++++++++++ .../com/azure/cosmos/CosmosClientBuilder.java | 44 ++++++++- 3 files changed, 169 insertions(+), 4 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index f803ddf91dcb..a7dd95664e07 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -18,6 +18,7 @@ import com.azure.cosmos.models.FeedOptions; import com.azure.cosmos.models.ModelBridgeInternal; import com.azure.cosmos.models.SqlQuerySpec; +import com.azure.cosmos.models.ThroughputProperties; import com.azure.cosmos.util.CosmosPagedFlux; import com.azure.cosmos.util.UtilBridgeInternal; import io.micrometer.core.instrument.MeterRegistry; @@ -53,6 +54,7 @@ public final class CosmosAsyncClient implements Closeable { private final boolean sessionCapturingOverride; private final boolean enableTransportClientSharing; private final TracerProvider tracerProvider; + private final boolean contentResponseOnWriteEnabled; CosmosAsyncClient(CosmosClientBuilder builder) { this.configs = builder.configs(); @@ -66,6 +68,7 @@ public final class CosmosAsyncClient implements Closeable { this.sessionCapturingOverride = builder.isSessionCapturingOverrideEnabled(); this.enableTransportClientSharing = builder.isConnectionReuseAcrossClientsEnabled(); this.tracerProvider = builder.getTracerProvider(); + this.contentResponseOnWriteEnabled = builder.isContentResponseOnWriteEnabled(); this.asyncDocumentClient = new AsyncDocumentClient.Builder() .withServiceEndpoint(this.serviceEndpoint) .withMasterKeyOrResourceToken(this.keyOrResourceToken) @@ -76,6 +79,7 @@ public final class CosmosAsyncClient implements Closeable { .withTokenResolver(this.cosmosAuthorizationTokenResolver) .withCosmosKeyCredential(this.cosmosKeyCredential) .withTransportClientSharing(this.enableTransportClientSharing) + .withContentResponseOnWriteEnabled(this.contentResponseOnWriteEnabled) .build(); } @@ -168,6 +172,22 @@ CosmosKeyCredential cosmosKeyCredential() { return cosmosKeyCredential; } + /** + * Gets the boolean which indicates whether to only return the headers and status code in Cosmos DB response + * in case of Create, Update and Delete operations on CosmosItem. + * + * If set to false (which is by default), this removes the resource from response. It reduces networking + * and CPU load by not sending the resource back over the network and serializing it + * on the client. + * + * By-default, this is false. + * + * @return a boolean indicating whether resource will be included in the response or not + */ + boolean isContentResponseOnWriteEnabled() { + return contentResponseOnWriteEnabled; + } + /** * CREATE a Database if it does not already exist on the service *

@@ -312,6 +332,19 @@ public Mono createDatabase(String id, int throughpu return createDatabase(new CosmosDatabaseProperties(id), options); } + /** + * Creates a database. + * + * @param id the id + * @param throughputProperties the throughputProperties + * @return the mono + */ + public Mono createDatabase(String id, ThroughputProperties throughputProperties) { + CosmosDatabaseRequestOptions options = new CosmosDatabaseRequestOptions(); + ModelBridgeInternal.setOfferProperties(options, throughputProperties); + return createDatabase(new CosmosDatabaseProperties(id), options); + } + /** * Reads all databases. *

diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java index 68d9c5a2562e..3dd11d243857 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java @@ -14,17 +14,24 @@ import com.azure.cosmos.models.CosmosAsyncUserResponse; import com.azure.cosmos.models.CosmosContainerProperties; import com.azure.cosmos.models.CosmosContainerRequestOptions; +import com.azure.cosmos.models.CosmosDatabaseProperties; import com.azure.cosmos.models.CosmosDatabaseRequestOptions; import com.azure.cosmos.models.CosmosUserProperties; import com.azure.cosmos.models.FeedOptions; import com.azure.cosmos.models.ModelBridgeInternal; +import com.azure.cosmos.models.SqlParameter; import com.azure.cosmos.models.SqlQuerySpec; +import com.azure.cosmos.models.ThroughputProperties; +import com.azure.cosmos.models.ThroughputResponse; import com.azure.cosmos.util.CosmosPagedFlux; import com.azure.cosmos.util.UtilBridgeInternal; import reactor.core.Exceptions; import reactor.core.publisher.Mono; import static com.azure.core.util.FluxUtil.withContext; +import java.util.Collections; +import java.util.List; + import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; /** @@ -159,6 +166,22 @@ public Mono createContainer( return createContainer(containerProperties, options); } + /** + * Creates a container. + * + * @param containerProperties the container properties + * @param throughputProperties the throughput properties + * @param options the request options + * @return the mono + */ + public Mono createContainer( + CosmosContainerProperties containerProperties, + ThroughputProperties throughputProperties, + CosmosContainerRequestOptions options){ + ModelBridgeInternal.setOfferProperties(options, throughputProperties); + return createContainer(containerProperties, options); + } + /** * Creates a document container. *

@@ -615,6 +638,79 @@ public Mono replaceProvisionedThroughput(int requestUnitsPerSecond) { return TracerProvider.cosmosWithContext(withContext(context -> replaceProvisionedThroughput(requestUnitsPerSecond, context)), client.getTracerProvider()); } + /** + * Sets throughput provisioned for a container in measurement of + * Requests-per-Unit in the Azure Cosmos service. + * + * @param throughputProperties the throughput properties + * @return the mono + */ + public Mono replaceThroughput(ThroughputProperties throughputProperties) { + return this.read() + .flatMap(response -> this.getDocClientWrapper() + .queryOffers(getOfferQuerySpecFromResourceId(response.getProperties() + .getResourceId()), + new FeedOptions()) + .single() + .flatMap(offerFeedResponse -> { + if (offerFeedResponse.getResults().isEmpty()) { + return Mono.error(BridgeInternal + .createCosmosClientException( + HttpConstants.StatusCodes.BADREQUEST, + "No offers found for the " + + "resource " + this.getId())); + } + + Offer existingOffer = offerFeedResponse.getResults().get(0); + Offer updatedOffer = + ModelBridgeInternal.updateOfferFromProperties(existingOffer, + throughputProperties); + + return this.getDocClientWrapper() + .replaceOffer(updatedOffer) + .single(); + }) + .map(ModelBridgeInternal::createThroughputRespose)); + } + + /** + * Gets the throughput of the database + * + * @return the mono containing throughput response + */ + public Mono readThroughput() { + return this.read() + .flatMap(response -> getDocClientWrapper() + .queryOffers(getOfferQuerySpecFromResourceId(response.getProperties() + .getResourceId()), + new FeedOptions()) + .single() + .flatMap(offerFeedResponse -> { + if (offerFeedResponse.getResults().isEmpty()) { + return Mono.error(BridgeInternal + .createCosmosClientException( + HttpConstants.StatusCodes.BADREQUEST, + "No offers found for the " + + "resource " + this.getId())); + } + return getDocClientWrapper() + .readOffer(offerFeedResponse.getResults() + .get(0) + .getSelfLink()) + .single(); + }) + .map(ModelBridgeInternal::createThroughputRespose)); + } + + SqlQuerySpec getOfferQuerySpecFromResourceId(String resourceId) { + String queryText = "select * from c where c.offerResourceId = @resourceId"; + SqlQuerySpec querySpec = new SqlQuerySpec(queryText); + List parameters = Collections + .singletonList(new SqlParameter("@resourceId", resourceId)); + querySpec.setParameters(parameters); + return querySpec; + } + CosmosAsyncClient getClient() { return client; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java index 9724f652fe5a..398b08cf65de 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java @@ -42,7 +42,9 @@ public class CosmosClientBuilder { private CosmosKeyCredential cosmosKeyCredential; private boolean sessionCapturingOverrideEnabled; private boolean connectionReuseAcrossClientsEnabled; + private TracerProvider tracerProvider; + private boolean contentResponseOnWriteEnabled; /** * Instantiates a new Cosmos client builder. @@ -51,6 +53,14 @@ public CosmosClientBuilder() { this.tracerProvider = new TracerProvider(ServiceLoader.load(Tracer.class)); } + /** + * Get the tracer provider + * @return tracerProvider + */ + public TracerProvider getTracerProvider() { + return tracerProvider; + } + /** * Session capturing is enabled by default for {@link ConsistencyLevel#SESSION}. * For other consistency levels, it is not needed, unless if you need occasionally send requests with Session @@ -297,12 +307,38 @@ public CosmosClientBuilder keyCredential(CosmosKeyCredential cosmosKeyCredential } /** - * Gets the tracer provider + * Gets the boolean which indicates whether to only return the headers and status code in Cosmos DB response + * in case of Create, Update and Delete operations on CosmosItem. * - * @return tracerProvider + * If set to false (which is by default), this removes the resource from response. It reduces networking + * and CPU load by not sending the resource back over the network and serializing it + * on the client. + * + * By-default, this is false. + * + * @return a boolean indicating whether resource will be included in the response or not */ - public TracerProvider getTracerProvider() { - return this.tracerProvider; + boolean isContentResponseOnWriteEnabled() { + return contentResponseOnWriteEnabled; + } + + /** + * Sets the boolean to only return the headers and status code in Cosmos DB response + * in case of Create, Update and Delete operations on CosmosItem. + * + * If set to false (which is by default), this removes the resource from response. It reduces networking + * and CPU load by not sending the resource back over the network and serializing it on the client. + * + * This feature does not impact RU usage for read or write operations. + * + * By-default, this is false. + * + * @param contentResponseOnWriteEnabled a boolean indicating whether resource will be included in the response or not + * @return current cosmosClientBuilder + */ + public CosmosClientBuilder contentResponseOnWriteEnabled(boolean contentResponseOnWriteEnabled) { + this.contentResponseOnWriteEnabled = contentResponseOnWriteEnabled; + return this; } /** From 8fbad2c39de56e6a74e0ea2014c274bd0d1abe2e Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Mon, 11 May 2020 10:36:28 -0400 Subject: [PATCH 13/40] check non opentelementry jar non existence in starting of all api --- .../com/azure/cosmos/CosmosAsyncClient.java | 63 +++-- .../azure/cosmos/CosmosAsyncContainer.java | 215 ++++++++++++------ .../com/azure/cosmos/CosmosAsyncDatabase.java | 148 ++++++++---- .../azure/cosmos/CosmosAsyncPermission.java | 56 +++-- .../com/azure/cosmos/CosmosAsyncScripts.java | 53 +++-- .../cosmos/CosmosAsyncStoredProcedure.java | 81 +++++-- .../com/azure/cosmos/CosmosAsyncTrigger.java | 51 +++-- .../com/azure/cosmos/CosmosAsyncUser.java | 90 +++++--- .../CosmosAsyncUserDefinedFunction.java | 52 +++-- .../com/azure/cosmos/CosmosClientBuilder.java | 11 +- .../cosmos/implementation/TracerProvider.java | 25 +- 11 files changed, 585 insertions(+), 260 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index a7dd95664e07..e03c64535981 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -4,6 +4,7 @@ import com.azure.core.annotation.ServiceClient; 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.CosmosAuthorizationTokenResolver; @@ -27,6 +28,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; @@ -67,8 +69,8 @@ public final class CosmosAsyncClient implements Closeable { this.cosmosKeyCredential = builder.getKeyCredential(); this.sessionCapturingOverride = builder.isSessionCapturingOverrideEnabled(); this.enableTransportClientSharing = builder.isConnectionReuseAcrossClientsEnabled(); - this.tracerProvider = builder.getTracerProvider(); this.contentResponseOnWriteEnabled = builder.isContentResponseOnWriteEnabled(); + this.tracerProvider = new TracerProvider(ServiceLoader.load(Tracer.class)); this.asyncDocumentClient = new AsyncDocumentClient.Builder() .withServiceEndpoint(this.serviceEndpoint) .withMasterKeyOrResourceToken(this.keyOrResourceToken) @@ -199,7 +201,12 @@ boolean isContentResponseOnWriteEnabled() { * an error. */ public Mono createDatabaseIfNotExists(CosmosDatabaseProperties databaseSettings) { - return createDatabaseIfNotExistsInternal(getDatabase(databaseSettings.getId())); + if(!getTracerProvider().isEnabled()) { + return createDatabaseIfNotExistsInternal(getDatabase(databaseSettings.getId())); + } + + return withContext(context -> createDatabaseIfNotExistsInternal(getDatabase(databaseSettings.getId()), + context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -212,7 +219,11 @@ public Mono createDatabaseIfNotExists(CosmosDatabas * an error */ public Mono createDatabaseIfNotExists(String id) { - return createDatabaseIfNotExistsInternal(getDatabase(id)); + if(!getTracerProvider().isEnabled()) { + return createDatabaseIfNotExistsInternal(getDatabase(id)); + } + + return withContext(context -> createDatabaseIfNotExistsInternal(getDatabase(id), context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -234,9 +245,12 @@ public Mono createDatabase(CosmosDatabaseProperties } Database wrappedDatabase = new Database(); wrappedDatabase.setId(databaseSettings.getId()); + if(!getTracerProvider().isEnabled()) { + return createDatabaseInternal(wrappedDatabase, options); + } + final CosmosDatabaseRequestOptions requestOptions = options; - return TracerProvider.cosmosWithContext(withContext(context -> createDatabase(wrappedDatabase, requestOptions - , context)), getTracerProvider()); + return withContext(context -> createDatabaseInternal(wrappedDatabase, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -291,9 +305,13 @@ public Mono createDatabase(CosmosDatabaseProperties ModelBridgeInternal.setOfferThroughput(options, throughput); Database wrappedDatabase = new Database(); wrappedDatabase.setId(databaseSettings.getId()); + if (!getTracerProvider().isEnabled()) { + return createDatabaseInternal(wrappedDatabase, options); + } + + final CosmosDatabaseRequestOptions requestOptions = options; - return TracerProvider.cosmosWithContext(withContext(context -> createDatabase(wrappedDatabase, requestOptions - , context)), getTracerProvider()); + return withContext(context -> createDatabaseInternal(wrappedDatabase, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -452,10 +470,6 @@ private CosmosPagedFlux queryDatabasesInternal(boolean }); } - private Mono createDatabaseIfNotExistsInternal(CosmosAsyncDatabase database) { - return TracerProvider.cosmosWithContext(withContext(context -> createDatabaseIfNotExistsInternal(database, - context)), getTracerProvider()); - } private Mono createDatabaseIfNotExistsInternal(CosmosAsyncDatabase database, Context context) { @@ -475,16 +489,35 @@ private Mono createDatabaseIfNotExistsInternal(Cosm this.serviceEndpoint); } - private Mono createDatabase(Database database, CosmosDatabaseRequestOptions options, + private Mono createDatabaseIfNotExistsInternal(CosmosAsyncDatabase database) { + return database.read().onErrorResume(exception -> { + final Throwable unwrappedException = Exceptions.unwrap(exception); + if (unwrappedException instanceof CosmosClientException) { + final CosmosClientException cosmosClientException = (CosmosClientException) unwrappedException; + if (cosmosClientException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) { + return createDatabase(new CosmosDatabaseProperties(database.getId()), + new CosmosDatabaseRequestOptions()); + } + } + return Mono.error(unwrappedException); + }); + } + + + private Mono createDatabaseInternal(Database database, CosmosDatabaseRequestOptions options, Context context) { String spanName = "createDatabase." + database.getId(); - Mono responseMono = asyncDocumentClient.createDatabase(database, + Mono responseMono = createDatabaseInternal(database, options); + return tracerProvider.traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), + this.serviceEndpoint); + } + + private Mono createDatabaseInternal(Database database, CosmosDatabaseRequestOptions options) { + return asyncDocumentClient.createDatabase(database, ModelBridgeInternal.toRequestOptions(options)) .map(databaseResourceResponse -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(databaseResourceResponse, this)) .single(); - return tracerProvider.traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), - this.serviceEndpoint); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index c02ae8fb0402..db819a4f97c2 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -90,7 +90,11 @@ public Mono read(CosmosContainerRequestOptions opt } final CosmosContainerRequestOptions requestOptions = options; - return TracerProvider.cosmosWithContext(withContext(context -> read(requestOptions, context)), database.getClient().getTracerProvider()); + if(!database.getClient().getTracerProvider().isEnabled()){ + return readInternal(options); + } + + return withContext(context -> read(requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -109,18 +113,12 @@ public Mono delete(CosmosContainerRequestOptions o options = new CosmosContainerRequestOptions(); } - final CosmosContainerRequestOptions requestOptions = options; - return TracerProvider.cosmosWithContext(withContext(context -> delete(requestOptions, context)), - database.getClient().getTracerProvider()); - } + if (!database.getClient().getTracerProvider().isEnabled()) { + return deleteInternal(options); + } - private Mono delete(CosmosContainerRequestOptions options, Context context) { - String spanName = "deleteContainer." + this.getId(); - Mono responseMono = database.getDocClientWrapper().deleteCollection(getLink(), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); - return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, spanName, database.getId(), database.getClient().getServiceEndpoint()); + final CosmosContainerRequestOptions requestOptions = options; + return withContext(context -> deleteInternal(requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -174,21 +172,12 @@ public Mono replace( options = new CosmosContainerRequestOptions(); } - final CosmosContainerRequestOptions requestOptions = options; - return TracerProvider.cosmosWithContext(withContext(context -> replace(containerProperties, requestOptions, - context)), database.getClient().getTracerProvider()); - } + if(!database.getClient().getTracerProvider().isEnabled()) { + return replaceInternal(containerProperties, options); + } - private Mono replace(CosmosContainerProperties containerProperties, - CosmosContainerRequestOptions options, - Context context) { - String spanName = "replaceContainer." + this.getId(); - Mono responseMono = database.getDocClientWrapper() - .replaceCollection(ModelBridgeInternal.getV2Collection(containerProperties), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); - return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, spanName, database.getId(), database.getClient().getServiceEndpoint()); + final CosmosContainerRequestOptions requestOptions = options; + return withContext(context -> replaceInternal(containerProperties, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /* CosmosAsyncItem operations */ @@ -246,25 +235,32 @@ public Mono> createItem(T item, CosmosItemRequest options = new CosmosItemRequestOptions(); } + if (!database.getClient().getTracerProvider().isEnabled()) { + return createItemInternal(item, options); + } + final CosmosItemRequestOptions requestOptions = options; - return TracerProvider.cosmosWithContext(withContext(context -> createItem(item, requestOptions, context)), - database.getClient().getTracerProvider()); + return withContext(context -> createItemInternal(item, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } - private Mono> createItem(T item, CosmosItemRequestOptions options, Context context) { + private Mono> createItemInternal(T item, CosmosItemRequestOptions options, Context context) { String spanName = "createItem." + this.getId(); + Mono> responseMono = createItemInternal(item, options); + return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, context + , spanName, database.getId(), database.getClient().getServiceEndpoint()); + } + + private Mono> createItemInternal(T item, CosmosItemRequestOptions options) { @SuppressWarnings("unchecked") Class itemType = (Class) item.getClass(); RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); - Mono> responseMono = database.getDocClientWrapper() + return database.getDocClientWrapper() .createDocument(getLink(), item, requestOptions, true) .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) .single(); - return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, context - , spanName, database.getId(), database.getClient().getServiceEndpoint()); } /** @@ -299,9 +295,12 @@ public Mono> upsertItem(T item, CosmosItemRequest options = new CosmosItemRequestOptions(); } + if(!database.getClient().getTracerProvider().isEnabled()) { + return upsertItemInternal(item, options); + } + final CosmosItemRequestOptions requestOptions = options; - return TracerProvider.cosmosWithContext(withContext(context -> upsertItem(item, requestOptions, context)), - database.getClient().getTracerProvider()); + return withContext(context -> upsertItemInternal(item, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -481,8 +480,12 @@ public Mono> readItem( ModelBridgeInternal.setPartitionKey(options, partitionKey); RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); - return TracerProvider.cosmosWithContext(withContext(context -> readItem(itemId, partitionKey, - requestOptions, itemType, context)), database.getClient().getTracerProvider()); + if(!database.getClient().getTracerProvider().isEnabled()) { + return readItemInternal(itemId, partitionKey, requestOptions, itemType); + } + + return withContext(context -> readItemInternal(itemId, partitionKey, + requestOptions, itemType, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -526,8 +529,12 @@ public Mono> replaceItem( ModelBridgeInternal.setPartitionKey(options, partitionKey); @SuppressWarnings("unchecked") Class itemType = (Class) item.getClass(); + if(!database.getClient().getTracerProvider().isEnabled()){ + return replaceItemInternal(itemType, itemId, doc, options); + } + final CosmosItemRequestOptions requestOptions = options; - return TracerProvider.cosmosWithContext(withContext(context -> replaceItem(itemType, itemId, doc, requestOptions,context)), database.getClient().getTracerProvider()); + return withContext(context -> replaceItemInternal(itemType, itemId, doc, requestOptions,context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -565,7 +572,11 @@ public Mono> deleteItem( } ModelBridgeInternal.setPartitionKey(options, partitionKey); RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); - return TracerProvider.cosmosWithContext(withContext(context -> deleteItem(itemId, requestOptions, context)), database.getClient().getTracerProvider()); + if(!database.getClient().getTracerProvider().isEnabled()) { + return deleteItemInternal(itemId, requestOptions); + } + + return withContext(context -> deleteItemInternal(itemId, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } private String getItemLink(String itemId) { @@ -658,7 +669,11 @@ public CosmosAsyncConflict getConflict(String id) { * @return a {@link Mono} containing throughput or an error. */ public Mono readProvisionedThroughput() { - return TracerProvider.cosmosWithContext(withContext(context -> readProvisionedThroughput(context)), database.getClient().getTracerProvider()); + if(!database.getClient().getTracerProvider().isEnabled()){ + return readProvisionedThroughputInternal(); + } + + return withContext(context -> readProvisionedThroughputInternal(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -670,7 +685,11 @@ public Mono readProvisionedThroughput() { * @return a {@link Mono} containing throughput or an error. */ public Mono replaceProvisionedThroughput(int requestUnitsPerSecond) { - return TracerProvider.cosmosWithContext(withContext(context -> replaceProvisionedThroughput(requestUnitsPerSecond, context)), database.getClient().getTracerProvider()); + if (!database.getClient().getTracerProvider().isEnabled()) { + return replaceProvisionedThroughputInternal(requestUnitsPerSecond); + } + + return withContext(context -> replaceProvisionedThroughput(requestUnitsPerSecond, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -758,75 +777,109 @@ String getLink() { return this.link; } - private Mono> deleteItem( + private Mono> deleteItemInternal( String itemId, RequestOptions requestOptions, Context context) { String spanName = "deleteItem." + this.getId(); - Mono> responseMono = this.getDatabase() + Mono> responseMono = deleteItemInternal(itemId, requestOptions); + return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, + context, spanName, database.getId(), database.getClient().getServiceEndpoint()); + } + + private Mono> deleteItemInternal( + String itemId, + RequestOptions requestOptions) { + return this.getDatabase() .getDocClientWrapper() .deleteDocument(getItemLink(itemId), requestOptions) .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponseWithObjectType(response)) .single(); - return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, - context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } - private Mono> replaceItem( + private Mono> replaceItemInternal( Class itemType, String itemId, Document doc, CosmosItemRequestOptions options, Context context) { String spanName = "replaceItem." + this.getId(); - Mono> responseMono = this.getDatabase() + Mono> responseMono = replaceItemInternal(itemType, itemId, doc, options); + return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, + context, spanName, database.getId(), database.getClient().getServiceEndpoint()); + } + + private Mono> replaceItemInternal( + Class itemType, + String itemId, + Document doc, + CosmosItemRequestOptions options) { + return this.getDatabase() .getDocClientWrapper() .replaceDocument(getItemLink(itemId), doc, ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) .single(); + } + + private Mono> upsertItemInternal(T item, CosmosItemRequestOptions options, Context context) { + String spanName = "upsertItem." + this.getId(); + Mono> responseMono = upsertItemInternal(item, options); return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } - private Mono> upsertItem(T item, CosmosItemRequestOptions options, Context context) { - String spanName = "upsertItem." + this.getId(); + private Mono> upsertItemInternal(T item, CosmosItemRequestOptions options) { @SuppressWarnings("unchecked") Class itemType = (Class) item.getClass(); - Mono> responseMono = this.getDatabase().getDocClientWrapper() + return this.getDatabase().getDocClientWrapper() .upsertDocument(this.getLink(), item, ModelBridgeInternal.toRequestOptions(options), true) .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) .single(); - return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, - context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } - private Mono> readItem( + private Mono> readItemInternal( String itemId, PartitionKey partitionKey, RequestOptions requestOptions, Class itemType, Context context) { String spanName = "readItem." + this.getId(); - Mono> responseMono = this.getDatabase().getDocClientWrapper() + Mono> responseMono = readItemInternal(itemId, partitionKey, requestOptions, itemType); + return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, + context, spanName, database.getId(), database.getClient().getServiceEndpoint()); + } + + private Mono> readItemInternal( + String itemId, PartitionKey partitionKey, + RequestOptions requestOptions, Class itemType) { + return this.getDatabase().getDocClientWrapper() .readDocument(getItemLink(itemId), requestOptions) .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) .single(); - return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, - context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } private Mono read(CosmosContainerRequestOptions options, Context context) { String spanName = "readContainer." + this.getId(); - Mono responseMono = database.getDocClientWrapper().readCollection(getLink(), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); + Mono responseMono = readInternal(options); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } - private Mono readProvisionedThroughput(Context context) { + private Mono readInternal(CosmosContainerRequestOptions options) { + return database.getDocClientWrapper().readCollection(getLink(), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); + } + + private Mono readProvisionedThroughputInternal(Context context) { String spanName = "readProvisionedThroughput." + this.getId(); - Mono responseMono = this.read() + Mono responseMono = readProvisionedThroughputInternal(); + return this.getDatabase().getClient().getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono + , context, spanName, database.getId(), database.getClient().getServiceEndpoint()); + } + + private Mono readProvisionedThroughputInternal() { + return this.read() .flatMap(cosmosContainerResponse -> database.getDocClientWrapper() .queryOffers("select * from c where c.offerResourceId = '" @@ -843,13 +896,17 @@ private Mono readProvisionedThroughput(Context context) { .readOffer(offerFeedResponse.getResults().get(0).getSelfLink()) .single(); }).map(cosmosOfferResponse -> cosmosOfferResponse.getResource().getThroughput()); - return this.getDatabase().getClient().getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono - , context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Context context) { String spanName = "replaceProvisionedThroughput." + this.getId(); - Mono responseMono = this.read() + Mono responseMono = replaceProvisionedThroughputInternal(requestUnitsPerSecond); + return this.getDatabase().getClient().getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono + , context, spanName, database.getId(), database.getClient().getServiceEndpoint()); + } + + private Mono replaceProvisionedThroughputInternal(int requestUnitsPerSecond) { + return this.read() .flatMap(cosmosContainerResponse -> database.getDocClientWrapper() .queryOffers("select * from c where c.offerResourceId = '" @@ -866,7 +923,35 @@ private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Co offer.setThroughput(requestUnitsPerSecond); return database.getDocClientWrapper().replaceOffer(offer).single(); }).map(offerResourceResponse -> offerResourceResponse.getResource().getThroughput()); - return this.getDatabase().getClient().getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono - , context, spanName, database.getId(), database.getClient().getServiceEndpoint()); + } + + private Mono deleteInternal(CosmosContainerRequestOptions options, Context context) { + String spanName = "deleteContainer." + this.getId(); + Mono responseMono = deleteInternal(options); + return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, + context, spanName, database.getId(), database.getClient().getServiceEndpoint()); + } + + private Mono deleteInternal(CosmosContainerRequestOptions options) { + return database.getDocClientWrapper().deleteCollection(getLink(), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); + } + + private Mono replaceInternal(CosmosContainerProperties containerProperties, + CosmosContainerRequestOptions options, + Context context) { + String spanName = "replaceContainer." + this.getId(); + Mono responseMono = replaceInternal(containerProperties, options); + return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, + context, spanName, database.getId(), database.getClient().getServiceEndpoint()); + } + + private Mono replaceInternal(CosmosContainerProperties containerProperties, + CosmosContainerRequestOptions options) { + return database.getDocClientWrapper() + .replaceCollection(ModelBridgeInternal.getV2Collection(containerProperties), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, database)).single(); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java index 3dd11d243857..d8ee89cb4bf1 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java @@ -86,8 +86,12 @@ public Mono read(CosmosDatabaseRequestOptions optio if (options == null) { options = new CosmosDatabaseRequestOptions(); } + + if (!client.getTracerProvider().isEnabled()) { + return readInternal(options); + } final CosmosDatabaseRequestOptions requestOptions = options; - return TracerProvider.cosmosWithContext(withContext(context -> read(requestOptions, context)), client.getTracerProvider()); + return withContext(context -> readInternal(requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -117,10 +121,12 @@ public Mono delete(CosmosDatabaseRequestOptions opt if (options == null) { options = new CosmosDatabaseRequestOptions(); } + if (!client.getTracerProvider().isEnabled()) { + return deleteInternal(options); + } final CosmosDatabaseRequestOptions requestOptions = options; - return TracerProvider.cosmosWithContext(withContext(context -> delete(requestOptions, context)), - client.getTracerProvider()); + return withContext(context -> deleteInternal(requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /* CosmosAsyncContainer operations */ @@ -206,9 +212,12 @@ public Mono createContainer( options = new CosmosContainerRequestOptions(); } + if (!client.getTracerProvider().isEnabled()) { + return createContainerInternal(containerProperties, options); + } + final CosmosContainerRequestOptions requestOptions = options; - return TracerProvider.cosmosWithContext(withContext(context -> createContainer(containerProperties, - requestOptions, context)), client.getTracerProvider()); + return withContext(context -> createContainerInternal(containerProperties, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } @@ -351,14 +360,13 @@ public Mono createContainerIfNotExists( CosmosContainerRequestOptions options = new CosmosContainerRequestOptions(); ModelBridgeInternal.setOfferThroughput(options, throughput); CosmosAsyncContainer container = getContainer(id); - return createContainerIfNotExistsInternal(new CosmosContainerProperties(id, partitionKeyPath), container, - options); - } + if (!client.getTracerProvider().isEnabled()) { + return createContainerIfNotExistsInternal(new CosmosContainerProperties(id, partitionKeyPath), container, + options); + } - private Mono createContainerIfNotExistsInternal( - CosmosContainerProperties containerProperties, CosmosAsyncContainer container, - CosmosContainerRequestOptions options) { - return TracerProvider.cosmosWithContext(withContext(context -> createContainerIfNotExistsInternal(containerProperties, container, options, context)), client.getTracerProvider()); + return withContext(context -> createContainerIfNotExistsInternal(new CosmosContainerProperties(id, + partitionKeyPath), container, options, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -485,8 +493,10 @@ public CosmosAsyncContainer getContainer(String id) { * created cosmos user or an error. */ public Mono createUser(CosmosUserProperties userProperties) { - return TracerProvider.cosmosWithContext(withContext(context -> createUserInternal(userProperties, context)), - client.getTracerProvider()); + if (!client.getTracerProvider().isEnabled()) { + return createUserInternal(userProperties); + } + return withContext(context -> createUserInternal(userProperties, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } @@ -502,8 +512,11 @@ public Mono createUser(CosmosUserProperties userPropert * upserted user or an error. */ public Mono upsertUser(CosmosUserProperties userProperties) { - return TracerProvider.cosmosWithContext(withContext(context -> upsertUserInternal(userProperties, context)), - client.getTracerProvider()); + if (!client.getTracerProvider().isEnabled()) { + return upsertUserInternal(userProperties); + } + + return withContext(context -> upsertUserInternal(userProperties, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -622,8 +635,11 @@ public CosmosAsyncUser getUser(String id) { * @return a {@link Mono} containing throughput or an error. */ public Mono readProvisionedThroughput() { - return TracerProvider.cosmosWithContext(withContext(context -> readProvisionedThroughput(context)), - client.getTracerProvider()); + if(!client.getTracerProvider().isEnabled()) { + return readProvisionedThroughputInternal(); + } + + return withContext(context -> readProvisionedThroughputInternal(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -635,7 +651,11 @@ public Mono readProvisionedThroughput() { * @return a {@link Mono} containing throughput or an error. */ public Mono replaceProvisionedThroughput(int requestUnitsPerSecond) { - return TracerProvider.cosmosWithContext(withContext(context -> replaceProvisionedThroughput(requestUnitsPerSecond, context)), client.getTracerProvider()); + if (!client.getTracerProvider().isEnabled()) { + return replaceProvisionedThroughputInternal(requestUnitsPerSecond); + } + + return withContext(context -> replaceProvisionedThroughputInternal(requestUnitsPerSecond, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -790,40 +810,76 @@ private Mono createContainerIfNotExistsInternal( spanName, getId(), getClient().getServiceEndpoint()); } - private Mono createContainer( + private Mono createContainerIfNotExistsInternal( + CosmosContainerProperties containerProperties, + CosmosAsyncContainer container, + CosmosContainerRequestOptions options) { + return container.read(options).onErrorResume(exception -> { + final Throwable unwrappedException = Exceptions.unwrap(exception); + if (unwrappedException instanceof CosmosClientException) { + final CosmosClientException cosmosClientException = (CosmosClientException) unwrappedException; + if (cosmosClientException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) { + return createContainer(containerProperties, options); + } + } + return Mono.error(unwrappedException); + }); + } + + private Mono createContainerInternal( CosmosContainerProperties containerProperties, CosmosContainerRequestOptions options, Context context) { String spanName = "createContainer." + containerProperties.getId(); - Mono responseMono = getDocClientWrapper() + Mono responseMono = createContainerInternal(containerProperties, options); + return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, getId(), getClient().getServiceEndpoint()); + } + + private Mono createContainerInternal( + CosmosContainerProperties containerProperties, + CosmosContainerRequestOptions options) { + return getDocClientWrapper() .createCollection(this.getLink(), ModelBridgeInternal.getV2Collection(containerProperties), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncContainerResponse(response, this)).single(); + } + + private Mono readInternal(CosmosDatabaseRequestOptions options, Context context) { + String spanName = "readDatabase." + this.getId(); + Mono responseMono = readInternal(options); return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, getId(), getClient().getServiceEndpoint()); } - private Mono read(CosmosDatabaseRequestOptions options, Context context) { - String spanName = "readDatabase." + this.getId(); - Mono responseMono = getDocClientWrapper().readDatabase(getLink(), + private Mono readInternal(CosmosDatabaseRequestOptions options) { + return getDocClientWrapper().readDatabase(getLink(), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single(); + } + + private Mono deleteInternal(CosmosDatabaseRequestOptions options, Context context) { + String spanName = "deleteDatabase." + this.getId(); + Mono responseMono = deleteInternal(options); return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, getId(), getClient().getServiceEndpoint()); } - private Mono delete(CosmosDatabaseRequestOptions options, Context context) { - String spanName = "deleteDatabase." + this.getId(); - Mono responseMono = getDocClientWrapper().deleteDatabase(getLink(), + private Mono deleteInternal(CosmosDatabaseRequestOptions options) { + return getDocClientWrapper().deleteDatabase(getLink(), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncDatabaseResponse(response, getClient())).single(); - return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, getId(), getClient().getServiceEndpoint()); } - private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Context context) { + private Mono replaceProvisionedThroughputInternal(int requestUnitsPerSecond, Context context) { String spanName = "replaceOffer." + this.getId(); - Mono responseMono = this.read() + Mono responseMono = replaceProvisionedThroughputInternal(requestUnitsPerSecond); + return this.client.getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono, context, spanName + , getId(), getClient().getServiceEndpoint()); + } + + private Mono replaceProvisionedThroughputInternal(int requestUnitsPerSecond) { + return this.read() .flatMap(cosmosDatabaseResponse -> this.getDocClientWrapper() .queryOffers("select * from c where c.offerResourceId = '" + cosmosDatabaseResponse.getProperties() @@ -844,13 +900,17 @@ private Mono replaceProvisionedThroughput(int requestUnitsPerSecond, Co }).map(offerResourceResponse -> offerResourceResponse .getResource() .getThroughput())); - return this.client.getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono, context, spanName - , getId(), getClient().getServiceEndpoint()); } - private Mono readProvisionedThroughput(Context context) { + private Mono readProvisionedThroughputInternal(Context context) { String spanName = "readProvisionedThroughput." + this.getId(); - Mono responseMono = this.read() + Mono responseMono = readProvisionedThroughputInternal(); + return this.client.getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono + , context, spanName, getId(), getClient().getServiceEndpoint()); + } + + private Mono readProvisionedThroughputInternal() { + return this.read() .flatMap(cosmosDatabaseResponse -> getDocClientWrapper() .queryOffers("select * from c where c.offerResourceId = '" + cosmosDatabaseResponse @@ -873,23 +933,29 @@ private Mono readProvisionedThroughput(Context context) { }).map(cosmosContainerResponse1 -> cosmosContainerResponse1 .getResource() .getThroughput())); - return this.client.getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono - , context, spanName, getId(), getClient().getServiceEndpoint()); } private Mono createUserInternal(CosmosUserProperties userProperties, Context context) { String spanName = "createUser." + this.getId(); - Mono responseMono = getDocClientWrapper().createUser(this.getLink(), ModelBridgeInternal.getV2User(userProperties), null) - .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, this)).single(); + Mono responseMono = createUserInternal(userProperties); return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, getId(), getClient().getServiceEndpoint()); } + private Mono createUserInternal(CosmosUserProperties userProperties) { + return getDocClientWrapper().createUser(this.getLink(), ModelBridgeInternal.getV2User(userProperties), null) + .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, this)).single(); + } + private Mono upsertUserInternal(CosmosUserProperties userProperties, Context context) { String spanName = "upsertUser." + this.getId(); - Mono responseMono = getDocClientWrapper().upsertUser(this.getLink(), ModelBridgeInternal.getV2User(userProperties), null) - .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, this)).single(); + Mono responseMono = upsertUserInternal(userProperties); return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, getId(), getClient().getServiceEndpoint()); } + + private Mono upsertUserInternal(CosmosUserProperties userProperties) { + return getDocClientWrapper().upsertUser(this.getLink(), ModelBridgeInternal.getV2User(userProperties), null) + .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, this)).single(); + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java index 5075c2ffc293..694303473c9d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java @@ -61,9 +61,12 @@ public Mono read(CosmosPermissionRequestOptions o options = new CosmosPermissionRequestOptions(); } + if (!cosmosUser.getDatabase().getClient().getTracerProvider().isEnabled()) { + return readInternal(options); + } + final CosmosPermissionRequestOptions requestOptions = options; - return TracerProvider.cosmosWithContext(withContext(context -> readInternal(requestOptions, context)), - cosmosUser.getDatabase().getClient().getTracerProvider()); + return withContext(context -> readInternal(requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -83,10 +86,12 @@ public Mono replace(CosmosPermissionProperties pe options = new CosmosPermissionRequestOptions(); } + if (!cosmosUser.getDatabase().getClient().getTracerProvider().isEnabled()) { + return replaceInternal(permissionSettings, options); + } + final CosmosPermissionRequestOptions requestOptions = options; - return TracerProvider.cosmosWithContext(withContext(context -> replaceInternal(permissionSettings, - requestOptions, context)), - cosmosUser.getDatabase().getClient().getTracerProvider()); + return withContext(context -> replaceInternal(permissionSettings, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -104,9 +109,12 @@ public Mono delete(CosmosPermissionRequestOptions options = new CosmosPermissionRequestOptions(); } + if (!cosmosUser.getDatabase().getClient().getTracerProvider().isEnabled()) { + return deleteInternal(options); + } + final CosmosPermissionRequestOptions requestOptions = options; - return TracerProvider.cosmosWithContext(withContext(context -> deleteInternal(requestOptions, context)), - cosmosUser.getDatabase().getClient().getTracerProvider()); + return withContext(context -> deleteInternal(requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } String getURIPathSegment() { @@ -130,13 +138,18 @@ String getLink() { private Mono readInternal(CosmosPermissionRequestOptions options, Context context) { String spanName = "readPermission." + cosmosUser.getId(); - Mono responseMono = cosmosUser.getDatabase() + Mono responseMono = readInternal(options); + return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, cosmosUser.getDatabase().getId(), cosmosUser.getDatabase().getClient().getServiceEndpoint()); + } + + private Mono readInternal(CosmosPermissionRequestOptions options) { + + return cosmosUser.getDatabase() .getDocClientWrapper() .readPermission(getLink(), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncPermissionResponse(response, cosmosUser)) .single(); - return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, cosmosUser.getDatabase().getId(), cosmosUser.getDatabase().getClient().getServiceEndpoint()); } private Mono replaceInternal(CosmosPermissionProperties permissionSettings, @@ -144,26 +157,37 @@ private Mono replaceInternal(CosmosPermissionProp Context context) { String spanName = "replacePermission." + cosmosUser.getId(); - Mono responseMono = cosmosUser.getDatabase() + Mono responseMono = replaceInternal(permissionSettings, options); + return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, cosmosUser.getDatabase().getId(), cosmosUser.getDatabase().getClient().getServiceEndpoint()); + } + + private Mono replaceInternal(CosmosPermissionProperties permissionSettings, + CosmosPermissionRequestOptions options) { + return cosmosUser.getDatabase() .getDocClientWrapper() .replacePermission(ModelBridgeInternal.getV2Permissions(permissionSettings), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncPermissionResponse(response, cosmosUser)) .single(); - return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, cosmosUser.getDatabase().getId(), cosmosUser.getDatabase().getClient().getServiceEndpoint()); } private Mono deleteInternal(CosmosPermissionRequestOptions options, Context context) { String spanName = "deletePermission." + cosmosUser.getId(); - Mono responseMono = cosmosUser.getDatabase() + Mono responseMono = deleteInternal(options); + return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, cosmosUser.getDatabase().getId(), cosmosUser.getDatabase().getClient().getServiceEndpoint()); + } + + private Mono deleteInternal(CosmosPermissionRequestOptions options) { + + return cosmosUser.getDatabase() .getDocClientWrapper() .deletePermission(getLink(), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncPermissionResponse(response, cosmosUser)) .single(); - return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, cosmosUser.getDatabase().getId(), cosmosUser.getDatabase().getClient().getServiceEndpoint()); } + } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java index 06f47bc15c76..ae5249e03bfd 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java @@ -74,9 +74,12 @@ public Mono createStoredProcedure( StoredProcedure sProc = new StoredProcedure(); sProc.setId(properties.getId()); sProc.setBody(properties.getBody()); + if (!container.getDatabase().getClient().getTracerProvider().isEnabled()) { + return createStoredProcedureInternal(sProc, options); + } + final CosmosStoredProcedureRequestOptions requestOptions = options; - return TracerProvider.cosmosWithContext(withContext(context -> createStoredProcedure(sProc, requestOptions, context)), - container.getDatabase().getClient().getTracerProvider()); + return withContext(context -> createStoredProcedureInternal(sProc, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } @@ -175,9 +178,11 @@ public Mono createUserDefinedFunction( UserDefinedFunction udf = new UserDefinedFunction(); udf.setId(properties.getId()); udf.setBody(properties.getBody()); + if (!container.getDatabase().getClient().getTracerProvider().isEnabled()) { + return createUserDefinedFunctionInternal(udf); + } - return TracerProvider.cosmosWithContext(withContext(context -> createUserDefinedFunction(udf, context)), - container.getDatabase().getClient().getTracerProvider()); + return withContext(context -> createUserDefinedFunctionInternal(udf, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -271,8 +276,11 @@ public CosmosAsyncUserDefinedFunction getUserDefinedFunction(String id) { * @return an {@link Mono} containing the single resource response with the created trigger or an error. */ public Mono createTrigger(CosmosTriggerProperties properties) { - return TracerProvider.cosmosWithContext(withContext(context -> createTrigger(properties, context)), - container.getDatabase().getClient().getTracerProvider()); + if (!container.getDatabase().getClient().getTracerProvider().isEnabled()) { + return createTriggerInternal(properties); + } + + return withContext(context -> createTriggerInternal(properties, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -423,33 +431,48 @@ private CosmosPagedFlux queryTriggersInternal( }); } - private Mono createStoredProcedure(StoredProcedure sProc, + private Mono createStoredProcedureInternal(StoredProcedure sProc, CosmosStoredProcedureRequestOptions options, Context context) { String spanName = "createStoredProcedure." + container.getId(); - Mono responseMono = database.getDocClientWrapper() + Mono responseMono = createStoredProcedureInternal(sProc, options); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); + } + + private Mono createStoredProcedureInternal(StoredProcedure sProc, + CosmosStoredProcedureRequestOptions options) { + return database.getDocClientWrapper() .createStoredProcedure(container.getLink(), sProc, ModelBridgeInternal.toRequestOptions(options)).map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, this.container)) .single(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } - private Mono createUserDefinedFunction( + private Mono createUserDefinedFunctionInternal( UserDefinedFunction udf, Context context) { String spanName = "createUserDefinedFunction." + container.getId(); - Mono responseMono = database.getDocClientWrapper() + Mono responseMono = createUserDefinedFunctionInternal(udf); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); + } + + private Mono createUserDefinedFunctionInternal( + UserDefinedFunction udf) { + return database.getDocClientWrapper() .createUserDefinedFunction(container.getLink(), udf, null).map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, this.container)).single(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } - private Mono createTrigger(CosmosTriggerProperties properties, Context context) { + private Mono createTriggerInternal(CosmosTriggerProperties properties, Context context) { String spanName = "createTrigger." + container.getId(); + Mono responseMono = createTriggerInternal(properties); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); + } + + private Mono createTriggerInternal(CosmosTriggerProperties properties) { Trigger trigger = new Trigger(ModelBridgeInternal.toJsonFromJsonSerializable(properties)); - Mono responseMono = database.getDocClientWrapper() + return database.getDocClientWrapper() .createTrigger(container.getLink(), trigger, null) .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, this.container)) .single(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } + } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java index 39e064790fa2..8acc71aefd15 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java @@ -74,8 +74,11 @@ public Mono read() { * @return an {@link Mono} containing the single resource response with the read stored procedure or an error. */ public Mono read(CosmosStoredProcedureRequestOptions options) { - return TracerProvider.cosmosWithContext(withContext(context -> read(options, context)), - cosmosContainer.getDatabase().getClient().getTracerProvider()); + if (!cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled()) { + return readInternal(options); + } + + return withContext(context -> readInternal(options, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -104,8 +107,11 @@ public Mono delete() { * @return an {@link Mono} containing the single resource response for the deleted stored procedure or an error. */ public Mono delete(CosmosStoredProcedureRequestOptions options) { - return TracerProvider.cosmosWithContext(withContext(context -> delete(options, context)), - cosmosContainer.getDatabase().getClient().getTracerProvider()); + if (!cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled()) { + return deleteInternal(options); + } + + return withContext(context -> deleteInternal(options, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -122,8 +128,11 @@ public Mono delete(CosmosStoredProcedureRequ */ public Mono execute(Object[] procedureParams, CosmosStoredProcedureRequestOptions options) { - return TracerProvider.cosmosWithContext(withContext(context -> execute(procedureParams, options, context)), - cosmosContainer.getDatabase().getClient().getTracerProvider()); + if (!cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled()) { + return executeInternal(procedureParams, options); + } + + return withContext(context -> executeInternal(procedureParams, options, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -155,8 +164,12 @@ public Mono replace(CosmosStoredProcedurePro */ public Mono replace(CosmosStoredProcedureProperties storedProcedureSettings, CosmosStoredProcedureRequestOptions options) { - return TracerProvider.cosmosWithContext(withContext(context -> replace(storedProcedureSettings, options, - context)), cosmosContainer.getDatabase().getClient().getTracerProvider()); + if (!cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled()) { + return replaceInternal(storedProcedureSettings, options); + } + + return withContext(context -> replaceInternal(storedProcedureSettings, options, + context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } String getURIPathSegment() { @@ -177,65 +190,83 @@ String getLink() { return builder.toString(); } - private Mono read(CosmosStoredProcedureRequestOptions options, + private Mono readInternal(CosmosStoredProcedureRequestOptions options, Context context) { String spanName = "readStoredProcedure." + cosmosContainer.getId(); + Mono responseMono = readInternal(options); + return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); + } + + private Mono readInternal(CosmosStoredProcedureRequestOptions options) { if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } - Mono responseMono = - cosmosContainer.getDatabase().getDocClientWrapper().readStoredProcedure(getLink(), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, - cosmosContainer)).single(); - return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); + return cosmosContainer.getDatabase().getDocClientWrapper().readStoredProcedure(getLink(), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, + cosmosContainer)).single(); } - private Mono delete(CosmosStoredProcedureRequestOptions options, + private Mono deleteInternal(CosmosStoredProcedureRequestOptions options, Context context) { String spanName = "deleteStoredProcedure." + cosmosContainer.getId(); + Mono responseMono = deleteInternal(options); + return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); + } + + private Mono deleteInternal(CosmosStoredProcedureRequestOptions options) { if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } - Mono responseMono = cosmosContainer.getDatabase() + return cosmosContainer.getDatabase() .getDocClientWrapper() .deleteStoredProcedure(getLink(), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer)) .single(); - return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); } - private Mono execute(Object[] procedureParams, + private Mono executeInternal(Object[] procedureParams, CosmosStoredProcedureRequestOptions options, Context context) { String spanName = "executeStoredProcedure." + cosmosContainer.getId(); + Mono responseMono = executeInternal(procedureParams, options); + return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); + } + + private Mono executeInternal(Object[] procedureParams, + CosmosStoredProcedureRequestOptions options) { if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } - Mono responseMono = cosmosContainer.getDatabase() + + return cosmosContainer.getDatabase() .getDocClientWrapper() .executeStoredProcedure(getLink(), ModelBridgeInternal.toRequestOptions(options), procedureParams) .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer, this.id)) .single(); - return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); } - private Mono replace(CosmosStoredProcedureProperties storedProcedureSettings, + private Mono replaceInternal(CosmosStoredProcedureProperties storedProcedureSettings, CosmosStoredProcedureRequestOptions options, Context context) { String spanName = "replaceStoredProcedure." + cosmosContainer.getId(); + Mono responseMono = replaceInternal(storedProcedureSettings, options); + return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); + } + + private Mono replaceInternal(CosmosStoredProcedureProperties storedProcedureSettings, + CosmosStoredProcedureRequestOptions options) { if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } - Mono responseMono = cosmosContainer.getDatabase() + + return cosmosContainer.getDatabase() .getDocClientWrapper() .replaceStoredProcedure(new StoredProcedure(ModelBridgeInternal.toJsonFromJsonSerializable(storedProcedureSettings)), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncStoredProcedureResponse(response, cosmosContainer)) .single(); - return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); } - } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java index e0efc005db0d..47ebc197cb5c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java @@ -56,8 +56,11 @@ CosmosAsyncTrigger setId(String id) { * @return an {@link Mono} containing the single resource response for the read cosmos trigger or an error. */ public Mono read() { - return TracerProvider.cosmosWithContext(withContext(context -> read(context)), - container.getDatabase().getClient().getTracerProvider()); + if (!container.getDatabase().getClient().getTracerProvider().isEnabled()) { + return readInternal(); + } + + return withContext(context -> readInternal(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -71,8 +74,11 @@ public Mono read() { * @return an {@link Mono} containing the single resource response with the replaced cosmos trigger or an error. */ public Mono replace(CosmosTriggerProperties triggerSettings) { - return TracerProvider.cosmosWithContext(withContext(context -> replace(triggerSettings, context)), - container.getDatabase().getClient().getTracerProvider()); + if (!container.getDatabase().getClient().getTracerProvider().isEnabled()) { + return replaceInternal(triggerSettings); + } + + return withContext(context -> replaceInternal(triggerSettings, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -85,8 +91,11 @@ public Mono replace(CosmosTriggerProperties triggerS * @return an {@link Mono} containing the single resource response for the deleted cosmos trigger or an error. */ public Mono delete() { - return TracerProvider.cosmosWithContext(withContext(context -> delete(context)), - container.getDatabase().getClient().getTracerProvider()); + if (!container.getDatabase().getClient().getTracerProvider().isEnabled()) { + return deleteInternal(); + } + + return withContext(context -> deleteInternal(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } String getURIPathSegment() { @@ -107,33 +116,45 @@ String getLink() { return builder.toString(); } - private Mono read(Context context) { + private Mono readInternal(Context context) { String spanName = "readTrigger." + container.getId(); - Mono responseMono = container.getDatabase() + Mono responseMono = readInternal(); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); + } + + private Mono readInternal() { + return container.getDatabase() .getDocClientWrapper() .readTrigger(getLink(), null) .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, container)) .single(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); } - private Mono replace(CosmosTriggerProperties triggerSettings, Context context) { + private Mono replaceInternal(CosmosTriggerProperties triggerSettings, Context context) { String spanName = "replaceTrigger." + container.getId(); - Mono responseMono = container.getDatabase() + Mono responseMono = replaceInternal(triggerSettings); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); + } + + private Mono replaceInternal(CosmosTriggerProperties triggerSettings) { + return container.getDatabase() .getDocClientWrapper() .replaceTrigger(new Trigger(ModelBridgeInternal.toJsonFromJsonSerializable(triggerSettings)), null) .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, container)) .single(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); } - private Mono delete(Context context) { + private Mono deleteInternal(Context context) { String spanName = "deleteTrigger." + container.getId(); - Mono responseMono = container.getDatabase() + Mono responseMono = deleteInternal(); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); + } + + private Mono deleteInternal() { + return container.getDatabase() .getDocClientWrapper() .deleteTrigger(getLink(), null) .map(response -> ModelBridgeInternal.createCosmosAsyncTriggerResponse(response, container)) .single(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java index f51e3e61af9c..c9125995f112 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java @@ -59,8 +59,11 @@ CosmosAsyncUser setId(String id) { * @return a {@link Mono} containing the single resource response with the read user or an error. */ public Mono read() { - return TracerProvider.cosmosWithContext(withContext(context -> readInternal(context)), - database.getClient().getTracerProvider()); + if (!database.getClient().getTracerProvider().isEnabled()) { + return readInternal(); + } + + return withContext(context -> readInternal(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -70,8 +73,11 @@ public Mono read() { * @return a {@link Mono} containing the single resource response with the replaced user or an error. */ public Mono replace(CosmosUserProperties userSettings) { - return TracerProvider.cosmosWithContext(withContext(context -> replaceInternal(userSettings, context)), - database.getClient().getTracerProvider()); + if (!database.getClient().getTracerProvider().isEnabled()) { + return replaceInternal(userSettings); + } + + return withContext(context -> replaceInternal(userSettings, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -80,8 +86,11 @@ public Mono replace(CosmosUserProperties userSettings) * @return a {@link Mono} containing the single resource response with the deleted user or an error. */ public Mono delete() { - return TracerProvider.cosmosWithContext(withContext(context -> deleteInternal(context)), - database.getClient().getTracerProvider()); + if (!database.getClient().getTracerProvider().isEnabled()) { + return deleteInternal(); + } + + return withContext(context -> deleteInternal(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -103,9 +112,12 @@ public Mono createPermission( } Permission permission = ModelBridgeInternal.getV2Permissions(permissionSettings); + if (!database.getClient().getTracerProvider().isEnabled()) { + return createPermissionInternal(permission, options); + } + final CosmosPermissionRequestOptions requesOptions = options; - return TracerProvider.cosmosWithContext(withContext(context -> createPermissionInternal(permission, requesOptions, context)), - database.getClient().getTracerProvider()); + return withContext(context -> createPermissionInternal(permission, requesOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -127,10 +139,12 @@ public Mono upsertPermission( options = new CosmosPermissionRequestOptions(); } + if (!database.getClient().getTracerProvider().isEnabled()) { + return upsertPermissionInternal(permission, options); + } + final CosmosPermissionRequestOptions requestOptions = options; - return TracerProvider.cosmosWithContext(withContext(context -> upsertPermissionInternal(permission, - requestOptions, context)), - database.getClient().getTracerProvider()); + return withContext(context -> upsertPermissionInternal(permission, requestOptions, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } @@ -243,42 +257,60 @@ public CosmosAsyncDatabase getDatabase() { private Mono readInternal(Context context) { String spanName = "readUser." + getId(); - Mono responseMono = this.database.getDocClientWrapper() - .readUser(getLink(), null) - .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, database)).single(); + Mono responseMono = readInternal(); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } + private Mono readInternal() { + return this.database.getDocClientWrapper() + .readUser(getLink(), null) + .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, database)).single(); + } + private Mono replaceInternal(CosmosUserProperties userSettings, Context context) { String spanName = "replaceUser." + getId(); - Mono responseMono = this.database.getDocClientWrapper() - .replaceUser(ModelBridgeInternal.getV2User(userSettings), null) - .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, database)).single(); + Mono responseMono = replaceInternal(userSettings); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } + private Mono replaceInternal(CosmosUserProperties userSettings) { + return this.database.getDocClientWrapper() + .replaceUser(ModelBridgeInternal.getV2User(userSettings), null) + .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, database)).single(); + } + private Mono deleteInternal(Context context) { String spanName = "deleteUser." + getId(); - Mono responseMono = this.database.getDocClientWrapper() - .deleteUser(getLink(), null) - .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, database)).single(); + Mono responseMono = deleteInternal(); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } + private Mono deleteInternal() { + return this.database.getDocClientWrapper() + .deleteUser(getLink(), null) + .map(response -> ModelBridgeInternal.createCosmosAsyncUserResponse(response, database)).single(); + } + private Mono createPermissionInternal( Permission permission, CosmosPermissionRequestOptions options, Context context) { String spanName = "createPermission." + getId(); - Mono responseMono = database.getDocClientWrapper() + Mono responseMono = createPermissionInternal(permission, options); + return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, database.getId(), database.getClient().getServiceEndpoint()); + } + + private Mono createPermissionInternal( + Permission permission, + CosmosPermissionRequestOptions options) { + return database.getDocClientWrapper() .createPermission(getLink(), permission, ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncPermissionResponse(response, this)) .single(); - return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, database.getId(), database.getClient().getServiceEndpoint()); } private Mono upsertPermissionInternal( @@ -286,11 +318,17 @@ private Mono upsertPermissionInternal( CosmosPermissionRequestOptions options, Context context) { String spanName = "upsertPermission." + getId(); - Mono responseMono = database.getDocClientWrapper() + Mono responseMono = upsertPermissionInternal(permission, options); + return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, database.getId(), database.getClient().getServiceEndpoint()); + } + + private Mono upsertPermissionInternal( + Permission permission, + CosmosPermissionRequestOptions options) { + return database.getDocClientWrapper() .upsertPermission(getLink(), permission, ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncPermissionResponse(response, this)) .single(); - return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, database.getId(), database.getClient().getServiceEndpoint()); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java index 6137dfb09468..e16ee192c4e8 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java @@ -58,8 +58,11 @@ CosmosAsyncUserDefinedFunction setId(String id) { * @return an {@link Mono} containing the single resource response for the read user defined function or an error. */ public Mono read() { - return TracerProvider.cosmosWithContext(withContext(context -> read(context)), - container.getDatabase().getClient().getTracerProvider()); + if(!container.getDatabase().getClient().getTracerProvider().isEnabled()) { + return readInternal(); + } + + return withContext(context -> readInternal(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -75,8 +78,11 @@ public Mono read() { * or an error. */ public Mono replace(CosmosUserDefinedFunctionProperties udfSettings) { - return TracerProvider.cosmosWithContext(withContext(context -> replace(udfSettings, context)), - container.getDatabase().getClient().getTracerProvider()); + if(!container.getDatabase().getClient().getTracerProvider().isEnabled()) { + return replaceInternal(udfSettings); + } + + return withContext(context -> replaceInternal(udfSettings, context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -91,8 +97,11 @@ public Mono replace(CosmosUserDefinedFun * an error. */ public Mono delete() { - return TracerProvider.cosmosWithContext(withContext(context -> delete(context)), - container.getDatabase().getClient().getTracerProvider()); + if(!container.getDatabase().getClient().getTracerProvider().isEnabled()) { + return deleteInternal(); + } + + return withContext(context -> deleteInternal(context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } String getURIPathSegment() { @@ -113,32 +122,43 @@ String getLink() { return builder.toString(); } - private Mono read(Context context) { + private Mono readInternal(Context context) { String spanName = "readUDF." + container.getId(); - Mono responseMono = - container.getDatabase().getDocClientWrapper().readUserDefinedFunction(getLink(), null) - .map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, container)).single(); + Mono responseMono = readInternal(); return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); } - private Mono replace(CosmosUserDefinedFunctionProperties udfSettings, + private Mono readInternal() { + return container.getDatabase().getDocClientWrapper().readUserDefinedFunction(getLink(), null) + .map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, container)).single(); + } + + private Mono replaceInternal(CosmosUserDefinedFunctionProperties udfSettings, Context context) { String spanName = "replaceUDF." + container.getId(); - Mono responseMono = container.getDatabase() + Mono responseMono = replaceInternal(udfSettings); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); + } + + private Mono replaceInternal(CosmosUserDefinedFunctionProperties udfSettings) { + return container.getDatabase() .getDocClientWrapper() .replaceUserDefinedFunction(new UserDefinedFunction(ModelBridgeInternal.toJsonFromJsonSerializable(udfSettings)), null) .map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, container)) .single(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); } - private Mono delete(Context context) { + private Mono deleteInternal(Context context) { String spanName = "deleteUDF." + container.getId(); - Mono responseMono = container.getDatabase() + Mono responseMono = deleteInternal(); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); + } + + private Mono deleteInternal() { + return container.getDatabase() .getDocClientWrapper() .deleteUserDefinedFunction(this.getLink(), null) .map(response -> ModelBridgeInternal.createCosmosAsyncUserDefinedFunctionResponse(response, container)) .single(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java index 398b08cf65de..2790b37dc3ac 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java @@ -43,22 +43,12 @@ public class CosmosClientBuilder { private boolean sessionCapturingOverrideEnabled; private boolean connectionReuseAcrossClientsEnabled; - private TracerProvider tracerProvider; private boolean contentResponseOnWriteEnabled; /** * Instantiates a new Cosmos client builder. */ public CosmosClientBuilder() { - this.tracerProvider = new TracerProvider(ServiceLoader.load(Tracer.class)); - } - - /** - * Get the tracer provider - * @return tracerProvider - */ - public TracerProvider getTracerProvider() { - return tracerProvider; } /** @@ -307,6 +297,7 @@ public CosmosClientBuilder keyCredential(CosmosKeyCredential cosmosKeyCredential } /** +<<<<<<< Updated upstream * Gets the boolean which indicates whether to only return the headers and status code in Cosmos DB response * in case of Create, Update and Delete operations on CosmosItem. * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java index 078c297daf29..37a74912594b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java @@ -15,6 +15,7 @@ import java.io.PrintWriter; import java.io.StringWriter; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.Objects; import java.util.concurrent.atomic.AtomicReference; @@ -24,6 +25,7 @@ public class TracerProvider { private final List tracers = new ArrayList<>(); + private final boolean isEnabled; public final static String DB_TYPE_VALUE = "cosmosdb"; public final static String DB_TYPE = "db.type"; public final static String DB_INSTANCE = "db.instance"; @@ -35,6 +37,7 @@ public class TracerProvider { public static final String COSMOS_CALL_DEPTH = "cosmosCallDepth"; public static final String MICROSOFT_DOCOMENTDB = "Microsoft.DocumentDB"; + public final static Function callDepthAttributeFunc = ( reactor.util.context.Context reactorContext) -> { CallDepth callDepth = reactorContext.getOrDefault(COSMOS_CALL_DEPTH, CallDepth.ZERO); @@ -48,21 +51,14 @@ public class TracerProvider { } }; - public static Mono cosmosWithContext(Mono mono, TracerProvider tracerProvider) { - if(tracerProvider.isEnabled()) { - return mono.subscriberContext(TracerProvider.callDepthAttributeFunc); - } else { - return mono; - } - } - public TracerProvider(Iterable tracers) { Objects.requireNonNull(tracers, "'tracers' cannot be null."); tracers.forEach(this.tracers::add); + isEnabled = this.tracers.size() > 0; } public boolean isEnabled() { - return tracers.size() > 0; + return isEnabled; } /** @@ -159,26 +155,23 @@ public Mono traceEnabledPublisher(Mono resultPublisher, String databaseId, String endpoint, Function statusCodeFunc) { - final boolean isTracingEnabled = this.isEnabled(); - final AtomicReference parentContext = isTracingEnabled - ? new AtomicReference<>(Context.NONE) - : null; + final AtomicReference parentContext = new AtomicReference<>(Context.NONE); CallDepth callDepth = FluxUtil.toReactorContext(context).getOrDefault(COSMOS_CALL_DEPTH, CallDepth.ZERO); assert callDepth != null; final boolean isNestedCall = callDepth.equals(CallDepth.TWO_OR_MORE); return resultPublisher .doOnSubscribe(ignoredValue -> { - if (isTracingEnabled && !isNestedCall) { + if (!isNestedCall) { parentContext.set(this.startSpan(spanName, databaseId, endpoint, context)); } }).doOnSuccess(response -> { - if (isTracingEnabled && !isNestedCall) { + if (!isNestedCall) { this.endSpan(parentContext.get(), Signal.complete(), statusCodeFunc.apply(response)); } }).doOnError(throwable -> { - if (isTracingEnabled && !isNestedCall) { + if (!isNestedCall) { this.endSpan(parentContext.get(), Signal.error(throwable), 0); } }); From 6b09e5c2504b94e80788d30048c96836bd20c1cb Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Mon, 11 May 2020 11:18:52 -0400 Subject: [PATCH 14/40] fixing checkstyle --- .../com/azure/cosmos/CosmosAsyncDatabase.java | 39 +++++++++++-------- .../com/azure/cosmos/CosmosClientBuilder.java | 5 --- .../cosmos/implementation/TracerProvider.java | 1 - 3 files changed, 22 insertions(+), 23 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java index d8ee89cb4bf1..fb09250cd672 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java @@ -14,7 +14,6 @@ import com.azure.cosmos.models.CosmosAsyncUserResponse; import com.azure.cosmos.models.CosmosContainerProperties; import com.azure.cosmos.models.CosmosContainerRequestOptions; -import com.azure.cosmos.models.CosmosDatabaseProperties; import com.azure.cosmos.models.CosmosDatabaseRequestOptions; import com.azure.cosmos.models.CosmosUserProperties; import com.azure.cosmos.models.FeedOptions; @@ -28,10 +27,10 @@ import reactor.core.Exceptions; import reactor.core.publisher.Mono; -import static com.azure.core.util.FluxUtil.withContext; import java.util.Collections; import java.util.List; +import static com.azure.core.util.FluxUtil.withContext; import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; /** @@ -297,7 +296,13 @@ public Mono createContainer(String id, String part public Mono createContainerIfNotExists( CosmosContainerProperties containerProperties) { CosmosAsyncContainer container = getContainer(containerProperties.getId()); - return createContainerIfNotExistsInternal(containerProperties, container, null); + if (!client.getTracerProvider().isEnabled()) { + return createContainerIfNotExistsInternal(containerProperties, container, + null); + } + + return withContext(context -> createContainerIfNotExistsInternal(containerProperties, container, null, + context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -319,7 +324,12 @@ public Mono createContainerIfNotExists( CosmosContainerRequestOptions options = new CosmosContainerRequestOptions(); ModelBridgeInternal.setOfferThroughput(options, throughput); CosmosAsyncContainer container = getContainer(containerProperties.getId()); - return createContainerIfNotExistsInternal(containerProperties, container, options); + if (!client.getTracerProvider().isEnabled()) { + return createContainerIfNotExistsInternal(containerProperties, container, options); + } + + return withContext(context -> createContainerIfNotExistsInternal(containerProperties, container, options, + context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -336,9 +346,13 @@ public Mono createContainerIfNotExists( */ public Mono createContainerIfNotExists(String id, String partitionKeyPath) { CosmosAsyncContainer container = getContainer(id); - return createContainerIfNotExistsInternal(new CosmosContainerProperties(id, partitionKeyPath), - container, - null); + if (!client.getTracerProvider().isEnabled()) { + return createContainerIfNotExistsInternal(new CosmosContainerProperties(id, partitionKeyPath), container, + null); + } + + return withContext(context -> createContainerIfNotExistsInternal(new CosmosContainerProperties(id, partitionKeyPath), container, null, + context)).subscriberContext(TracerProvider.callDepthAttributeFunc); } /** @@ -796,16 +810,7 @@ private Mono createContainerIfNotExistsInternal( CosmosContainerRequestOptions options, Context context) { String spanName = "createContainerIfNotExistsInternal." + containerProperties.getId(); - Mono responseMono = container.read(options).onErrorResume(exception -> { - final Throwable unwrappedException = Exceptions.unwrap(exception); - if (unwrappedException instanceof CosmosClientException) { - final CosmosClientException cosmosClientException = (CosmosClientException) unwrappedException; - if (cosmosClientException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) { - return createContainer(containerProperties, options); - } - } - return Mono.error(unwrappedException); - }); + Mono responseMono = createContainerIfNotExistsInternal(containerProperties, container, options); return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, getId(), getClient().getServiceEndpoint()); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java index 2790b37dc3ac..0435020dc3e8 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java @@ -3,15 +3,12 @@ package com.azure.cosmos; import com.azure.core.annotation.ServiceClientBuilder; -import com.azure.core.util.tracing.Tracer; import com.azure.cosmos.implementation.Configs; -import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.implementation.CosmosAuthorizationTokenResolver; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.models.CosmosPermissionProperties; import java.util.List; -import java.util.ServiceLoader; /** * Helper class to buildAsyncClient {@link CosmosAsyncClient} instances @@ -42,7 +39,6 @@ public class CosmosClientBuilder { private CosmosKeyCredential cosmosKeyCredential; private boolean sessionCapturingOverrideEnabled; private boolean connectionReuseAcrossClientsEnabled; - private boolean contentResponseOnWriteEnabled; /** @@ -297,7 +293,6 @@ public CosmosClientBuilder keyCredential(CosmosKeyCredential cosmosKeyCredential } /** -<<<<<<< Updated upstream * Gets the boolean which indicates whether to only return the headers and status code in Cosmos DB response * in case of Create, Update and Delete operations on CosmosItem. * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java index 37a74912594b..2661004a9e92 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java @@ -15,7 +15,6 @@ import java.io.PrintWriter; import java.io.StringWriter; import java.util.ArrayList; -import java.util.Collections; import java.util.List; import java.util.Objects; import java.util.concurrent.atomic.AtomicReference; From 81040b270c3a369ae5289428143c144b9e57d54c Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Mon, 11 May 2020 14:09:41 -0400 Subject: [PATCH 15/40] resolving comments --- .../src/main/java/com/azure/cosmos/CosmosAsyncClient.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index e03c64535981..bbe11ede3be0 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -473,7 +473,7 @@ private CosmosPagedFlux queryDatabasesInternal(boolean private Mono createDatabaseIfNotExistsInternal(CosmosAsyncDatabase database, Context context) { - String spanName = "createDatabaseIfNotExistsInternal." + database.getId(); + String spanName = "createDatabaseIfNotExists." + database.getId(); Mono responseMono = database.read().onErrorResume(exception -> { final Throwable unwrappedException = Exceptions.unwrap(exception); if (unwrappedException instanceof CosmosClientException) { From 806cdf154862a36147c6bc3cfebe582e6fa323b3 Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Mon, 11 May 2020 14:56:09 -0400 Subject: [PATCH 16/40] check style fix as per java 8 --- .../java/com/azure/cosmos/implementation/TracerProvider.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java index 2661004a9e92..2a08d5e2fb8d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java @@ -128,7 +128,7 @@ public > Mono traceEnabledCosmos String databaseId, String endpoint) { return traceEnabledPublisher(resultPublisher, context, spanName,databaseId, endpoint, - CosmosResponse::getStatusCode); + (T response) -> response.getStatusCode()); } public Mono traceEnabledNonCosmosResponsePublisher(Mono resultPublisher, From d061db5abcbf948ae4159576f7ee94a75eca574f Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Mon, 11 May 2020 17:11:36 -0400 Subject: [PATCH 17/40] adding azure-core-tracing-opentelemetry in test scope --- sdk/cosmos/azure-cosmos/pom.xml | 13 +++++++++++++ .../java/com/azure/cosmos/CosmosTracerTest.java | 6 +++--- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/pom.xml b/sdk/cosmos/azure-cosmos/pom.xml index 763d5827f705..af40f300d83e 100644 --- a/sdk/cosmos/azure-cosmos/pom.xml +++ b/sdk/cosmos/azure-cosmos/pom.xml @@ -132,6 +132,19 @@ Licensed under the MIT License. + + com.azure + azure-core-tracing-opentelemetry + test + 1.0.0-beta.4 + + + io.netty + * + + + + com.fasterxml.jackson.core jackson-core diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java index f53375fcb012..5178f2099dec 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java @@ -53,21 +53,21 @@ public void cosmosAsyncClient() { client.createDatabaseIfNotExists(cosmosAsyncDatabase.getId()).block(); Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); - Mockito.verify(tracer, Mockito.times(2)).endSpan(Matchers.any(Context.class), Matchers.>>any(), + Mockito.verify(tracer, Mockito.times(1)).endSpan(Matchers.any(Context.class), Matchers.>>any(), Matchers.anyInt()); //endSpan invocation is 1 greater than startSpan, because createDatabaseIfNotExists // uses public read api client.readAllDatabases(new FeedOptions()).byPage().single().block(); Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); - Mockito.verify(tracer, Mockito.atLeast(3)).endSpan(Matchers.any(Context.class), Matchers.>>any(), + Mockito.verify(tracer, Mockito.atLeast(2)).endSpan(Matchers.any(Context.class), Matchers.>>any(), Matchers.anyInt()); //We used atLeast in endSpan because for query we are calling endSpan on every doOnEach String query = "select * from c where c.id = '" + cosmosAsyncDatabase.getId() + "'"; client.queryDatabases(query, new FeedOptions()).byPage().single().block(); Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); - Mockito.verify(tracer, Mockito.atLeast(4)).endSpan(Matchers.any(Context.class), Matchers.>>any(), + Mockito.verify(tracer, Mockito.atLeast(3)).endSpan(Matchers.any(Context.class), Matchers.>>any(), Matchers.anyInt()); //We used atLeast in endSpan because for query we are calling endSpan on every doOnEach } From 0fc0cc923557735ac777c7037e5a5491e6675467 Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Mon, 11 May 2020 19:38:49 -0400 Subject: [PATCH 18/40] build error fix --- .../src/test/java/com/azure/cosmos/rx/TestSuiteBase.java | 1 - 1 file changed, 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java index 2b58b0fb2d99..e6c4d05a50d4 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java @@ -22,7 +22,6 @@ import com.azure.cosmos.CosmosClientException; import com.azure.cosmos.models.CosmosContainerProperties; import com.azure.cosmos.models.CosmosContainerRequestOptions; -import com.azure.cosmos.models.JsonSerializable; import com.azure.cosmos.models.ModelBridgeInternal; import com.azure.cosmos.util.CosmosPagedFlux; import com.azure.cosmos.CosmosDatabase; From 1290ae13cd5eb1995ddd3a4bbc9a705456fafd9a Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Mon, 11 May 2020 21:19:20 -0400 Subject: [PATCH 19/40] test failure fix --- .../java/com/azure/cosmos/CosmosAsyncStoredProcedure.java | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java index 8f20fba7d01e..65283817fa57 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java @@ -258,6 +258,10 @@ private Mono replaceInternal(CosmosStoredPro private Mono replaceInternal(CosmosStoredProcedureProperties storedProcedureSettings, CosmosStoredProcedureRequestOptions options) { + if (options == null) { + options = new CosmosStoredProcedureRequestOptions(); + } + return cosmosContainer.getDatabase() .getDocClientWrapper() .replaceStoredProcedure(new StoredProcedure(ModelBridgeInternal.toJsonFromJsonSerializable( From d2a815459cf553a0f24126b0765e5b3a375ccd4b Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Mon, 11 May 2020 22:00:54 -0400 Subject: [PATCH 20/40] test fix --- .../azure/cosmos/rx/ReadFeedExceptionHandlingTest.java | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/ReadFeedExceptionHandlingTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/ReadFeedExceptionHandlingTest.java index 2fa9ff8b01c0..a38000b6edaa 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/ReadFeedExceptionHandlingTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/ReadFeedExceptionHandlingTest.java @@ -2,13 +2,15 @@ // Licensed under the MIT License. package com.azure.cosmos.rx; +import com.azure.core.util.tracing.Tracer; import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.CosmosAsyncClient; import com.azure.cosmos.CosmosClientBuilder; -import com.azure.cosmos.util.CosmosPagedFlux; +import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.models.CosmosDatabaseProperties; import com.azure.cosmos.models.FeedOptions; import com.azure.cosmos.models.FeedResponse; +import com.azure.cosmos.util.CosmosPagedFlux; import com.azure.cosmos.util.UtilBridgeInternal; import io.reactivex.subscribers.TestSubscriber; import org.mockito.Mockito; @@ -19,6 +21,7 @@ import reactor.core.publisher.Flux; import java.util.ArrayList; +import java.util.ServiceLoader; import static org.assertj.core.api.Assertions.assertThat; @@ -47,7 +50,10 @@ public void readFeedException() throws Exception { .mergeWith(Flux.fromIterable(frps)); final CosmosAsyncClientWrapper mockedClientWrapper = Mockito.spy(new CosmosAsyncClientWrapper(client)); - Mockito.when(mockedClientWrapper.readAllDatabases(null)).thenReturn(UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> response)); + Mockito.when(mockedClientWrapper.readAllDatabases(null)).thenReturn(UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + pagedFluxOptions.setTracerInformation(new TracerProvider(ServiceLoader.load(Tracer.class)), "testSpan", "testEndpoint,", "testDb"); + return response; + })); TestSubscriber> subscriber = new TestSubscriber<>(); mockedClientWrapper.readAllDatabases(null).byPage().subscribe(subscriber); assertThat(subscriber.valueCount()).isEqualTo(2); From e7fa831bfcdc284440ff14a26ec84bd02732c79c Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Wed, 13 May 2020 13:38:31 -0400 Subject: [PATCH 21/40] complie error fix --- .../java/com/azure/cosmos/implementation/TracerProvider.java | 2 +- .../src/test/java/com/azure/cosmos/CosmosTracerTest.java | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java index fddd4e502b99..42c87add6fd1 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java @@ -25,7 +25,7 @@ public class TracerProvider { private final List tracers = new ArrayList<>(); private final boolean isEnabled; - public final static String DB_TYPE_VALUE = "cosmos"; + public final static String DB_TYPE_VALUE = "Cosmos"; public final static String DB_TYPE = "db.type"; public final static String DB_INSTANCE = "db.instance"; public final static String DB_URL = "db.url"; diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java index 43ae1e1c8143..6069057f99d6 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java @@ -38,6 +38,7 @@ public void beforeClass() { client = new CosmosClientBuilder() .endpoint(TestConfigurations.HOST) .key(TestConfigurations.MASTER_KEY) + .directMode(DirectConnectionConfig.getDefaultConfig()) .buildAsyncClient(); cosmosAsyncDatabase = getSharedCosmosDatabase(client); cosmosAsyncContainer = getSharedMultiPartitionCosmosContainer(client); From 3aa851d64c0ee1d43d1d86994d4e19a677ac7dc8 Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Fri, 15 May 2020 10:51:03 -0400 Subject: [PATCH 22/40] removing query text from tracer --- sdk/cosmos/azure-cosmos/pom.xml | 4 +-- .../com/azure/cosmos/CosmosAsyncClient.java | 14 +++----- .../azure/cosmos/CosmosAsyncContainer.java | 18 ++++------ .../com/azure/cosmos/CosmosAsyncDatabase.java | 34 ++++++------------- .../com/azure/cosmos/CosmosAsyncScripts.java | 24 +++---------- 5 files changed, 28 insertions(+), 66 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/pom.xml b/sdk/cosmos/azure-cosmos/pom.xml index fcb4f63dbf4c..29ecd116d22c 100644 --- a/sdk/cosmos/azure-cosmos/pom.xml +++ b/sdk/cosmos/azure-cosmos/pom.xml @@ -134,9 +134,9 @@ Licensed under the MIT License. com.azure - azure-core-tracing-opentelemetry + azure-core-tracing-opentelemetry test - 1.0.0-beta.4 + 1.0.0-beta.4 io.netty diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index 8d735b5c2849..5bf4aee03b5b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -413,7 +413,7 @@ public CosmosPagedFlux readAllDatabases() { * @return a {@link CosmosPagedFlux} containing one or several feed response pages of read databases or an error. */ public CosmosPagedFlux queryDatabases(String query, FeedOptions options) { - return queryDatabasesInternal(false, new SqlQuerySpec(query), options); + return queryDatabasesInternal(new SqlQuerySpec(query), options); } /** @@ -428,7 +428,7 @@ public CosmosPagedFlux queryDatabases(String query, Fe * @return a {@link CosmosPagedFlux} containing one or several feed response pages of read databases or an error. */ public CosmosPagedFlux queryDatabases(SqlQuerySpec querySpec, FeedOptions options) { - return queryDatabasesInternal(true, querySpec, options); + return queryDatabasesInternal(querySpec, options); } /** @@ -453,15 +453,9 @@ TracerProvider getTracerProvider(){ return this.tracerProvider; } - private CosmosPagedFlux queryDatabasesInternal(boolean isParameterised, SqlQuerySpec querySpec, FeedOptions options){ + private CosmosPagedFlux queryDatabasesInternal(SqlQuerySpec querySpec, FeedOptions options){ return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName; - if (isParameterised) { - spanName = "queryDatabases." + querySpec.getQueryText(); - } else { - spanName = "queryDatabases"; - } - + String spanName = "queryDatabases"; pagedFluxOptions.setTracerInformation(this.tracerProvider, spanName, this.serviceEndpoint, null); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().queryDatabases(querySpec, options) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index a99658973234..4abba56c2bd6 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -381,7 +381,7 @@ public CosmosPagedFlux readAllItems(FeedOptions options, Class classTy * error. */ public CosmosPagedFlux queryItems(String query, Class classType) { - return queryItemsInternal(false, new SqlQuerySpec(query), new FeedOptions(), classType); + return queryItemsInternal(new SqlQuerySpec(query), new FeedOptions(), classType); } /** @@ -399,7 +399,7 @@ public CosmosPagedFlux queryItems(String query, Class classType) { * error. */ public CosmosPagedFlux queryItems(String query, FeedOptions options, Class classType) { - return queryItemsInternal(false, new SqlQuerySpec(query), options, classType); + return queryItemsInternal(new SqlQuerySpec(query), options, classType); } /** @@ -416,7 +416,7 @@ public CosmosPagedFlux queryItems(String query, FeedOptions options, Clas * error. */ public CosmosPagedFlux queryItems(SqlQuerySpec querySpec, Class classType) { - return queryItemsInternal(true, querySpec, new FeedOptions(), classType); + return queryItemsInternal(querySpec, new FeedOptions(), classType); } /** @@ -434,19 +434,13 @@ public CosmosPagedFlux queryItems(SqlQuerySpec querySpec, Class classT * error. */ public CosmosPagedFlux queryItems(SqlQuerySpec querySpec, FeedOptions options, Class classType) { - return queryItemsInternal(true, querySpec, options, classType); + return queryItemsInternal(querySpec, options, classType); } private CosmosPagedFlux queryItemsInternal( - boolean isParameterised, SqlQuerySpec sqlQuerySpec, FeedOptions feedOptions, Class classType) { + SqlQuerySpec sqlQuerySpec, FeedOptions feedOptions, Class classType) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName; - if (isParameterised) { - spanName = this.queryItemsSpanName + "." + sqlQuerySpec.getQueryText(); - } else { - spanName = this.queryItemsSpanName; - } - + String spanName = this.queryItemsSpanName; pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), spanName, this.getDatabase().getClient().getServiceEndpoint(), database.getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, feedOptions); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java index f00ad63fd615..3d5474bda9f8 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java @@ -433,7 +433,7 @@ public CosmosPagedFlux readAllContainers() { * obtained containers or an error. */ public CosmosPagedFlux queryContainers(String query) { - return queryContainersInternal(false, new SqlQuerySpec(query), new FeedOptions()); + return queryContainersInternal(new SqlQuerySpec(query), new FeedOptions()); } /** @@ -449,7 +449,7 @@ public CosmosPagedFlux queryContainers(String query) * obtained containers or an error. */ public CosmosPagedFlux queryContainers(String query, FeedOptions options) { - return queryContainersInternal(false, new SqlQuerySpec(query), options); + return queryContainersInternal(new SqlQuerySpec(query), options); } /** @@ -464,7 +464,7 @@ public CosmosPagedFlux queryContainers(String query, * obtained containers or an error. */ public CosmosPagedFlux queryContainers(SqlQuerySpec querySpec) { - return queryContainersInternal(true, querySpec, new FeedOptions()); + return queryContainersInternal(querySpec, new FeedOptions()); } /** @@ -481,7 +481,7 @@ public CosmosPagedFlux queryContainers(SqlQuerySpec q */ public CosmosPagedFlux queryContainers(SqlQuerySpec querySpec , FeedOptions options) { - return queryContainersInternal(true, querySpec, options); + return queryContainersInternal(querySpec, options); } /** @@ -599,7 +599,7 @@ public CosmosPagedFlux queryUsers(String query) { * obtained users or an error. */ public CosmosPagedFlux queryUsers(String query, FeedOptions options) { - return queryUsersInternal(false, new SqlQuerySpec(query), options); + return queryUsersInternal(new SqlQuerySpec(query), options); } /** @@ -614,7 +614,7 @@ public CosmosPagedFlux queryUsers(String query, FeedOption * obtained users or an error. */ public CosmosPagedFlux queryUsers(SqlQuerySpec querySpec) { - return queryUsersInternal(true, querySpec, new FeedOptions()); + return queryUsersInternal(querySpec, new FeedOptions()); } /** @@ -630,7 +630,7 @@ public CosmosPagedFlux queryUsers(SqlQuerySpec querySpec) * obtained users or an error. */ public CosmosPagedFlux queryUsers(SqlQuerySpec querySpec, FeedOptions options) { - return queryUsersInternal(true, querySpec, options); + return queryUsersInternal(querySpec, options); } /** @@ -765,16 +765,10 @@ String getLink() { return this.link; } - public CosmosPagedFlux queryContainersInternal(boolean isParameterised, SqlQuerySpec querySpec + public CosmosPagedFlux queryContainersInternal(SqlQuerySpec querySpec , FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName; - if (isParameterised) { - spanName = "queryContainers." + this.getId() + "." + querySpec.getQueryText(); - } else { - spanName = "queryContainers." + this.getId(); - } - + String spanName = "queryContainers." + this.getId(); pagedFluxOptions.setTracerInformation(this.getClient().getTracerProvider(), spanName, this.getClient().getServiceEndpoint(), getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); @@ -785,15 +779,9 @@ public CosmosPagedFlux queryContainersInternal(boolea }); } - private CosmosPagedFlux queryUsersInternal(boolean isParameterised, SqlQuerySpec querySpec, FeedOptions options) { + private CosmosPagedFlux queryUsersInternal(SqlQuerySpec querySpec, FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName; - if (isParameterised) { - spanName = "queryUsers." + this.getId() + "." + querySpec.getQueryText(); - } else { - spanName = "queryUsers." + this.getId(); - } - + String spanName = "queryUsers." + this.getId(); pagedFluxOptions.setTracerInformation(this.getClient().getTracerProvider(), spanName, this.getClient().getServiceEndpoint(), getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java index 6ed36f6b0888..34adcd145f82 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java @@ -128,7 +128,7 @@ public CosmosPagedFlux readAllStoredProcedures( public CosmosPagedFlux queryStoredProcedures( String query, FeedOptions options) { - return queryStoredProceduresInternal(false, new SqlQuerySpec(query), options); + return queryStoredProceduresInternal(new SqlQuerySpec(query), options); } /** @@ -147,7 +147,7 @@ public CosmosPagedFlux queryStoredProcedures( public CosmosPagedFlux queryStoredProcedures( SqlQuerySpec querySpec, FeedOptions options) { - return queryStoredProceduresInternal(true, querySpec, options); + return queryStoredProceduresInternal(querySpec, options); } /** @@ -250,7 +250,7 @@ public CosmosPagedFlux queryUserDefinedFunc public CosmosPagedFlux queryUserDefinedFunctions( SqlQuerySpec querySpec, FeedOptions options) { - return queryUserDefinedFunctionsInternal(true, querySpec, options); + return queryUserDefinedFunctionsInternal(querySpec, options); } /** @@ -357,17 +357,10 @@ public CosmosAsyncTrigger getTrigger(String id) { } private CosmosPagedFlux queryStoredProceduresInternal( - boolean isParameterised, SqlQuerySpec querySpec, FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName; - if (isParameterised) { - spanName = "queryStoredProcedures." + this.container.getId() + "." + querySpec.getQueryText(); - } else { - spanName = "queryStoredProcedures." + this.container.getId(); - } - + String spanName = "queryStoredProcedures." + this.container.getId(); pagedFluxOptions.setTracerInformation(this.container.getDatabase().getClient().getTracerProvider(), spanName, this.container.getDatabase().getClient().getServiceEndpoint(), @@ -382,17 +375,10 @@ private CosmosPagedFlux queryStoredProceduresIn } private CosmosPagedFlux queryUserDefinedFunctionsInternal( - boolean isParameterised, SqlQuerySpec querySpec, FeedOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName; - if (isParameterised) { - spanName = "queryUserDefinedFunctions." + this.container.getId() + "." + querySpec.getQueryText(); - } else { - spanName = "queryUserDefinedFunctions." + this.container.getId(); - } - + String spanName = "queryUserDefinedFunctions." + this.container.getId(); pagedFluxOptions.setTracerInformation(this.container.getDatabase().getClient().getTracerProvider(), spanName, this.container.getDatabase().getClient().getServiceEndpoint(), From 0984a9ceadd61b2c0862a4a6ae3b968693952d57 Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Thu, 11 Jun 2020 18:11:06 -0400 Subject: [PATCH 23/40] code formating change and some comment resolution --- .../com/azure/cosmos/CosmosAsyncClient.java | 60 +++++++++++------- .../com/azure/cosmos/CosmosAsyncConflict.java | 45 ++++++++++++-- .../azure/cosmos/CosmosAsyncContainer.java | 44 ++++++++++--- .../com/azure/cosmos/CosmosAsyncDatabase.java | 62 +++++++++++++------ .../azure/cosmos/CosmosAsyncPermission.java | 12 +++- .../com/azure/cosmos/CosmosAsyncScripts.java | 24 +++++-- .../cosmos/CosmosAsyncStoredProcedure.java | 12 +++- .../com/azure/cosmos/CosmosAsyncTrigger.java | 18 +++++- .../com/azure/cosmos/CosmosAsyncUser.java | 20 ++++-- .../CosmosAsyncUserDefinedFunction.java | 18 +++++- 10 files changed, 236 insertions(+), 79 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index 59d6294cbcf7..473c45d6d25d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -14,9 +14,9 @@ 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; @@ -204,11 +204,11 @@ boolean isContentResponseOnWriteEnabled() { public Mono createDatabaseIfNotExists(CosmosDatabaseProperties databaseProperties) { if(!getTracerProvider().isEnabled()) { CosmosAsyncDatabase database = getDatabase(databaseProperties.getId()); - return createDatabaseIfNotExistsInternal(database.read(), database); + return createDatabaseIfNotExistsInternal(database.read(), database, null, null); } return withContext(context -> createDatabaseIfNotExistsInternal(getDatabase(databaseProperties.getId()), - context)); + null, context)); } /** @@ -224,10 +224,10 @@ public Mono createDatabaseIfNotExists(CosmosDatabaseProp public Mono createDatabaseIfNotExists(String id) { if(!getTracerProvider().isEnabled()) { CosmosAsyncDatabase database = getDatabase(id); - return createDatabaseIfNotExistsInternal(database.read(), database); + return createDatabaseIfNotExistsInternal(database.read(), database, null,null); } - return withContext(context -> createDatabaseIfNotExistsInternal(getDatabase(id), context)); + return withContext(context -> createDatabaseIfNotExistsInternal(getDatabase(id), null, context)); } /** @@ -244,19 +244,13 @@ public Mono createDatabaseIfNotExists(String id) { * @return the mono. */ public Mono createDatabaseIfNotExists(String id, ThroughputProperties throughputProperties) { - 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); - }); + if(!getTracerProvider().isEnabled()) { + CosmosAsyncDatabase database = getDatabase(id); + return createDatabaseIfNotExistsInternal(database.read(), database, throughputProperties, null); + } + + return withContext(context -> createDatabaseIfNotExistsInternal(getDatabase(id), + throughputProperties, context)); } /** @@ -499,22 +493,37 @@ private CosmosPagedFlux queryDatabasesInternal(SqlQuer private Mono createDatabaseIfNotExistsInternal(CosmosAsyncDatabase database, - Context context) { + ThroughputProperties throughputProperties, Context context) { String spanName = "createDatabaseIfNotExists." + database.getId(); Context nestedContext = context.addData(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.COSMOS_CALL_DEPTH_VAL); - Mono responseMono = createDatabaseIfNotExistsInternal(database.readInternal(new CosmosDatabaseRequestOptions(), nestedContext), database); - return tracerProvider.traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), + Mono responseMono = createDatabaseIfNotExistsInternal(database.readInternal(new CosmosDatabaseRequestOptions(), nestedContext), database, throughputProperties, nestedContext); + return tracerProvider.traceEnabledCosmosResponsePublisher(responseMono, + context, + spanName, + database.getId(), this.serviceEndpoint); } - private Mono createDatabaseIfNotExistsInternal(Mono responseMono, CosmosAsyncDatabase database) { + private Mono createDatabaseIfNotExistsInternal(Mono 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()), - new CosmosDatabaseRequestOptions()); + requestOptions); } } return Mono.error(unwrappedException); @@ -526,7 +535,10 @@ private Mono createDatabaseInternal(Database database, C Context context) { String spanName = "createDatabase." + database.getId(); Mono responseMono = createDatabaseInternal(database, options); - return tracerProvider.traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), + return tracerProvider.traceEnabledCosmosResponsePublisher(responseMono, + context, + spanName, + database.getId(), this.serviceEndpoint); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncConflict.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncConflict.java index 9e7a336e782f..96f1edaaa7ed 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncConflict.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncConflict.java @@ -2,6 +2,7 @@ // 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; @@ -9,6 +10,8 @@ import com.azure.cosmos.models.ModelBridgeInternal; import reactor.core.publisher.Mono; +import static com.azure.core.util.FluxUtil.withContext; + /** * Read and delete conflicts */ @@ -64,9 +67,11 @@ public Mono read(CosmosConflictRequestOptions options) { options = new CosmosConflictRequestOptions(); } RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); - return this.container.getDatabase().getDocClientWrapper().readConflict(getLink(), requestOptions) - .map(response -> ModelBridgeInternal.createCosmosConflictResponse(response)).single(); + if (!this.container.getDatabase().getClient().getTracerProvider().isEnabled()) { + return readInternal(requestOptions); + } + return withContext(context -> readInternal(requestOptions, context)); } /** @@ -85,8 +90,11 @@ public Mono delete(CosmosConflictRequestOptions options) options = new CosmosConflictRequestOptions(); } RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); - return this.container.getDatabase().getDocClientWrapper().deleteConflict(getLink(), requestOptions) - .map(response -> ModelBridgeInternal.createCosmosConflictResponse(response)).single(); + if (!this.container.getDatabase().getClient().getTracerProvider().isEnabled()) { + return deleteInternal(requestOptions); + } + + return withContext(context -> deleteInternal(requestOptions, context)); } String getURIPathSegment() { @@ -106,4 +114,33 @@ String getLink() { builder.append(getId()); return builder.toString(); } + + private Mono readInternal(RequestOptions options, Context context) { + String spanName = "readConflict." + getId(); + Mono 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 readInternal(RequestOptions options) { + return this.container.getDatabase().getDocClientWrapper().readConflict(getLink(), options) + .map(response -> ModelBridgeInternal.createCosmosConflictResponse(response)).single(); + } + + private Mono deleteInternal(RequestOptions options, Context context) { + String spanName = "deleteConflict." + getId(); + Mono responseMono = deleteInternal(options); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, + this.container.getDatabase().getId(), + this.container.getDatabase().getClient().getServiceEndpoint()); + } + + private Mono deleteInternal(RequestOptions options) { + return this.container.getDatabase().getDocClientWrapper().deleteConflict(getLink(), options) + .map(response -> ModelBridgeInternal.createCosmosConflictResponse(response)).single(); + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index d52028cc3ceb..e08efdd4d8c4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -754,7 +754,10 @@ private Mono> deleteItemInternal( Context context) { Mono> responseMono = deleteItemInternal(itemId, requestOptions); return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, - context, this.deleteItemSpanName, database.getId(), database.getClient().getServiceEndpoint()); + context, + this.deleteItemSpanName, + database.getId(), + database.getClient().getServiceEndpoint()); } private Mono> deleteItemInternal( @@ -793,7 +796,10 @@ private Mono> replaceItemInternal( private Mono> upsertItemInternal(T item, CosmosItemRequestOptions options, Context context) { Mono> responseMono = upsertItemInternal(item, options); return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, - context, this.upsertItemSpanName, database.getId(), database.getClient().getServiceEndpoint()); + context, + this.upsertItemSpanName, + database.getId(), + database.getClient().getServiceEndpoint()); } private Mono> upsertItemInternal(T item, CosmosItemRequestOptions options) { @@ -813,7 +819,10 @@ private Mono> readItemInternal( Context context) { Mono> responseMono = readItemInternal(itemId, requestOptions, itemType); return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, - context, this.readItemSpanName, database.getId(), database.getClient().getServiceEndpoint()); + context, + this.readItemSpanName, + database.getId(), + database.getClient().getServiceEndpoint()); } private Mono> readItemInternal( @@ -828,7 +837,10 @@ private Mono> readItemInternal( Mono read(CosmosContainerRequestOptions options, Context context) { Mono responseMono = readInternal(options); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, this.readContainerSpanName, database.getId(), database.getClient().getServiceEndpoint()); + context, + this.readContainerSpanName, + database.getId(), + database.getClient().getServiceEndpoint()); } private Mono readInternal(CosmosContainerRequestOptions options) { @@ -840,7 +852,10 @@ private Mono readInternal(CosmosContainerRequestOptions private Mono deleteInternal(CosmosContainerRequestOptions options, Context context) { Mono responseMono = deleteInternal(options); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, this.deleteContainerSpanName, database.getId(), database.getClient().getServiceEndpoint()); + context, + this.deleteContainerSpanName, + database.getId(), + database.getClient().getServiceEndpoint()); } private Mono deleteInternal(CosmosContainerRequestOptions options) { @@ -854,7 +869,10 @@ private Mono replaceInternal(CosmosContainerProperties Context context) { Mono responseMono = replaceInternal(containerProperties, options); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, this.replaceContainerSpanName, database.getId(), database.getClient().getServiceEndpoint()); + context, + this.replaceContainerSpanName, + database.getId(), + database.getClient().getServiceEndpoint()); } private Mono replaceInternal(CosmosContainerProperties containerProperties, @@ -869,8 +887,11 @@ private Mono readThroughputInternal(Context context) { Context nestedContext = context.addData(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.COSMOS_CALL_DEPTH_VAL); Mono responseMono = readThroughputInternal(this.read(new CosmosContainerRequestOptions(), nestedContext)); - return this.getDatabase().getClient().getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono - , context, this.readThroughputSpanName, database.getId(), database.getClient().getServiceEndpoint()); + return this.getDatabase().getClient().getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono, + context, + this.readThroughputSpanName, + database.getId(), + database.getClient().getServiceEndpoint()); } private Mono readThroughputInternal(Mono responseMono) { @@ -903,8 +924,11 @@ private Mono replaceThroughputInternal(ThroughputProperties Mono responseMono = replaceThroughputInternal(this.read(new CosmosContainerRequestOptions(), nestedContext), throughputProperties); - return this.getDatabase().getClient().getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono - , context, this.replaceThroughputSpanName, database.getId(), database.getClient().getServiceEndpoint()); + return this.getDatabase().getClient().getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono, + context, + this.replaceThroughputSpanName, + database.getId(), + database.getClient().getServiceEndpoint()); } private Mono replaceThroughputInternal(Mono responseMono, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java index 6eb23ab98a74..6f09a2f43f61 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java @@ -293,7 +293,7 @@ public Mono createContainerIfNotExists( CosmosContainerProperties containerProperties) { CosmosAsyncContainer container = getContainer(containerProperties.getId()); if (!client.getTracerProvider().isEnabled()) { - return createContainerIfNotExistsInternal(container.read(), containerProperties, null); + return createContainerIfNotExistsInternal(container.read(), containerProperties, null, null); } return withContext(context -> createContainerIfNotExistsInternal(containerProperties, container, null, @@ -323,7 +323,7 @@ Mono createContainerIfNotExists( ModelBridgeInternal.setThroughputProperties(options, ThroughputProperties.createManualThroughput(throughput)); CosmosAsyncContainer container = getContainer(containerProperties.getId()); if (!client.getTracerProvider().isEnabled()) { - return createContainerIfNotExistsInternal(container.read(), containerProperties, options); + return createContainerIfNotExistsInternal(container.read(), containerProperties, options, null); } return withContext(context -> createContainerIfNotExistsInternal(containerProperties, container, options, @@ -353,7 +353,7 @@ public Mono createContainerIfNotExists( ModelBridgeInternal.setThroughputProperties(options, throughputProperties); CosmosAsyncContainer container = getContainer(containerProperties.getId()); if (!client.getTracerProvider().isEnabled()) { - return createContainerIfNotExistsInternal(container.read(), containerProperties, options); + return createContainerIfNotExistsInternal(container.read(), containerProperties, options, null); } return withContext(context -> createContainerIfNotExistsInternal(containerProperties, container, options, @@ -375,10 +375,12 @@ public Mono createContainerIfNotExists( public Mono createContainerIfNotExists(String id, String partitionKeyPath) { CosmosAsyncContainer container = getContainer(id); if (!client.getTracerProvider().isEnabled()) { - return createContainerIfNotExistsInternal(container.read(), new CosmosContainerProperties(id, partitionKeyPath), null); + return createContainerIfNotExistsInternal(container.read(), new CosmosContainerProperties(id, + partitionKeyPath), null, null); } - return withContext(context -> createContainerIfNotExistsInternal(new CosmosContainerProperties(id, partitionKeyPath), container, null, + return withContext(context -> createContainerIfNotExistsInternal(new CosmosContainerProperties(id, + partitionKeyPath), container, null, context)); } @@ -405,7 +407,7 @@ public Mono createContainerIfNotExists( ModelBridgeInternal.setThroughputProperties(options, throughputProperties); CosmosAsyncContainer container = getContainer(id); if (!client.getTracerProvider().isEnabled()) { - return createContainerIfNotExistsInternal(container.read(), new CosmosContainerProperties(id, partitionKeyPath), options); + return createContainerIfNotExistsInternal(container.read(), new CosmosContainerProperties(id, partitionKeyPath), options, null); } return withContext(context -> createContainerIfNotExistsInternal(new CosmosContainerProperties(id, @@ -435,8 +437,9 @@ Mono createContainerIfNotExists( ModelBridgeInternal.setThroughputProperties(options, ThroughputProperties.createManualThroughput(throughput)); CosmosAsyncContainer container = getContainer(id); if (!client.getTracerProvider().isEnabled()) { - return createContainerIfNotExistsInternal(container.read(), new CosmosContainerProperties(id, partitionKeyPath), - options); + return createContainerIfNotExistsInternal(container.read(), new CosmosContainerProperties(id, + partitionKeyPath), + options, null); } return withContext(context -> createContainerIfNotExistsInternal(new CosmosContainerProperties(id, @@ -795,20 +798,27 @@ private Mono createContainerIfNotExistsInternal( options = new CosmosContainerRequestOptions(); } - Mono responseMono = createContainerIfNotExistsInternal(container.read(options, nestedContext), containerProperties, options); + Mono responseMono = createContainerIfNotExistsInternal(container.read(options, nestedContext), containerProperties, options, nestedContext); return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, getId(), getClient().getServiceEndpoint()); + spanName, + getId(), + getClient().getServiceEndpoint()); } private Mono createContainerIfNotExistsInternal( Mono responseMono, CosmosContainerProperties containerProperties, - CosmosContainerRequestOptions options) { + CosmosContainerRequestOptions options, + 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) { + if(context != null) { + return createContainerInternal(containerProperties, options, context); + } + return createContainer(containerProperties, options); } } @@ -823,7 +833,9 @@ private Mono createContainerInternal( String spanName = "createContainer." + containerProperties.getId(); Mono responseMono = createContainerInternal(containerProperties, options); return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, getId(), getClient().getServiceEndpoint()); + spanName, + getId(), + getClient().getServiceEndpoint()); } private Mono createContainerInternal( @@ -839,7 +851,9 @@ Mono readInternal(CosmosDatabaseRequestOptions options, String spanName = "readDatabase." + this.getId(); Mono responseMono = readInternal(options); return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, getId(), getClient().getServiceEndpoint()); + spanName, + getId(), + getClient().getServiceEndpoint()); } private Mono readInternal(CosmosDatabaseRequestOptions options) { @@ -852,7 +866,9 @@ private Mono deleteInternal(CosmosDatabaseRequestOptions String spanName = "deleteDatabase." + this.getId(); Mono responseMono = deleteInternal(options); return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, getId(), getClient().getServiceEndpoint()); + spanName, + getId(), + getClient().getServiceEndpoint()); } private Mono deleteInternal(CosmosDatabaseRequestOptions options) { @@ -865,7 +881,9 @@ private Mono createUserInternal(CosmosUserProperties userPro String spanName = "createUser." + this.getId(); Mono responseMono = createUserInternal(userProperties); return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, getId(), getClient().getServiceEndpoint()); + spanName, + getId(), + getClient().getServiceEndpoint()); } private Mono createUserInternal(CosmosUserProperties userProperties) { @@ -889,8 +907,11 @@ private Mono replaceThroughputInternal(ThroughputProperties String spanName = "replaceThroughput." + this.getId(); Context nestedContext = context.addData(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.COSMOS_CALL_DEPTH_VAL); Mono responseMono = replaceThroughputInternal(this.readInternal(new CosmosDatabaseRequestOptions(), nestedContext), throughputProperties); - return this.client.getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono - , context, spanName, getId(), getClient().getServiceEndpoint()); + return this.client.getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono, + context, + spanName, + getId(), + getClient().getServiceEndpoint()); } private Mono replaceThroughputInternal(Mono responseMono, ThroughputProperties throughputProperties) { @@ -924,8 +945,11 @@ private Mono readThroughputInternal(Context context){ String spanName = "readThroughput." + this.getId(); Context nestedContext = context.addData(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.COSMOS_CALL_DEPTH_VAL); Mono responseMono = readThroughputInternal(this.readInternal(new CosmosDatabaseRequestOptions(), nestedContext)); - return this.client.getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono - , context, spanName, getId(), getClient().getServiceEndpoint()); + return this.client.getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono, + context, + spanName, + getId(), + getClient().getServiceEndpoint()); } private Mono readThroughputInternal(Mono responseMono) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java index 96b8060e7ada..128937f4f762 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java @@ -139,7 +139,9 @@ private Mono readInternal(CosmosPermissionRequestOptio String spanName = "readPermission." + cosmosUser.getId(); Mono responseMono = readInternal(options); return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, cosmosUser.getDatabase().getId(), cosmosUser.getDatabase().getClient().getServiceEndpoint()); + spanName, + cosmosUser.getDatabase().getId(), + cosmosUser.getDatabase().getClient().getServiceEndpoint()); } private Mono readInternal(CosmosPermissionRequestOptions options) { @@ -158,7 +160,9 @@ private Mono replaceInternal(CosmosPermissionPropertie String spanName = "replacePermission." + cosmosUser.getId(); Mono responseMono = replaceInternal(permissionProperties, options); return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, cosmosUser.getDatabase().getId(), cosmosUser.getDatabase().getClient().getServiceEndpoint()); + spanName, + cosmosUser.getDatabase().getId(), + cosmosUser.getDatabase().getClient().getServiceEndpoint()); } private Mono replaceInternal(CosmosPermissionProperties permissionProperties, @@ -178,7 +182,9 @@ private Mono deleteInternal(CosmosPermissionRequestOpt String spanName = "deletePermission." + cosmosUser.getId(); Mono responseMono = deleteInternal(options); return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, cosmosUser.getDatabase().getId(), cosmosUser.getDatabase().getClient().getServiceEndpoint()); + spanName, + cosmosUser.getDatabase().getId(), + cosmosUser.getDatabase().getClient().getServiceEndpoint()); } private Mono deleteInternal(CosmosPermissionRequestOptions options) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java index 9b0f2ec70c83..7d4e22244948 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java @@ -61,7 +61,7 @@ public Mono createStoredProcedure(CosmosStoredPro * In case of failure the {@link Mono} will error. * * @param properties the cosmos stored procedure properties. - * @param options the stored procedure request options. + * @param options the stored procedure request options. * @return an {@link Mono} containing the single cosmos stored procedure resource response or an error. */ public Mono createStoredProcedure( @@ -364,8 +364,8 @@ CosmosPagedFlux readAllTriggers(CosmosQueryRequestOptio * The {@link CosmosPagedFlux} will contain one or several feed response pages of the obtained triggers. * In case of failure the {@link CosmosPagedFlux} will error. * - * @param query the query. - * @param options the feed options. + * @param query the query. + * @param options the query request options. * @return a {@link CosmosPagedFlux} containing one or several feed response pages of the obtained triggers or an * error. */ @@ -467,7 +467,11 @@ private Mono createStoredProcedureInternal(Stored Context context) { String spanName = "createStoredProcedure." + container.getId(); Mono responseMono = createStoredProcedureInternal(sProc, options); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, + context, + spanName, + database.getId(), + database.getClient().getServiceEndpoint()); } private Mono createStoredProcedureInternal(StoredProcedure sProc, @@ -482,7 +486,11 @@ private Mono createUserDefinedFunctionInterna Context context) { String spanName = "createUserDefinedFunction." + container.getId(); Mono responseMono = createUserDefinedFunctionInternal(udf); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, + context, + spanName, + database.getId(), + database.getClient().getServiceEndpoint()); } private Mono createUserDefinedFunctionInternal( @@ -494,7 +502,11 @@ private Mono createUserDefinedFunctionInterna private Mono createTriggerInternal(CosmosTriggerProperties properties, Context context) { String spanName = "createTrigger." + container.getId(); Mono responseMono = createTriggerInternal(properties); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, + context, + spanName, + database.getId(), + database.getClient().getServiceEndpoint()); } private Mono createTriggerInternal(CosmosTriggerProperties properties) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java index 71650e825e69..e6d4433de989 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java @@ -194,7 +194,11 @@ private Mono readInternal(CosmosStoredProcedureRe Context context) { String spanName = "readStoredProcedure." + cosmosContainer.getId(); Mono responseMono = readInternal(options); - return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); + return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, + context, + spanName, + cosmosContainer.getDatabase().getId(), + cosmosContainer.getDatabase().getClient().getServiceEndpoint()); } private Mono readInternal(CosmosStoredProcedureRequestOptions options) { @@ -210,7 +214,11 @@ private Mono deleteInternal(CosmosStoredProcedure Context context) { String spanName = "deleteStoredProcedure." + cosmosContainer.getId(); Mono responseMono = deleteInternal(options); - return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); + return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, + context, + spanName, + cosmosContainer.getDatabase().getId(), + cosmosContainer.getDatabase().getClient().getServiceEndpoint()); } private Mono deleteInternal(CosmosStoredProcedureRequestOptions options) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java index eef66c14e3ff..666ddad0a646 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java @@ -118,7 +118,11 @@ String getLink() { private Mono readInternal(Context context) { String spanName = "readTrigger." + container.getId(); Mono responseMono = readInternal(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, + context, + spanName, + container.getDatabase().getId(), + container.getDatabase().getClient().getServiceEndpoint()); } private Mono readInternal() { @@ -132,7 +136,11 @@ private Mono readInternal() { private Mono replaceInternal(CosmosTriggerProperties triggerSettings, Context context) { String spanName = "replaceTrigger." + container.getId(); Mono responseMono = replaceInternal(triggerSettings); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, + context, + spanName, + container.getDatabase().getId(), + container.getDatabase().getClient().getServiceEndpoint()); } private Mono replaceInternal(CosmosTriggerProperties triggerSettings) { @@ -147,7 +155,11 @@ private Mono replaceInternal(CosmosTriggerProperties trig private Mono deleteInternal(Context context) { String spanName = "deleteTrigger." + container.getId(); Mono responseMono = deleteInternal(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, + context, + spanName, + container.getDatabase().getId(), + container.getDatabase().getClient().getServiceEndpoint()); } private Mono deleteInternal() { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java index 43442f117d22..f2ea9c442a9a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java @@ -272,7 +272,9 @@ private Mono readInternal(Context context) { String spanName = "readUser." + getId(); Mono responseMono = readInternal(); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, database.getId(), database.getClient().getServiceEndpoint()); + spanName, + database.getId(), + database.getClient().getServiceEndpoint()); } private Mono readInternal() { @@ -285,7 +287,9 @@ private Mono replaceInternal(CosmosUserProperties userSettin String spanName = "replaceUser." + getId(); Mono responseMono = replaceInternal(userSettings); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, database.getId(), database.getClient().getServiceEndpoint()); + spanName, + database.getId(), + database.getClient().getServiceEndpoint()); } private Mono replaceInternal(CosmosUserProperties userSettings) { @@ -298,7 +302,9 @@ private Mono deleteInternal(Context context) { String spanName = "deleteUser." + getId(); Mono responseMono = deleteInternal(); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, database.getId(), database.getClient().getServiceEndpoint()); + spanName, + database.getId(), + database.getClient().getServiceEndpoint()); } private Mono deleteInternal() { @@ -314,7 +320,9 @@ private Mono createPermissionInternal( String spanName = "createPermission." + getId(); Mono responseMono = createPermissionInternal(permission, options); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, database.getId(), database.getClient().getServiceEndpoint()); + spanName, + database.getId(), + database.getClient().getServiceEndpoint()); } private Mono createPermissionInternal( @@ -333,7 +341,9 @@ private Mono upsertPermissionInternal( String spanName = "upsertPermission." + getId(); Mono responseMono = upsertPermissionInternal(permission, options); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, database.getId(), database.getClient().getServiceEndpoint()); + spanName, + database.getId(), + database.getClient().getServiceEndpoint()); } private Mono upsertPermissionInternal( diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java index ff524c1783e7..45b6803fdaac 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java @@ -124,7 +124,11 @@ String getLink() { private Mono readInternal(Context context) { String spanName = "readUDF." + container.getId(); Mono responseMono = readInternal(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, + context, + spanName, + container.getDatabase().getId(), + container.getDatabase().getClient().getServiceEndpoint()); } private Mono readInternal() { @@ -136,7 +140,11 @@ private Mono replaceInternal(CosmosUserDefine Context context) { String spanName = "replaceUDF." + container.getId(); Mono responseMono = replaceInternal(udfSettings); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, + context, + spanName, + container.getDatabase().getId(), + container.getDatabase().getClient().getServiceEndpoint()); } private Mono replaceInternal(CosmosUserDefinedFunctionProperties udfSettings) { @@ -151,7 +159,11 @@ private Mono replaceInternal(CosmosUserDefine private Mono deleteInternal(Context context) { String spanName = "deleteUDF." + container.getId(); Mono responseMono = deleteInternal(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, container.getDatabase().getId(), container.getDatabase().getClient().getServiceEndpoint()); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, + context, + spanName, + container.getDatabase().getId(), + container.getDatabase().getClient().getServiceEndpoint()); } private Mono deleteInternal() { From b9c60a2431add13e597f3c7874f590cc8695945f Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Thu, 11 Jun 2020 19:30:10 -0400 Subject: [PATCH 24/40] resolving comments --- .../com/azure/cosmos/CosmosAsyncContainer.java | 4 ++-- .../com/azure/cosmos/CosmosAsyncDatabase.java | 4 ++-- .../cosmos/implementation/TracerProvider.java | 15 ++++----------- 3 files changed, 8 insertions(+), 15 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index e08efdd4d8c4..81055b0dc932 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -887,7 +887,7 @@ private Mono readThroughputInternal(Context context) { Context nestedContext = context.addData(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.COSMOS_CALL_DEPTH_VAL); Mono responseMono = readThroughputInternal(this.read(new CosmosContainerRequestOptions(), nestedContext)); - return this.getDatabase().getClient().getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono, + return this.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, this.readThroughputSpanName, database.getId(), @@ -924,7 +924,7 @@ private Mono replaceThroughputInternal(ThroughputProperties Mono responseMono = replaceThroughputInternal(this.read(new CosmosContainerRequestOptions(), nestedContext), throughputProperties); - return this.getDatabase().getClient().getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono, + return this.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, this.replaceThroughputSpanName, database.getId(), diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java index 6f09a2f43f61..1da2affd0916 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java @@ -907,7 +907,7 @@ private Mono replaceThroughputInternal(ThroughputProperties String spanName = "replaceThroughput." + this.getId(); Context nestedContext = context.addData(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.COSMOS_CALL_DEPTH_VAL); Mono responseMono = replaceThroughputInternal(this.readInternal(new CosmosDatabaseRequestOptions(), nestedContext), throughputProperties); - return this.client.getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono, + return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, getId(), @@ -945,7 +945,7 @@ private Mono readThroughputInternal(Context context){ String spanName = "readThroughput." + this.getId(); Context nestedContext = context.addData(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.COSMOS_CALL_DEPTH_VAL); Mono responseMono = readThroughputInternal(this.readInternal(new CosmosDatabaseRequestOptions(), nestedContext)); - return this.client.getTracerProvider().traceEnabledNonCosmosResponsePublisher(responseMono, + return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, getId(), diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java index 0c8a6716e432..344cef63c391 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java @@ -17,6 +17,8 @@ import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; +import static com.azure.core.util.tracing.Tracer.AZ_TRACING_NAMESPACE_KEY; + public class TracerProvider { private Tracer tracer; public final static String DB_TYPE_VALUE = "Cosmos"; @@ -30,6 +32,7 @@ public class TracerProvider { public static final String COSMOS_CALL_DEPTH = "cosmosCallDepth"; public static final String COSMOS_CALL_DEPTH_VAL = "nested"; public static final int ERROR_CODE = 0; + public static final String RESOURCE_PROVIDER_NAME = "Microsoft.DocumentDB"; public TracerProvider(Iterable tracers) { Objects.requireNonNull(tracers, "'tracers' cannot be null."); @@ -59,6 +62,7 @@ public Context startSpan(String methodName, String databaseId, String endpoint, tracer.setAttribute(TracerProvider.DB_INSTANCE, databaseId, local); } + tracer.setAttribute(AZ_TRACING_NAMESPACE_KEY, RESOURCE_PROVIDER_NAME, local); tracer.setAttribute(TracerProvider.DB_TYPE, DB_TYPE_VALUE, local); tracer.setAttribute(TracerProvider.DB_URL, endpoint, local); tracer.setAttribute(TracerProvider.DB_STATEMENT, methodName, local); @@ -108,14 +112,6 @@ public > Mono traceEnabledCosmosResponsePublisher (T response) -> response.getStatusCode()); } - public Mono traceEnabledNonCosmosResponsePublisher(Mono resultPublisher, - Context context, - String spanName, - String databaseId, - String endpoint) { - return traceEnabledPublisher(resultPublisher, context, spanName, databaseId, endpoint, (T response) -> HttpConstants.StatusCodes.OK); - } - public Mono> traceEnabledCosmosItemResponsePublisher(Mono> resultPublisher, Context context, String spanName, @@ -155,9 +151,6 @@ private void end(int statusCode, Throwable throwable, Context context) { if (throwable != null) { tracer.setAttribute(TracerProvider.ERROR_MSG, throwable.getMessage(), context); tracer.setAttribute(TracerProvider.ERROR_TYPE, throwable.getClass().getName(), context); - StringWriter errorStack = new StringWriter(); - throwable.printStackTrace(new PrintWriter(errorStack)); - tracer.setAttribute(TracerProvider.ERROR_STACK, errorStack.toString(), context); } tracer.end(statusCode, throwable, context); } From aaabc72ed458360d2e758f05e8744f2ba22d2d8e Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Mon, 15 Jun 2020 12:33:48 -0400 Subject: [PATCH 25/40] moving AZ_TRACING_NAMESPACE_KEY to context from span attribute --- .../java/com/azure/cosmos/implementation/TracerProvider.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java index 344cef63c391..677e87db4e2f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java @@ -57,12 +57,12 @@ public boolean isEnabled() { */ public Context startSpan(String methodName, String databaseId, String endpoint, Context context) { Context local = Objects.requireNonNull(context, "'context' cannot be null."); + local = local.addData(AZ_TRACING_NAMESPACE_KEY, RESOURCE_PROVIDER_NAME); local = tracer.start(methodName, local); // start the span and return the started span if (databaseId != null) { tracer.setAttribute(TracerProvider.DB_INSTANCE, databaseId, local); } - tracer.setAttribute(AZ_TRACING_NAMESPACE_KEY, RESOURCE_PROVIDER_NAME, local); tracer.setAttribute(TracerProvider.DB_TYPE, DB_TYPE_VALUE, local); tracer.setAttribute(TracerProvider.DB_URL, endpoint, local); tracer.setAttribute(TracerProvider.DB_STATEMENT, methodName, local); From 8024999b269086cdc804cf1b2212abde9f21ca6a Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Mon, 15 Jun 2020 14:40:41 -0400 Subject: [PATCH 26/40] updating core opentelementry jar --- sdk/cosmos/azure-cosmos/pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/pom.xml b/sdk/cosmos/azure-cosmos/pom.xml index 588a958e95b6..ae8dee16f2b3 100644 --- a/sdk/cosmos/azure-cosmos/pom.xml +++ b/sdk/cosmos/azure-cosmos/pom.xml @@ -131,7 +131,7 @@ Licensed under the MIT License. com.azure azure-core-tracing-opentelemetry test - 1.0.0-beta.4 + 1.0.0-beta.5 io.netty From 1f4c6ac37696c0ce0b7e15677955ede58afcf04b Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Tue, 7 Jul 2020 10:41:19 -0400 Subject: [PATCH 27/40] Tracer code without conditional check for withContext --- .../com/azure/cosmos/CosmosAsyncClient.java | 38 ++------- .../com/azure/cosmos/CosmosAsyncConflict.java | 8 -- .../azure/cosmos/CosmosAsyncContainer.java | 77 +++++-------------- .../com/azure/cosmos/CosmosAsyncDatabase.java | 75 +++--------------- .../azure/cosmos/CosmosAsyncPermission.java | 12 --- .../com/azure/cosmos/CosmosAsyncScripts.java | 24 ++---- .../cosmos/CosmosAsyncStoredProcedure.java | 16 ---- .../com/azure/cosmos/CosmosAsyncTrigger.java | 12 --- .../com/azure/cosmos/CosmosAsyncUser.java | 24 +----- .../CosmosAsyncUserDefinedFunction.java | 12 --- .../cosmos/implementation/TracerProvider.java | 8 +- .../azure/cosmos/util/CosmosPagedFlux.java | 41 ++++------ .../azure/cosmos/util/UtilBridgeInternal.java | 5 +- .../com/azure/cosmos/CosmosTracerTest.java | 12 +-- .../rx/ReadFeedExceptionHandlingTest.java | 2 +- 15 files changed, 73 insertions(+), 293 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index 473c45d6d25d..7fefa8938ecf 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -202,11 +202,6 @@ boolean isContentResponseOnWriteEnabled() { * an error. */ public Mono 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)); } @@ -222,11 +217,6 @@ public Mono createDatabaseIfNotExists(CosmosDatabaseProp * an error. */ public Mono createDatabaseIfNotExists(String id) { - if(!getTracerProvider().isEnabled()) { - CosmosAsyncDatabase database = getDatabase(id); - return createDatabaseIfNotExistsInternal(database.read(), database, null,null); - } - return withContext(context -> createDatabaseIfNotExistsInternal(getDatabase(id), null, context)); } @@ -244,11 +234,6 @@ public Mono createDatabaseIfNotExists(String id) { * @return the mono. */ public Mono 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)); } @@ -272,10 +257,6 @@ public Mono 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)); } @@ -332,11 +313,6 @@ public Mono 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)); } @@ -406,11 +382,11 @@ CosmosPagedFlux readAllDatabases(CosmosQueryRequestOpt 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()); + .map(response -> + BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosDatabasePropertiesFromV2Results(response.getResults()), + response.getResponseHeaders())); + }); } /** @@ -481,14 +457,14 @@ TracerProvider getTracerProvider(){ private CosmosPagedFlux queryDatabasesInternal(SqlQuerySpec querySpec, CosmosQueryRequestOptions options){ return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "queryDatabases"; + 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()); + }); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncConflict.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncConflict.java index 96f1edaaa7ed..08177c49043c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncConflict.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncConflict.java @@ -67,10 +67,6 @@ public Mono read(CosmosConflictRequestOptions options) { options = new CosmosConflictRequestOptions(); } RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); - if (!this.container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return readInternal(requestOptions); - } - return withContext(context -> readInternal(requestOptions, context)); } @@ -90,10 +86,6 @@ public Mono 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)); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 6c1fb405c039..e8625d310857 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -4,9 +4,9 @@ import com.azure.core.util.Context; import com.azure.cosmos.implementation.Constants; -import com.azure.cosmos.implementation.InternalObjectNode; import com.azure.cosmos.implementation.Document; import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.InternalObjectNode; import com.azure.cosmos.implementation.Offer; import com.azure.cosmos.implementation.Paths; import com.azure.cosmos.implementation.RequestOptions; @@ -122,10 +122,6 @@ public Mono read(CosmosContainerRequestOptions options) } final CosmosContainerRequestOptions requestOptions = options; - if(!database.getClient().getTracerProvider().isEnabled()){ - return readInternal(options); - } - return withContext(context -> read(requestOptions, context)); } @@ -145,10 +141,6 @@ public Mono delete(CosmosContainerRequestOptions option options = new CosmosContainerRequestOptions(); } - if (!database.getClient().getTracerProvider().isEnabled()) { - return deleteInternal(options); - } - final CosmosContainerRequestOptions requestOptions = options; return withContext(context -> deleteInternal(requestOptions, context)); } @@ -203,10 +195,6 @@ public Mono replace( options = new CosmosContainerRequestOptions(); } - if(!database.getClient().getTracerProvider().isEnabled()) { - return replaceInternal(containerProperties, options); - } - final CosmosContainerRequestOptions requestOptions = options; return withContext(context -> replaceInternal(containerProperties, requestOptions, context)); } @@ -266,10 +254,6 @@ public Mono> createItem(T item, CosmosItemRequestOptio options = new CosmosItemRequestOptions(); } - if (!database.getClient().getTracerProvider().isEnabled()) { - return createItemInternal(item, options); - } - final CosmosItemRequestOptions requestOptions = options; return withContext(context -> createItemInternal(item, requestOptions, context)); } @@ -328,10 +312,6 @@ public Mono> upsertItem(T item, CosmosItemRequestOptio options = new CosmosItemRequestOptions(); } - if(!database.getClient().getTracerProvider().isEnabled()) { - return upsertItemInternal(item, options); - } - final CosmosItemRequestOptions requestOptions = options; return withContext(context -> upsertItemInternal(item, requestOptions, context)); } @@ -367,12 +347,13 @@ CosmosPagedFlux readAllItems(Class classType) { */ CosmosPagedFlux readAllItems(CosmosQueryRequestOptions options, Class classType) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), this.readAllItemsSpanName, + pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), + this.readAllItemsSpanName, this.getDatabase().getClient().getServiceEndpoint(), database.getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDatabase().getDocClientWrapper().readDocuments(getLink(), options).map( response -> prepareFeedResponse(response, classType)); - }, this.getDatabase().getClient().getTracerProvider().isEnabled()); + }); } /** @@ -453,10 +434,10 @@ private CosmosPagedFlux queryItemsInternal( this.getDatabase().getClient().getServiceEndpoint(), database.getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, cosmosQueryRequestOptions); return getDatabase().getDocClientWrapper() - .queryDocuments(CosmosAsyncContainer.this.getLink(), sqlQuerySpec, cosmosQueryRequestOptions) - .map(response -> - prepareFeedResponse(response, classType)); - }, this.getDatabase().getClient().getTracerProvider().isEnabled()); + .queryDocuments(CosmosAsyncContainer.this.getLink(), sqlQuerySpec, cosmosQueryRequestOptions) + .map(response -> + prepareFeedResponse(response, classType)); + }); } private FeedResponse prepareFeedResponse(FeedResponse response, Class classType) { @@ -523,10 +504,6 @@ public Mono> readItem( ModelBridgeInternal.setPartitionKey(options, partitionKey); RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); - if(!database.getClient().getTracerProvider().isEnabled()) { - return readItemInternal(itemId, requestOptions, itemType); - } - return withContext(context -> readItemInternal(itemId, requestOptions, itemType, context)); } @@ -569,10 +546,6 @@ public Mono> replaceItem( ModelBridgeInternal.setPartitionKey(options, partitionKey); @SuppressWarnings("unchecked") Class itemType = (Class) item.getClass(); - if(!database.getClient().getTracerProvider().isEnabled()){ - return replaceItemInternal(itemType, itemId, doc, options); - } - final CosmosItemRequestOptions requestOptions = options; return withContext(context -> replaceItemInternal(itemType, itemId, doc, requestOptions,context)); } @@ -610,10 +583,6 @@ public Mono> deleteItem( } ModelBridgeInternal.setPartitionKey(options, partitionKey); RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); - if(!database.getClient().getTracerProvider().isEnabled()) { - return deleteItemInternal(itemId, requestOptions); - } - return withContext(context -> deleteItemInternal(itemId, requestOptions, context)); } @@ -650,14 +619,15 @@ public CosmosAsyncScripts getScripts() { */ public CosmosPagedFlux readAllConflicts(CosmosQueryRequestOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), this.readAllConflictsSpanName, + pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), + this.readAllConflictsSpanName, this.getDatabase().getClient().getServiceEndpoint(), database.getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper().readConflicts(getLink(), options) - .map(response -> BridgeInternal.createFeedResponse( - ModelBridgeInternal.getCosmosConflictPropertiesFromV2Results(response.getResults()), - response.getResponseHeaders())); - }, this.getDatabase().getClient().getTracerProvider().isEnabled()); + .map(response -> BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosConflictPropertiesFromV2Results(response.getResults()), + response.getResponseHeaders())); + }); } /** @@ -681,14 +651,15 @@ public CosmosPagedFlux queryConflicts(String query) { */ public CosmosPagedFlux queryConflicts(String query, CosmosQueryRequestOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), this.queryConflictsSpanName, + pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), + this.queryConflictsSpanName, this.getDatabase().getClient().getServiceEndpoint(), database.getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper().queryConflicts(getLink(), query, options) - .map(response -> BridgeInternal.createFeedResponse( - ModelBridgeInternal.getCosmosConflictPropertiesFromV2Results(response.getResults()), - response.getResponseHeaders())); - }, this.getDatabase().getClient().getTracerProvider().isEnabled()); + .map(response -> BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosConflictPropertiesFromV2Results(response.getResults()), + response.getResponseHeaders())); + }); } /** @@ -708,10 +679,6 @@ public CosmosAsyncConflict getConflict(String id) { * @return the mono containing throughput response */ public Mono replaceThroughput(ThroughputProperties throughputProperties) { - if (!this.database.getClient().getTracerProvider().isEnabled()) { - return replaceThroughputInternal(this.read(), throughputProperties); - } - return withContext(context -> replaceThroughputInternal(throughputProperties, context)); } @@ -721,10 +688,6 @@ public Mono replaceThroughput(ThroughputProperties throughpu * @return the mono containing throughput response. */ public Mono readThroughput() { - if (!this.database.getClient().getTracerProvider().isEnabled()) { - return readThroughputInternal(this.read()); - } - return withContext(context -> readThroughputInternal(context)); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java index 1da2affd0916..9a341e92e774 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java @@ -86,9 +86,6 @@ public Mono read(CosmosDatabaseRequestOptions options) { options = new CosmosDatabaseRequestOptions(); } - if (!client.getTracerProvider().isEnabled()) { - return readInternal(options); - } final CosmosDatabaseRequestOptions requestOptions = options; return withContext(context -> readInternal(requestOptions, context)); } @@ -120,9 +117,6 @@ public Mono delete(CosmosDatabaseRequestOptions options) if (options == null) { options = new CosmosDatabaseRequestOptions(); } - if (!client.getTracerProvider().isEnabled()) { - return deleteInternal(options); - } final CosmosDatabaseRequestOptions requestOptions = options; return withContext(context -> deleteInternal(requestOptions, context)); @@ -209,10 +203,6 @@ public Mono createContainer( options = new CosmosContainerRequestOptions(); } - if (!client.getTracerProvider().isEnabled()) { - return createContainerInternal(containerProperties, options); - } - final CosmosContainerRequestOptions requestOptions = options; return withContext(context -> createContainerInternal(containerProperties, requestOptions, context)); } @@ -292,10 +282,6 @@ public Mono createContainer(String id, String partition public Mono createContainerIfNotExists( CosmosContainerProperties containerProperties) { CosmosAsyncContainer container = getContainer(containerProperties.getId()); - if (!client.getTracerProvider().isEnabled()) { - return createContainerIfNotExistsInternal(container.read(), containerProperties, null, null); - } - return withContext(context -> createContainerIfNotExistsInternal(containerProperties, container, null, context)); } @@ -322,10 +308,6 @@ Mono createContainerIfNotExists( CosmosContainerRequestOptions options = new CosmosContainerRequestOptions(); ModelBridgeInternal.setThroughputProperties(options, ThroughputProperties.createManualThroughput(throughput)); CosmosAsyncContainer container = getContainer(containerProperties.getId()); - if (!client.getTracerProvider().isEnabled()) { - return createContainerIfNotExistsInternal(container.read(), containerProperties, options, null); - } - return withContext(context -> createContainerIfNotExistsInternal(containerProperties, container, options, context)); } @@ -352,10 +334,6 @@ public Mono createContainerIfNotExists( CosmosContainerRequestOptions options = new CosmosContainerRequestOptions(); ModelBridgeInternal.setThroughputProperties(options, throughputProperties); CosmosAsyncContainer container = getContainer(containerProperties.getId()); - if (!client.getTracerProvider().isEnabled()) { - return createContainerIfNotExistsInternal(container.read(), containerProperties, options, null); - } - return withContext(context -> createContainerIfNotExistsInternal(containerProperties, container, options, context)); } @@ -374,11 +352,6 @@ public Mono createContainerIfNotExists( */ public Mono createContainerIfNotExists(String id, String partitionKeyPath) { CosmosAsyncContainer container = getContainer(id); - if (!client.getTracerProvider().isEnabled()) { - return createContainerIfNotExistsInternal(container.read(), new CosmosContainerProperties(id, - partitionKeyPath), null, null); - } - return withContext(context -> createContainerIfNotExistsInternal(new CosmosContainerProperties(id, partitionKeyPath), container, null, context)); @@ -406,10 +379,6 @@ public Mono createContainerIfNotExists( CosmosContainerRequestOptions options = new CosmosContainerRequestOptions(); ModelBridgeInternal.setThroughputProperties(options, throughputProperties); CosmosAsyncContainer container = getContainer(id); - if (!client.getTracerProvider().isEnabled()) { - return createContainerIfNotExistsInternal(container.read(), new CosmosContainerProperties(id, partitionKeyPath), options, null); - } - return withContext(context -> createContainerIfNotExistsInternal(new CosmosContainerProperties(id, partitionKeyPath), container, options, context)); } @@ -436,12 +405,6 @@ Mono createContainerIfNotExists( CosmosContainerRequestOptions options = new CosmosContainerRequestOptions(); ModelBridgeInternal.setThroughputProperties(options, ThroughputProperties.createManualThroughput(throughput)); CosmosAsyncContainer container = getContainer(id); - if (!client.getTracerProvider().isEnabled()) { - return createContainerIfNotExistsInternal(container.read(), new CosmosContainerProperties(id, - partitionKeyPath), - options, null); - } - return withContext(context -> createContainerIfNotExistsInternal(new CosmosContainerProperties(id, partitionKeyPath), container, options, context)); } @@ -464,10 +427,10 @@ public CosmosPagedFlux readAllContainers(CosmosQueryR this.getClient().getServiceEndpoint(), getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().readCollections(getLink(), options) - .map(response -> BridgeInternal.createFeedResponse( - ModelBridgeInternal.getCosmosContainerPropertiesFromV2Results(response.getResults()), - response.getResponseHeaders())); - }, this.getClient().getTracerProvider().isEnabled()); + .map(response -> BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosContainerPropertiesFromV2Results(response.getResults()), + response.getResponseHeaders())); + }); } /** @@ -568,9 +531,6 @@ public CosmosAsyncContainer getContainer(String id) { * created cosmos user or an error. */ public Mono createUser(CosmosUserProperties userProperties) { - if (!client.getTracerProvider().isEnabled()) { - return createUserInternal(userProperties); - } return withContext(context -> createUserInternal(userProperties, context)); } @@ -586,10 +546,6 @@ public Mono createUser(CosmosUserProperties userProperties) * upserted user or an error. */ public Mono upsertUser(CosmosUserProperties userProperties) { - if (!client.getTracerProvider().isEnabled()) { - return upsertUserInternal(userProperties); - } - return withContext(context -> upsertUserInternal(userProperties, context)); } @@ -625,10 +581,10 @@ CosmosPagedFlux readAllUsers(CosmosQueryRequestOptions opt this.getClient().getServiceEndpoint(), getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().readUsers(getLink(), options) - .map(response -> BridgeInternal.createFeedResponse( - ModelBridgeInternal.getCosmosUserPropertiesFromV2Results(response.getResults()), response - .getResponseHeaders())); - }, this.getClient().getTracerProvider().isEnabled()); + .map(response -> BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosUserPropertiesFromV2Results(response.getResults()), response + .getResponseHeaders())); + }); } /** @@ -711,10 +667,6 @@ public CosmosAsyncUser getUser(String id) { * @return the mono. */ public Mono replaceThroughput(ThroughputProperties throughputProperties) { - if(!this.client.getTracerProvider().isEnabled()) { - return replaceThroughputInternal(this.read(), throughputProperties); - } - return withContext(context -> replaceThroughputInternal(throughputProperties, context)); } @@ -724,10 +676,6 @@ public Mono replaceThroughput(ThroughputProperties throughpu * @return the mono containing throughput response. */ public Mono readThroughput() { - if(!this.client.getTracerProvider().isEnabled()) { - return readThroughputInternal(this.read()); - } - return withContext(context -> readThroughputInternal(context)); } @@ -771,7 +719,7 @@ private CosmosPagedFlux queryContainersInternal(SqlQu .map(response -> BridgeInternal.createFeedResponse( ModelBridgeInternal.getCosmosContainerPropertiesFromV2Results(response.getResults()), response.getResponseHeaders())); - }, this.getClient().getTracerProvider().isEnabled()); + }); } private CosmosPagedFlux queryUsersInternal(SqlQuerySpec querySpec, CosmosQueryRequestOptions options) { @@ -782,9 +730,10 @@ private CosmosPagedFlux queryUsersInternal(SqlQuerySpec qu setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().queryUsers(getLink(), querySpec, options) .map(response -> BridgeInternal.createFeedResponseWithQueryMetrics( - ModelBridgeInternal.getCosmosUserPropertiesFromV2Results(response.getResults()), response.getResponseHeaders(), + ModelBridgeInternal.getCosmosUserPropertiesFromV2Results(response.getResults()), + response.getResponseHeaders(), ModelBridgeInternal.queryMetrics(response))); - }, this.getClient().getTracerProvider().isEnabled()); + }); } private Mono createContainerIfNotExistsInternal( diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java index 128937f4f762..a1f8c868fc88 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java @@ -60,10 +60,6 @@ public Mono read(CosmosPermissionRequestOptions option options = new CosmosPermissionRequestOptions(); } - if (!cosmosUser.getDatabase().getClient().getTracerProvider().isEnabled()) { - return readInternal(options); - } - final CosmosPermissionRequestOptions requestOptions = options; return withContext(context -> readInternal(requestOptions, context)); } @@ -85,10 +81,6 @@ public Mono replace(CosmosPermissionProperties permiss options = new CosmosPermissionRequestOptions(); } - if (!cosmosUser.getDatabase().getClient().getTracerProvider().isEnabled()) { - return replaceInternal(permissionProperties, options); - } - final CosmosPermissionRequestOptions requestOptions = options; return withContext(context -> replaceInternal(permissionProperties, requestOptions, context)); } @@ -108,10 +100,6 @@ public Mono delete(CosmosPermissionRequestOptions opti options = new CosmosPermissionRequestOptions(); } - if (!cosmosUser.getDatabase().getClient().getTracerProvider().isEnabled()) { - return deleteInternal(options); - } - final CosmosPermissionRequestOptions requestOptions = options; return withContext(context -> deleteInternal(requestOptions, context)); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java index 7d4e22244948..7079baa3e225 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java @@ -73,10 +73,6 @@ public Mono createStoredProcedure( StoredProcedure sProc = new StoredProcedure(); sProc.setId(properties.getId()); sProc.setBody(properties.getBody()); - if (!container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return createStoredProcedureInternal(sProc, options); - } - final CosmosStoredProcedureRequestOptions requestOptions = options; return withContext(context -> createStoredProcedureInternal(sProc, requestOptions, context)); } @@ -123,7 +119,7 @@ CosmosPagedFlux readAllStoredProcedures(CosmosQ .map(response -> BridgeInternal.createFeedResponse( ModelBridgeInternal.getCosmosStoredProcedurePropertiesFromV2Results(response.getResults()), response.getResponseHeaders())); - }, this.container.getDatabase().getClient().getTracerProvider().isEnabled()); + }); } /** @@ -192,10 +188,6 @@ public Mono createUserDefinedFunction( UserDefinedFunction udf = new UserDefinedFunction(); udf.setId(properties.getId()); udf.setBody(properties.getBody()); - if (!container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return createUserDefinedFunctionInternal(udf); - } - return withContext(context -> createUserDefinedFunctionInternal(udf, context)); } @@ -239,7 +231,7 @@ CosmosPagedFlux readAllUserDefinedFunctions .map(response -> BridgeInternal.createFeedResponse( ModelBridgeInternal.getCosmosUserDefinedFunctionPropertiesFromV2Results(response.getResults()), response.getResponseHeaders())); - }, this.container.getDatabase().getClient().getTracerProvider().isEnabled()); + }); } /** @@ -305,10 +297,6 @@ public CosmosAsyncUserDefinedFunction getUserDefinedFunction(String id) { * @return an {@link Mono} containing the single resource response with the created trigger or an error. */ public Mono createTrigger(CosmosTriggerProperties properties) { - if (!container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return createTriggerInternal(properties); - } - return withContext(context -> createTriggerInternal(properties, context)); } @@ -354,7 +342,7 @@ CosmosPagedFlux readAllTriggers(CosmosQueryRequestOptio .map(response -> BridgeInternal.createFeedResponse( ModelBridgeInternal.getCosmosTriggerPropertiesFromV2Results(response.getResults()), response.getResponseHeaders())); - }, this.container.getDatabase().getClient().getTracerProvider().isEnabled()); + }); } /** @@ -416,7 +404,7 @@ private CosmosPagedFlux queryStoredProceduresIn .map(response -> BridgeInternal.createFeedResponse( ModelBridgeInternal.getCosmosStoredProcedurePropertiesFromV2Results(response.getResults()), response.getResponseHeaders())); - }, this.container.getDatabase().getClient().getTracerProvider().isEnabled()); + }); } private CosmosPagedFlux queryUserDefinedFunctionsInternal( @@ -434,7 +422,7 @@ private CosmosPagedFlux queryUserDefinedFun .map(response -> BridgeInternal.createFeedResponse( ModelBridgeInternal.getCosmosUserDefinedFunctionPropertiesFromV2Results(response.getResults()), response.getResponseHeaders())); - }, this.container.getDatabase().getClient().getTracerProvider().isEnabled()); + }); } private CosmosPagedFlux queryTriggersInternal( @@ -459,7 +447,7 @@ private CosmosPagedFlux queryTriggersInternal( .map(response -> BridgeInternal.createFeedResponse( ModelBridgeInternal.getCosmosTriggerPropertiesFromV2Results(response.getResults()), response.getResponseHeaders())); - }, this.container.getDatabase().getClient().getTracerProvider().isEnabled()); + }); } private Mono createStoredProcedureInternal(StoredProcedure sProc, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java index e6d4433de989..a79396f27195 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java @@ -74,10 +74,6 @@ public Mono read() { * @return an {@link Mono} containing the single resource response with the read stored procedure or an error. */ public Mono read(CosmosStoredProcedureRequestOptions options) { - if (!cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled()) { - return readInternal(options); - } - return withContext(context -> readInternal(options, context)); } @@ -107,10 +103,6 @@ public Mono delete() { * @return an {@link Mono} containing the single resource response for the deleted stored procedure or an error. */ public Mono delete(CosmosStoredProcedureRequestOptions options) { - if (!cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled()) { - return deleteInternal(options); - } - return withContext(context -> deleteInternal(options, context)); } @@ -128,10 +120,6 @@ public Mono delete(CosmosStoredProcedureRequestOp */ public Mono execute(List procedureParams, CosmosStoredProcedureRequestOptions options) { - if (!cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled()) { - return executeInternal(procedureParams, options); - } - return withContext(context -> executeInternal(procedureParams, options, context)); } @@ -164,10 +152,6 @@ public Mono replace(CosmosStoredProcedureProperti */ public Mono replace(CosmosStoredProcedureProperties storedProcedureProperties, CosmosStoredProcedureRequestOptions options) { - if (!cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled()) { - return replaceInternal(storedProcedureProperties, options); - } - return withContext(context -> replaceInternal(storedProcedureProperties, options, context)); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java index 666ddad0a646..e3d627b49542 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java @@ -55,10 +55,6 @@ CosmosAsyncTrigger setId(String id) { * @return an {@link Mono} containing the single resource response for the read cosmos trigger or an error. */ public Mono read() { - if (!container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return readInternal(); - } - return withContext(context -> readInternal(context)); } @@ -73,10 +69,6 @@ public Mono read() { * @return an {@link Mono} containing the single resource response with the replaced cosmos trigger or an error. */ public Mono replace(CosmosTriggerProperties triggerProperties) { - if (!container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return replaceInternal(triggerProperties); - } - return withContext(context -> replaceInternal(triggerProperties, context)); } @@ -90,10 +82,6 @@ public Mono replace(CosmosTriggerProperties triggerProper * @return an {@link Mono} containing the single resource response for the deleted cosmos trigger or an error. */ public Mono delete() { - if (!container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return deleteInternal(); - } - return withContext(context -> deleteInternal(context)); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java index f2ea9c442a9a..52fc7988a3f2 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java @@ -58,10 +58,6 @@ CosmosAsyncUser setId(String id) { * @return a {@link Mono} containing the single resource response with the read user or an error. */ public Mono read() { - if (!database.getClient().getTracerProvider().isEnabled()) { - return readInternal(); - } - return withContext(context -> readInternal(context)); } @@ -72,10 +68,6 @@ public Mono read() { * @return a {@link Mono} containing the single resource response with the replaced user or an error. */ public Mono replace(CosmosUserProperties userProperties) { - if (!database.getClient().getTracerProvider().isEnabled()) { - return replaceInternal(userProperties); - } - return withContext(context -> replaceInternal(userProperties, context)); } @@ -85,10 +77,6 @@ public Mono replace(CosmosUserProperties userProperties) { * @return a {@link Mono} containing the single resource response with the deleted user or an error. */ public Mono delete() { - if (!database.getClient().getTracerProvider().isEnabled()) { - return deleteInternal(); - } - return withContext(context -> deleteInternal(context)); } @@ -111,10 +99,6 @@ public Mono createPermission( } Permission permission = ModelBridgeInternal.getPermission(permissionProperties, database.getId()); - if (!database.getClient().getTracerProvider().isEnabled()) { - return createPermissionInternal(permission, options); - } - final CosmosPermissionRequestOptions requesOptions = options; return withContext(context -> createPermissionInternal(permission, requesOptions, context)); } @@ -138,10 +122,6 @@ public Mono upsertPermission( options = new CosmosPermissionRequestOptions(); } - if (!database.getClient().getTracerProvider().isEnabled()) { - return upsertPermissionInternal(permission, options); - } - final CosmosPermissionRequestOptions requestOptions = options; return withContext(context -> upsertPermissionInternal(permission, requestOptions, context)); } @@ -185,7 +165,7 @@ CosmosPagedFlux readAllPermissions(CosmosQueryReques .map(response -> BridgeInternal.createFeedResponse( ModelBridgeInternal.getCosmosPermissionPropertiesFromResults(response.getResults()), response.getResponseHeaders())); - }, this.getDatabase().getClient().getTracerProvider().isEnabled()); + }); } /** @@ -228,7 +208,7 @@ public CosmosPagedFlux queryPermissions(String query .map(response -> BridgeInternal.createFeedResponse( ModelBridgeInternal.getCosmosPermissionPropertiesFromResults(response.getResults()), response.getResponseHeaders())); - }, this.getDatabase().getClient().getTracerProvider().isEnabled()); + }); } /** diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java index 45b6803fdaac..f0a35931713a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java @@ -57,10 +57,6 @@ CosmosAsyncUserDefinedFunction setId(String id) { * @return an {@link Mono} containing the single resource response for the read user defined function or an error. */ public Mono read() { - if(!container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return readInternal(); - } - return withContext(context -> readInternal(context)); } @@ -77,10 +73,6 @@ public Mono read() { * or an error. */ public Mono replace(CosmosUserDefinedFunctionProperties udfSettings) { - if(!container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return replaceInternal(udfSettings); - } - return withContext(context -> replaceInternal(udfSettings, context)); } @@ -96,10 +88,6 @@ public Mono replace(CosmosUserDefinedFunction * an error. */ public Mono delete() { - if(!container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return deleteInternal(); - } - return withContext(context -> deleteInternal(context)); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java index 677e87db4e2f..665fba9e184a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java @@ -10,8 +10,6 @@ import reactor.core.publisher.Mono; import reactor.core.publisher.Signal; -import java.io.PrintWriter; -import java.io.StringWriter; import java.util.Objects; import java.util.Optional; import java.util.concurrent.atomic.AtomicReference; @@ -132,16 +130,16 @@ public Mono traceEnabledPublisher(Mono resultPublisher, final boolean isNestedCall = callDepth.isPresent(); return resultPublisher .doOnSubscribe(ignoredValue -> { - if (!isNestedCall) { + if (isEnabled() && !isNestedCall) { parentContext.set(this.startSpan(spanName, databaseId, endpoint, context)); } }).doOnSuccess(response -> { - if (!isNestedCall) { + if (isEnabled() && !isNestedCall) { this.endSpan(parentContext.get(), Signal.complete(), statusCodeFunc.apply(response)); } }).doOnError(throwable -> { - if (!isNestedCall) { + if (isEnabled() && !isNestedCall) { this.endSpan(parentContext.get(), Signal.error(throwable), ERROR_CODE); } }); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java index 77167bc5970d..91b9c4af81d2 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java @@ -35,21 +35,14 @@ public final class CosmosPagedFlux extends ContinuablePagedFlux> { private final Function>> optionsFluxFunction; - private final boolean isTracerEnabled; - CosmosPagedFlux(Function>> optionsFluxFunction, - boolean isTracerEnable) { + CosmosPagedFlux(Function>> optionsFluxFunction) { this.optionsFluxFunction = optionsFluxFunction; - this.isTracerEnabled = isTracerEnable; } @Override public Flux> byPage() { CosmosPagedFluxOptions cosmosPagedFluxOptions = new CosmosPagedFluxOptions(); - if (!this.isTracerEnabled) { - return this.optionsFluxFunction.apply(cosmosPagedFluxOptions); - } - return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)); } @@ -57,10 +50,6 @@ public Flux> byPage() { public Flux> byPage(String continuationToken) { CosmosPagedFluxOptions cosmosPagedFluxOptions = new CosmosPagedFluxOptions(); cosmosPagedFluxOptions.setRequestContinuation(continuationToken); - if (!this.isTracerEnabled) { - return this.optionsFluxFunction.apply(cosmosPagedFluxOptions); - } - return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)); } @@ -68,10 +57,6 @@ public Flux> byPage(String continuationToken) { public Flux> byPage(int preferredPageSize) { CosmosPagedFluxOptions cosmosPagedFluxOptions = new CosmosPagedFluxOptions(); cosmosPagedFluxOptions.setMaxItemCount(preferredPageSize); - if (!this.isTracerEnabled) { - return this.optionsFluxFunction.apply(cosmosPagedFluxOptions); - } - return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)); } @@ -80,10 +65,6 @@ public Flux> byPage(String continuationToken, int preferredPageS CosmosPagedFluxOptions cosmosPagedFluxOptions = new CosmosPagedFluxOptions(); cosmosPagedFluxOptions.setRequestContinuation(continuationToken); cosmosPagedFluxOptions.setMaxItemCount(preferredPageSize); - if (!this.isTracerEnabled) { - return this.optionsFluxFunction.apply(cosmosPagedFluxOptions); - } - return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)); } @@ -109,15 +90,21 @@ public void subscribe(CoreSubscriber coreSubscriber) { private Flux> byPage(CosmosPagedFluxOptions pagedFluxOptions, Context context) { final AtomicReference parentContext = new AtomicReference<>(Context.NONE); return this.optionsFluxFunction.apply(pagedFluxOptions).doOnSubscribe(ignoredValue -> { - parentContext.set(pagedFluxOptions.getTracerProvider().startSpan(pagedFluxOptions.getTracerSpanName(), - pagedFluxOptions.getDatabaseId(), pagedFluxOptions.getServiceEndpoint(), - context)); + if (pagedFluxOptions.getTracerProvider().isEnabled()) { + parentContext.set(pagedFluxOptions.getTracerProvider().startSpan(pagedFluxOptions.getTracerSpanName(), + pagedFluxOptions.getDatabaseId(), pagedFluxOptions.getServiceEndpoint(), + context)); + } }).doOnComplete(() -> { - pagedFluxOptions.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), - HttpConstants.StatusCodes.OK); + if (pagedFluxOptions.getTracerProvider().isEnabled()) { + pagedFluxOptions.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), + HttpConstants.StatusCodes.OK); + } }).doOnError(throwable -> { - pagedFluxOptions.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), - TracerProvider.ERROR_CODE); + if (pagedFluxOptions.getTracerProvider().isEnabled()) { + pagedFluxOptions.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), + TracerProvider.ERROR_CODE); + } }); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/UtilBridgeInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/UtilBridgeInternal.java index 141797b27ed3..1282cd7a56de 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/UtilBridgeInternal.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/UtilBridgeInternal.java @@ -24,11 +24,10 @@ public final class UtilBridgeInternal { private UtilBridgeInternal() {} @Warning(value = INTERNAL_USE_ONLY_WARNING) - public static CosmosPagedFlux createCosmosPagedFlux(Function>> pagedFluxOptionsFluxFunction, boolean isTracerEnabled) { - return new CosmosPagedFlux<>(pagedFluxOptionsFluxFunction, isTracerEnabled); + public static CosmosPagedFlux createCosmosPagedFlux(Function>> pagedFluxOptionsFluxFunction) { + return new CosmosPagedFlux<>(pagedFluxOptionsFluxFunction); } - @Warning(value = INTERNAL_USE_ONLY_WARNING) public static CosmosPagedIterable createCosmosPagedIterable(ContinuablePagedFlux> pagedFlux) { return new CosmosPagedIterable<>(pagedFlux); diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java index d7ec6300e224..2307a26c781a 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java @@ -4,7 +4,7 @@ import com.azure.core.util.Context; import com.azure.core.util.tracing.Tracer; -import com.azure.cosmos.implementation.CosmosItemProperties; +import com.azure.cosmos.implementation.InternalObjectNode; import com.azure.cosmos.implementation.LifeCycleUtils; import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.TracerProvider; @@ -109,19 +109,19 @@ public void cosmosAsyncContainer() { Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); - CosmosItemProperties properties = new CosmosItemProperties(); - properties.setId(ITEM_ID); - cosmosAsyncContainer.createItem(properties).block(); + InternalObjectNode item = new InternalObjectNode(); + item.setId(ITEM_ID); + cosmosAsyncContainer.createItem(item).block(); Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); - cosmosAsyncContainer.upsertItem(properties, + cosmosAsyncContainer.upsertItem(item, new CosmosItemRequestOptions()).block(); Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.readItem(ITEM_ID, PartitionKey.NONE, - CosmosItemProperties.class).block(); + InternalObjectNode.class).block(); Mockito.verify(tracer, Mockito.times(5)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/ReadFeedExceptionHandlingTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/ReadFeedExceptionHandlingTest.java index a812c50cecd7..1183e9ae6771 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/ReadFeedExceptionHandlingTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/ReadFeedExceptionHandlingTest.java @@ -52,7 +52,7 @@ public void readFeedException() throws Exception { Mockito.when(mockedClientWrapper.readAllDatabases()).thenReturn(UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { pagedFluxOptions.setTracerInformation(new TracerProvider(ServiceLoader.load(Tracer.class)), "testSpan", "testEndpoint,", "testDb"); return response; - }, false)); + })); TestSubscriber> subscriber = new TestSubscriber<>(); mockedClientWrapper.readAllDatabases().byPage().subscribe(subscriber); assertThat(subscriber.valueCount()).isEqualTo(2); From 3548215370fea7ff0fb9b4e6a79c0130f629cb37 Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Tue, 7 Jul 2020 11:03:37 -0400 Subject: [PATCH 28/40] formating chage --- .../com/azure/cosmos/rx/TestSuiteBase.java | 30 +++++++++---------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java index 698df6b50b8b..ec318e228294 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java @@ -1068,23 +1068,23 @@ static protected CosmosClientBuilder createGatewayHouseKeepingDocumentClient(boo options.setMaxRetryWaitTime(Duration.ofSeconds(SUITE_SETUP_TIMEOUT)); GatewayConnectionConfig gatewayConnectionConfig = new GatewayConnectionConfig(); return new CosmosClientBuilder().endpoint(TestConfigurations.HOST) - .credential(credential) - .gatewayMode(gatewayConnectionConfig) - .throttlingRetryOptions(options) - .contentResponseOnWriteEnabled(contentResponseOnWriteEnabled) - .consistencyLevel(ConsistencyLevel.SESSION); + .credential(credential) + .gatewayMode(gatewayConnectionConfig) + .throttlingRetryOptions(options) + .contentResponseOnWriteEnabled(contentResponseOnWriteEnabled) + .consistencyLevel(ConsistencyLevel.SESSION); } static protected CosmosClientBuilder createGatewayRxDocumentClient(ConsistencyLevel consistencyLevel, boolean multiMasterEnabled, List preferredRegions, boolean contentResponseOnWriteEnabled) { GatewayConnectionConfig gatewayConnectionConfig = new GatewayConnectionConfig(); return new CosmosClientBuilder().endpoint(TestConfigurations.HOST) - .credential(credential) - .gatewayMode(gatewayConnectionConfig) - .multipleWriteRegionsEnabled(multiMasterEnabled) - .preferredRegions(preferredRegions) - .contentResponseOnWriteEnabled(contentResponseOnWriteEnabled) - .consistencyLevel(consistencyLevel); + .credential(credential) + .gatewayMode(gatewayConnectionConfig) + .multipleWriteRegionsEnabled(multiMasterEnabled) + .preferredRegions(preferredRegions) + .contentResponseOnWriteEnabled(contentResponseOnWriteEnabled) + .consistencyLevel(consistencyLevel); } static protected CosmosClientBuilder createGatewayRxDocumentClient() { @@ -1097,10 +1097,10 @@ static protected CosmosClientBuilder createDirectRxDocumentClient(ConsistencyLev List preferredRegions, boolean contentResponseOnWriteEnabled) { CosmosClientBuilder builder = new CosmosClientBuilder().endpoint(TestConfigurations.HOST) - .credential(credential) - .directMode(DirectConnectionConfig.getDefaultConfig()) - .contentResponseOnWriteEnabled(contentResponseOnWriteEnabled) - .consistencyLevel(consistencyLevel); + .credential(credential) + .directMode(DirectConnectionConfig.getDefaultConfig()) + .contentResponseOnWriteEnabled(contentResponseOnWriteEnabled) + .consistencyLevel(consistencyLevel); if (preferredRegions != null) { builder.preferredRegions(preferredRegions); } From eb398d9091cf47354129e61a5d90040fadd6f0d0 Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Tue, 7 Jul 2020 11:07:27 -0400 Subject: [PATCH 29/40] formating chage --- .../com/azure/cosmos/rx/TestSuiteBase.java | 38 +++++++++---------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java index ec318e228294..a5e0e0682e93 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java @@ -484,7 +484,7 @@ public Flux> bulkInsert(CosmosAsyncContainer cosmosCon return Flux.merge(Flux.fromIterable(result), concurrencyLevel); } public List bulkInsertBlocking(CosmosAsyncContainer cosmosContainer, - List documentDefinitionList) { + List documentDefinitionList) { return bulkInsert(cosmosContainer, documentDefinitionList, DEFAULT_BULK_INSERT_CONCURRENCY_LEVEL) .publishOn(Schedulers.parallel()) .map(itemResponse -> itemResponse.getItem()) @@ -815,12 +815,12 @@ public void validateItemFailure( } public void validateQuerySuccess(Flux> flowable, - FeedResponseListValidator validator) { + FeedResponseListValidator validator) { validateQuerySuccess(flowable, validator, subscriberValidationTimeout); } public static void validateQuerySuccess(Flux> flowable, - FeedResponseListValidator validator, long timeout) { + FeedResponseListValidator validator, long timeout) { TestSubscriber> testSubscriber = new TestSubscriber<>(); @@ -836,7 +836,7 @@ public void validateQueryFailure(Flux> flowable, FailureVali } public static void validateQueryFailure(Flux> flowable, - FailureValidator validator, long timeout) { + FailureValidator validator, long timeout) { TestSubscriber> testSubscriber = new TestSubscriber<>(); @@ -1068,23 +1068,23 @@ static protected CosmosClientBuilder createGatewayHouseKeepingDocumentClient(boo options.setMaxRetryWaitTime(Duration.ofSeconds(SUITE_SETUP_TIMEOUT)); GatewayConnectionConfig gatewayConnectionConfig = new GatewayConnectionConfig(); return new CosmosClientBuilder().endpoint(TestConfigurations.HOST) - .credential(credential) - .gatewayMode(gatewayConnectionConfig) - .throttlingRetryOptions(options) - .contentResponseOnWriteEnabled(contentResponseOnWriteEnabled) - .consistencyLevel(ConsistencyLevel.SESSION); + .credential(credential) + .gatewayMode(gatewayConnectionConfig) + .throttlingRetryOptions(options) + .contentResponseOnWriteEnabled(contentResponseOnWriteEnabled) + .consistencyLevel(ConsistencyLevel.SESSION); } static protected CosmosClientBuilder createGatewayRxDocumentClient(ConsistencyLevel consistencyLevel, boolean multiMasterEnabled, List preferredRegions, boolean contentResponseOnWriteEnabled) { GatewayConnectionConfig gatewayConnectionConfig = new GatewayConnectionConfig(); return new CosmosClientBuilder().endpoint(TestConfigurations.HOST) - .credential(credential) - .gatewayMode(gatewayConnectionConfig) - .multipleWriteRegionsEnabled(multiMasterEnabled) - .preferredRegions(preferredRegions) - .contentResponseOnWriteEnabled(contentResponseOnWriteEnabled) - .consistencyLevel(consistencyLevel); + .credential(credential) + .gatewayMode(gatewayConnectionConfig) + .multipleWriteRegionsEnabled(multiMasterEnabled) + .preferredRegions(preferredRegions) + .contentResponseOnWriteEnabled(contentResponseOnWriteEnabled) + .consistencyLevel(consistencyLevel); } static protected CosmosClientBuilder createGatewayRxDocumentClient() { @@ -1097,10 +1097,10 @@ static protected CosmosClientBuilder createDirectRxDocumentClient(ConsistencyLev List preferredRegions, boolean contentResponseOnWriteEnabled) { CosmosClientBuilder builder = new CosmosClientBuilder().endpoint(TestConfigurations.HOST) - .credential(credential) - .directMode(DirectConnectionConfig.getDefaultConfig()) - .contentResponseOnWriteEnabled(contentResponseOnWriteEnabled) - .consistencyLevel(consistencyLevel); + .credential(credential) + .directMode(DirectConnectionConfig.getDefaultConfig()) + .contentResponseOnWriteEnabled(contentResponseOnWriteEnabled) + .consistencyLevel(consistencyLevel); if (preferredRegions != null) { builder.preferredRegions(preferredRegions); } From 9c1d31f8b46f43c312fbb3dfcecc3837369d5e08 Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Thu, 9 Jul 2020 15:55:35 -0400 Subject: [PATCH 30/40] resolving comments --- .../com/azure/cosmos/CosmosAsyncClient.java | 14 ++++----- .../azure/cosmos/CosmosAsyncContainer.java | 30 +++++-------------- .../com/azure/cosmos/CosmosAsyncDatabase.java | 17 ++--------- .../com/azure/cosmos/CosmosAsyncUser.java | 14 ++------- 4 files changed, 19 insertions(+), 56 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index 7fefa8938ecf..341b3b92a8c5 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -49,6 +49,8 @@ public final class CosmosAsyncClient implements Closeable { private final AsyncDocumentClient asyncDocumentClient; private final String serviceEndpoint; private final String keyOrResourceToken; + private final String queryDatabasesSpanName = "queryDatabases"; + private final String readAllDatabasesSpanName = "readAllDatabases"; private final ConnectionPolicy connectionPolicy; private final ConsistencyLevel desiredConsistencyLevel; private final List permissions; @@ -252,12 +254,9 @@ public Mono createDatabaseIfNotExists(String id, Through */ public Mono createDatabase(CosmosDatabaseProperties databaseProperties, CosmosDatabaseRequestOptions options) { - if (options == null) { - options = new CosmosDatabaseRequestOptions(); - } + final CosmosDatabaseRequestOptions requestOptions = options == null ? new CosmosDatabaseRequestOptions() : options; Database wrappedDatabase = new Database(); wrappedDatabase.setId(databaseProperties.getId()); - final CosmosDatabaseRequestOptions requestOptions = options; return withContext(context -> createDatabaseInternal(wrappedDatabase, requestOptions, context)); } @@ -310,6 +309,7 @@ public Mono createDatabase(CosmosDatabaseProperties data if (options == null) { options = new CosmosDatabaseRequestOptions(); } + ModelBridgeInternal.setThroughputProperties(options, throughputProperties); Database wrappedDatabase = new Database(); wrappedDatabase.setId(databaseProperties.getId()); @@ -378,8 +378,7 @@ public Mono createDatabase(String id, ThroughputProperti */ CosmosPagedFlux readAllDatabases(CosmosQueryRequestOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "readAllDatabases"; - pagedFluxOptions.setTracerInformation(this.tracerProvider, spanName, this.serviceEndpoint, null); + pagedFluxOptions.setTracerInformation(this.tracerProvider, readAllDatabasesSpanName, this.serviceEndpoint, null); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().readDatabases(options) .map(response -> @@ -457,8 +456,7 @@ TracerProvider getTracerProvider(){ private CosmosPagedFlux queryDatabasesInternal(SqlQuerySpec querySpec, CosmosQueryRequestOptions options){ return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "queryDatabases"; - pagedFluxOptions.setTracerInformation(this.tracerProvider, spanName, this.serviceEndpoint, null); + pagedFluxOptions.setTracerInformation(this.tracerProvider, queryDatabasesSpanName, this.serviceEndpoint, null); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().queryDatabases(querySpec, options) .map(response -> BridgeInternal.createFeedResponse( diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index e8625d310857..f937d955b576 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -117,11 +117,7 @@ public Mono read() { * the read container or an error. */ public Mono read(CosmosContainerRequestOptions options) { - if (options == null) { - options = new CosmosContainerRequestOptions(); - } - - final CosmosContainerRequestOptions requestOptions = options; + final CosmosContainerRequestOptions requestOptions = options == null ? new CosmosContainerRequestOptions() : options; return withContext(context -> read(requestOptions, context)); } @@ -137,11 +133,7 @@ public Mono read(CosmosContainerRequestOptions options) * the deleted database or an error. */ public Mono delete(CosmosContainerRequestOptions options) { - if (options == null) { - options = new CosmosContainerRequestOptions(); - } - - final CosmosContainerRequestOptions requestOptions = options; + final CosmosContainerRequestOptions requestOptions = options == null ? new CosmosContainerRequestOptions() : options; return withContext(context -> deleteInternal(requestOptions, context)); } @@ -191,11 +183,7 @@ public Mono replace(CosmosContainerProperties container public Mono replace( CosmosContainerProperties containerProperties, CosmosContainerRequestOptions options) { - if (options == null) { - options = new CosmosContainerRequestOptions(); - } - - final CosmosContainerRequestOptions requestOptions = options; + final CosmosContainerRequestOptions requestOptions = options == null ? new CosmosContainerRequestOptions() : options; return withContext(context -> replaceInternal(containerProperties, requestOptions, context)); } @@ -308,11 +296,7 @@ public Mono> upsertItem(T item) { * @return an {@link Mono} containing the single resource response with the upserted item or an error. */ public Mono> upsertItem(T item, CosmosItemRequestOptions options) { - if (options == null) { - options = new CosmosItemRequestOptions(); - } - - final CosmosItemRequestOptions requestOptions = options; + final CosmosItemRequestOptions requestOptions = options == null ? new CosmosItemRequestOptions() : options; return withContext(context -> upsertItemInternal(item, requestOptions, context)); } @@ -673,10 +657,10 @@ public CosmosAsyncConflict getConflict(String id) { } /** - * Replace the throughput . + * Replace the throughput. * - * @param throughputProperties the throughput properties - * @return the mono containing throughput response + * @param throughputProperties the throughput properties. + * @return the mono containing throughput response. */ public Mono replaceThroughput(ThroughputProperties throughputProperties) { return withContext(context -> replaceThroughputInternal(throughputProperties, context)); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java index 9a341e92e774..91be7ecbfe6d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java @@ -82,11 +82,7 @@ public Mono read() { * the read database or an error. */ public Mono read(CosmosDatabaseRequestOptions options) { - if (options == null) { - options = new CosmosDatabaseRequestOptions(); - } - - final CosmosDatabaseRequestOptions requestOptions = options; + final CosmosDatabaseRequestOptions requestOptions = options == null ? new CosmosDatabaseRequestOptions() : options; return withContext(context -> readInternal(requestOptions, context)); } @@ -114,11 +110,7 @@ public Mono delete() { * @return an {@link Mono} containing the single cosmos database response. */ public Mono delete(CosmosDatabaseRequestOptions options) { - if (options == null) { - options = new CosmosDatabaseRequestOptions(); - } - - final CosmosDatabaseRequestOptions requestOptions = options; + final CosmosDatabaseRequestOptions requestOptions = options == null ? new CosmosDatabaseRequestOptions() : options; return withContext(context -> deleteInternal(requestOptions, context)); } @@ -199,11 +191,8 @@ public Mono createContainer( if (containerProperties == null) { throw new IllegalArgumentException("containerProperties"); } - if (options == null) { - options = new CosmosContainerRequestOptions(); - } - final CosmosContainerRequestOptions requestOptions = options; + final CosmosContainerRequestOptions requestOptions = options == null ? new CosmosContainerRequestOptions() : options; return withContext(context -> createContainerInternal(containerProperties, requestOptions, context)); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java index 52fc7988a3f2..9a449665540f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java @@ -94,13 +94,9 @@ public Mono delete() { public Mono createPermission( CosmosPermissionProperties permissionProperties, CosmosPermissionRequestOptions options) { - if (options == null) { - options = new CosmosPermissionRequestOptions(); - } - + final CosmosPermissionRequestOptions requestOptions = options == null ? new CosmosPermissionRequestOptions() : options; Permission permission = ModelBridgeInternal.getPermission(permissionProperties, database.getId()); - final CosmosPermissionRequestOptions requesOptions = options; - return withContext(context -> createPermissionInternal(permission, requesOptions, context)); + return withContext(context -> createPermissionInternal(permission, requestOptions, context)); } /** @@ -118,11 +114,7 @@ public Mono upsertPermission( CosmosPermissionProperties permissionProperties, CosmosPermissionRequestOptions options) { Permission permission = ModelBridgeInternal.getPermission(permissionProperties, database.getId()); - if (options == null) { - options = new CosmosPermissionRequestOptions(); - } - - final CosmosPermissionRequestOptions requestOptions = options; + final CosmosPermissionRequestOptions requestOptions = options == null ? new CosmosPermissionRequestOptions() : options; return withContext(context -> upsertPermissionInternal(permission, requestOptions, context)); } From 0703a2810d684c09b7738f1c5695df0f823a362a Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Thu, 9 Jul 2020 16:31:59 -0400 Subject: [PATCH 31/40] build error fix --- .../src/main/java/com/azure/cosmos/CosmosAsyncClient.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index 341b3b92a8c5..a41594e1a8ba 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -45,12 +45,12 @@ public final class CosmosAsyncClient implements Closeable { // Async Cosmos client wrapper + private final static String QUERY_DB_SPAN = "queryDatabases"; + private final static String READ_ALL_DB_SPAN = "readAllDatabases"; private final Configs configs; private final AsyncDocumentClient asyncDocumentClient; private final String serviceEndpoint; private final String keyOrResourceToken; - private final String queryDatabasesSpanName = "queryDatabases"; - private final String readAllDatabasesSpanName = "readAllDatabases"; private final ConnectionPolicy connectionPolicy; private final ConsistencyLevel desiredConsistencyLevel; private final List permissions; @@ -378,7 +378,7 @@ public Mono createDatabase(String id, ThroughputProperti */ CosmosPagedFlux readAllDatabases(CosmosQueryRequestOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - pagedFluxOptions.setTracerInformation(this.tracerProvider, readAllDatabasesSpanName, this.serviceEndpoint, null); + pagedFluxOptions.setTracerInformation(this.tracerProvider, READ_ALL_DB_SPAN, this.serviceEndpoint, null); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().readDatabases(options) .map(response -> @@ -456,7 +456,7 @@ TracerProvider getTracerProvider(){ private CosmosPagedFlux queryDatabasesInternal(SqlQuerySpec querySpec, CosmosQueryRequestOptions options){ return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - pagedFluxOptions.setTracerInformation(this.tracerProvider, queryDatabasesSpanName, this.serviceEndpoint, null); + pagedFluxOptions.setTracerInformation(this.tracerProvider, QUERY_DB_SPAN, this.serviceEndpoint, null); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().queryDatabases(querySpec, options) .map(response -> BridgeInternal.createFeedResponse( From deb6978f7642675623ae06dcb34d2a186d66cadf Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Thu, 9 Jul 2020 20:15:21 -0400 Subject: [PATCH 32/40] resolving comment --- .../src/main/java/com/azure/cosmos/CosmosAsyncClient.java | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index a41594e1a8ba..67f87c910c2d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -45,8 +45,6 @@ public final class CosmosAsyncClient implements Closeable { // Async Cosmos client wrapper - private final static String QUERY_DB_SPAN = "queryDatabases"; - private final static String READ_ALL_DB_SPAN = "readAllDatabases"; private final Configs configs; private final AsyncDocumentClient asyncDocumentClient; private final String serviceEndpoint; @@ -378,7 +376,7 @@ public Mono createDatabase(String id, ThroughputProperti */ CosmosPagedFlux readAllDatabases(CosmosQueryRequestOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - pagedFluxOptions.setTracerInformation(this.tracerProvider, READ_ALL_DB_SPAN, this.serviceEndpoint, null); + pagedFluxOptions.setTracerInformation(this.tracerProvider, "readAllDatabases", this.serviceEndpoint, null); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().readDatabases(options) .map(response -> @@ -456,7 +454,7 @@ TracerProvider getTracerProvider(){ private CosmosPagedFlux queryDatabasesInternal(SqlQuerySpec querySpec, CosmosQueryRequestOptions options){ return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - pagedFluxOptions.setTracerInformation(this.tracerProvider, QUERY_DB_SPAN, this.serviceEndpoint, null); + pagedFluxOptions.setTracerInformation(this.tracerProvider, "queryDatabases", this.serviceEndpoint, null); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().queryDatabases(querySpec, options) .map(response -> BridgeInternal.createFeedResponse( From 8af1cc9d56a405a9ec2bb72bac142b9d8c02890e Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Fri, 10 Jul 2020 11:24:50 -0400 Subject: [PATCH 33/40] resolving comments --- .../com/azure/cosmos/CosmosAsyncClient.java | 33 +++++++++++++------ .../com/azure/cosmos/CosmosAsyncDatabase.java | 33 ++++++------------- 2 files changed, 33 insertions(+), 33 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index 67f87c910c2d..6970c91cbda8 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -468,7 +468,25 @@ private Mono createDatabaseIfNotExistsInternal(CosmosAsy ThroughputProperties throughputProperties, Context context) { String spanName = "createDatabaseIfNotExists." + database.getId(); Context nestedContext = context.addData(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.COSMOS_CALL_DEPTH_VAL); - Mono responseMono = createDatabaseIfNotExistsInternal(database.readInternal(new CosmosDatabaseRequestOptions(), nestedContext), database, throughputProperties, nestedContext); + Mono responseMono = database.readInternal(new CosmosDatabaseRequestOptions(), + nestedContext).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); + } + + Database wrappedDatabase = new Database(); + wrappedDatabase.setId(database.getId()); + return createDatabaseInternal(wrappedDatabase, + requestOptions, nestedContext); + } + } + return Mono.error(unwrappedException); + }); return tracerProvider.traceEnabledCosmosResponsePublisher(responseMono, context, spanName, @@ -487,15 +505,10 @@ private Mono createDatabaseIfNotExistsInternal(Mono createContainerIfNotExists( * Creates a Cosmos container if it does not exist on the service. *

* The throughput setting will only be used if the specified container - * does not exist and therefor a new container will be created. + * does not exist and therefore a new container will be created. * * After subscription the operation will be performed. The {@link Mono} upon * successful completion will contain a cosmos container response with the @@ -732,36 +732,23 @@ private Mono createContainerIfNotExistsInternal( Context context) { String spanName = "createContainerIfNotExistsInternal." + containerProperties.getId(); Context nestedContext = context.addData(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.COSMOS_CALL_DEPTH_VAL); - if (options == null) { - options = new CosmosContainerRequestOptions(); - } - - Mono responseMono = createContainerIfNotExistsInternal(container.read(options, nestedContext), containerProperties, options, nestedContext); - return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, - getId(), - getClient().getServiceEndpoint()); - } - - private Mono createContainerIfNotExistsInternal( - Mono responseMono, - CosmosContainerProperties containerProperties, - CosmosContainerRequestOptions options, - Context context) { - return responseMono.onErrorResume(exception -> { + final CosmosContainerRequestOptions requestOptions = options == null ? new CosmosContainerRequestOptions() : + options; + Mono responseMono = + container.read(requestOptions, nestedContext).onErrorResume(exception -> { final Throwable unwrappedException = Exceptions.unwrap(exception); if (unwrappedException instanceof CosmosException) { final CosmosException cosmosException = (CosmosException) unwrappedException; if (cosmosException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) { - if(context != null) { - return createContainerInternal(containerProperties, options, context); - } - - return createContainer(containerProperties, options); + return createContainerInternal(containerProperties, requestOptions, nestedContext); } } return Mono.error(unwrappedException); }); + return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, + getId(), + getClient().getServiceEndpoint()); } private Mono createContainerInternal( From 798839c218d086be39a3463b17a681d39c179084 Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Fri, 10 Jul 2020 11:33:32 -0400 Subject: [PATCH 34/40] resolving comment --- .../com/azure/cosmos/CosmosAsyncPermission.java | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java index a1f8c868fc88..ead21b5b8cb5 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java @@ -4,9 +4,9 @@ import com.azure.core.util.Context; import com.azure.cosmos.implementation.Paths; -import com.azure.cosmos.models.CosmosPermissionResponse; import com.azure.cosmos.models.CosmosPermissionProperties; import com.azure.cosmos.models.CosmosPermissionRequestOptions; +import com.azure.cosmos.models.CosmosPermissionResponse; import com.azure.cosmos.models.ModelBridgeInternal; import reactor.core.publisher.Mono; @@ -125,20 +125,15 @@ String getLink() { private Mono readInternal(CosmosPermissionRequestOptions options, Context context) { String spanName = "readPermission." + cosmosUser.getId(); - Mono responseMono = readInternal(options); - return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, - cosmosUser.getDatabase().getId(), - cosmosUser.getDatabase().getClient().getServiceEndpoint()); - } - - private Mono readInternal(CosmosPermissionRequestOptions options) { - - return cosmosUser.getDatabase() + Mono responseMono = cosmosUser.getDatabase() .getDocClientWrapper() .readPermission(getLink(), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosPermissionResponse(response)) .single(); + return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, + cosmosUser.getDatabase().getId(), + cosmosUser.getDatabase().getClient().getServiceEndpoint()); } private Mono replaceInternal(CosmosPermissionProperties permissionProperties, From 21c5743709725058254427079d898d2ad810da90 Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Fri, 10 Jul 2020 11:34:00 -0400 Subject: [PATCH 35/40] resolving comment --- .../java/com/azure/cosmos/implementation/TracerProvider.java | 1 - 1 file changed, 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java index 665fba9e184a..432943a11b76 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java @@ -26,7 +26,6 @@ public class TracerProvider { public static final String DB_STATEMENT = "db.statement"; public static final String ERROR_MSG = "error.msg"; public static final String ERROR_TYPE = "error.type"; - public static final String ERROR_STACK = "error.stack"; public static final String COSMOS_CALL_DEPTH = "cosmosCallDepth"; public static final String COSMOS_CALL_DEPTH_VAL = "nested"; public static final int ERROR_CODE = 0; From 42e97890b47c2cfb06745bc4593003bd8f944b47 Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Fri, 10 Jul 2020 13:06:25 -0400 Subject: [PATCH 36/40] removing extra method --- .../com/azure/cosmos/CosmosAsyncClient.java | 40 ------------------- 1 file changed, 40 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index 6970c91cbda8..f8d694da9239 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -333,24 +333,6 @@ public Mono createDatabase(CosmosDatabaseProperties data return createDatabase(databaseProperties, options); } - /** - * Creates a database. - *

- * After subscription the operation will be performed. - * The {@link Mono} upon successful completion will contain a single resource response with the - * created database. - * In case of failure the {@link Mono} will error. - * - * @param id id of the database. - * @param throughput the throughput for the database. - * @return a {@link Mono} containing the single cosmos database response with the created database or an error. - */ - Mono createDatabase(String id, int throughput) { - CosmosDatabaseRequestOptions options = new CosmosDatabaseRequestOptions(); - ModelBridgeInternal.setThroughputProperties(options, ThroughputProperties.createManualThroughput(throughput)); - return createDatabase(new CosmosDatabaseProperties(id), options); - } - /** * Creates a database. * @@ -494,28 +476,6 @@ private Mono createDatabaseIfNotExistsInternal(CosmosAsy this.serviceEndpoint); } - private Mono createDatabaseIfNotExistsInternal(Mono 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); - } - - Database wrappedDatabase = new Database(); - wrappedDatabase.setId(database.getId()); - return createDatabaseInternal(wrappedDatabase, - requestOptions, context); - } - } - return Mono.error(unwrappedException); - }); - } - - private Mono createDatabaseInternal(Database database, CosmosDatabaseRequestOptions options, Context context) { String spanName = "createDatabase." + database.getId(); From be8d7ebbdd6ca856946093d058fa3045b56316fd Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Fri, 10 Jul 2020 16:16:12 -0400 Subject: [PATCH 37/40] removing tets dependency --- sdk/cosmos/azure-cosmos/pom.xml | 12 --------- .../com/azure/cosmos/CosmosTracerTest.java | 27 ++++++++++++------- 2 files changed, 18 insertions(+), 21 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/pom.xml b/sdk/cosmos/azure-cosmos/pom.xml index 280b7ee1f194..370b7dd718f4 100644 --- a/sdk/cosmos/azure-cosmos/pom.xml +++ b/sdk/cosmos/azure-cosmos/pom.xml @@ -128,18 +128,6 @@ Licensed under the MIT License. - - com.azure - azure-core-tracing-opentelemetry - test - 1.0.0-beta.5 - - - io.netty - netty-tcnative-boringssl-static - - - com.fasterxml.jackson.module jackson-module-afterburner diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java index 2307a26c781a..963581755186 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java @@ -24,7 +24,8 @@ import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; -import java.util.ServiceLoader; +import java.util.ArrayList; +import java.util.List; import java.util.UUID; public class CosmosTracerTest extends TestSuiteBase { @@ -47,26 +48,26 @@ public void beforeClass() { @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void cosmosAsyncClient() { - TracerProvider tracer = Mockito.spy(new TracerProvider(ServiceLoader.load(Tracer.class))); - ReflectionUtils.setTracerProvider(client, tracer); + TracerProvider tracerProvider = Mockito.spy(new TracerProvider(getMockTracer())); + ReflectionUtils.setTracerProvider(client, tracerProvider); client.createDatabaseIfNotExists(cosmosAsyncDatabase.getId()).block(); - Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); client.readAllDatabases(new CosmosQueryRequestOptions()).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); String query = "select * from c where c.id = '" + cosmosAsyncDatabase.getId() + "'"; client.queryDatabases(query, new CosmosQueryRequestOptions()).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void cosmosAsyncDatabase() { - TracerProvider tracer = Mockito.spy(new TracerProvider(ServiceLoader.load(Tracer.class))); + TracerProvider tracer = Mockito.spy(new TracerProvider(getMockTracer())); ReflectionUtils.setTracerProvider(client, tracer); cosmosAsyncDatabase.createContainerIfNotExists(cosmosAsyncContainer.getId(), @@ -94,7 +95,7 @@ public void cosmosAsyncDatabase() { @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void cosmosAsyncContainer() { - TracerProvider tracer = Mockito.spy(new TracerProvider(ServiceLoader.load(Tracer.class))); + TracerProvider tracer = Mockito.spy(new TracerProvider(getMockTracer())); ReflectionUtils.setTracerProvider(client, tracer); cosmosAsyncContainer.read().block(); @@ -140,7 +141,7 @@ public void cosmosAsyncContainer() { @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void cosmosAsyncScripts() { - TracerProvider tracer = Mockito.spy(new TracerProvider(ServiceLoader.load(Tracer.class))); + TracerProvider tracer = Mockito.spy(new TracerProvider(getMockTracer())); ReflectionUtils.setTracerProvider(client, tracer); cosmosAsyncContainer.getScripts().readAllStoredProcedures(new CosmosQueryRequestOptions()).byPage().single().block(); @@ -242,4 +243,12 @@ private static CosmosStoredProcedureProperties getCosmosStoredProcedurePropertie CosmosStoredProcedureProperties storedProcedureDef = new CosmosStoredProcedureProperties(UUID.randomUUID().toString(), "function() {var x = 10;}"); return storedProcedureDef; } + + private List getMockTracer() { + Tracer tracer = Mockito.mock(Tracer.class); + List tracerList = new ArrayList<>(); + tracerList.add(tracer); + Mockito.when(tracer.start(Matchers.anyString(), Matchers.any(Context.class))).thenReturn(Context.NONE); + return tracerList; + } } From 1f9725c8608fb80cbb446f5f4f0c94b5754b97db Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Mon, 13 Jul 2020 13:02:12 -0400 Subject: [PATCH 38/40] adding span attribute unit test --- .../com/azure/cosmos/CosmosAsyncConflict.java | 20 +- .../azure/cosmos/CosmosAsyncContainer.java | 2 +- .../com/azure/cosmos/CosmosAsyncDatabase.java | 2 +- .../com/azure/cosmos/CosmosTracerTest.java | 205 ++++++++++++++---- 4 files changed, 172 insertions(+), 57 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncConflict.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncConflict.java index 08177c49043c..7a69191a459a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncConflict.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncConflict.java @@ -5,8 +5,8 @@ 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.CosmosConflictResponse; import com.azure.cosmos.models.ModelBridgeInternal; import reactor.core.publisher.Mono; @@ -109,7 +109,9 @@ String getLink() { private Mono readInternal(RequestOptions options, Context context) { String spanName = "readConflict." + getId(); - Mono responseMono = this.readInternal(options); + Mono responseMono = + this.container.getDatabase().getDocClientWrapper().readConflict(getLink(), options) + .map(response -> ModelBridgeInternal.createCosmosConflictResponse(response)).single(); return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, this.container.getDatabase().getId(), @@ -117,22 +119,14 @@ private Mono readInternal(RequestOptions options, Contex } - private Mono readInternal(RequestOptions options) { - return this.container.getDatabase().getDocClientWrapper().readConflict(getLink(), options) - .map(response -> ModelBridgeInternal.createCosmosConflictResponse(response)).single(); - } - private Mono deleteInternal(RequestOptions options, Context context) { String spanName = "deleteConflict." + getId(); - Mono responseMono = deleteInternal(options); + Mono responseMono = + this.container.getDatabase().getDocClientWrapper().deleteConflict(getLink(), options) + .map(response -> ModelBridgeInternal.createCosmosConflictResponse(response)).single(); return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, this.container.getDatabase().getId(), this.container.getDatabase().getClient().getServiceEndpoint()); } - - private Mono deleteInternal(RequestOptions options) { - return this.container.getDatabase().getDocClientWrapper().deleteConflict(getLink(), options) - .map(response -> ModelBridgeInternal.createCosmosConflictResponse(response)).single(); - } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index f937d955b576..b967eb8b73d1 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -531,7 +531,7 @@ public Mono> replaceItem( @SuppressWarnings("unchecked") Class itemType = (Class) item.getClass(); final CosmosItemRequestOptions requestOptions = options; - return withContext(context -> replaceItemInternal(itemType, itemId, doc, requestOptions,context)); + return withContext(context -> replaceItemInternal(itemType, itemId, doc, requestOptions, context)); } /** diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java index 48514a5900d7..933847f78210 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java @@ -730,7 +730,7 @@ private Mono createContainerIfNotExistsInternal( CosmosAsyncContainer container, CosmosContainerRequestOptions options, Context context) { - String spanName = "createContainerIfNotExistsInternal." + containerProperties.getId(); + String spanName = "createContainerIfNotExists." + containerProperties.getId(); Context nestedContext = context.addData(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.COSMOS_CALL_DEPTH_VAL); final CosmosContainerRequestOptions requestOptions = options == null ? new CosmosContainerRequestOptions() : options; diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java index 963581755186..29ef071f0867 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java @@ -20,6 +20,8 @@ import com.azure.cosmos.rx.TestSuiteBase; import org.mockito.Matchers; import org.mockito.Mockito; +import org.mockito.invocation.InvocationOnMock; +import org.mockito.stubbing.Answer; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; @@ -48,32 +50,52 @@ public void beforeClass() { @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void cosmosAsyncClient() { - TracerProvider tracerProvider = Mockito.spy(new TracerProvider(getMockTracer())); + Tracer mockTracer = Mockito.mock(Tracer.class); + TracerProvider tracerProvider = Mockito.spy(new TracerProvider(getMockTracer(mockTracer))); ReflectionUtils.setTracerProvider(client, tracerProvider); + TracerProviderCapture tracerProviderCapture = new TracerProviderCapture(); + Mockito.doAnswer(tracerProviderCapture).when(tracerProvider).startSpan(Matchers.anyString(), + Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); + client.createDatabaseIfNotExists(cosmosAsyncDatabase.getId()).block(); + Context context = tracerProviderCapture.getResult(); Mockito.verify(tracerProvider, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "createDatabaseIfNotExists." + cosmosAsyncDatabase.getId(), context, + cosmosAsyncDatabase.getId(), 1); client.readAllDatabases(new CosmosQueryRequestOptions()).byPage().single().block(); Mockito.verify(tracerProvider, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "readAllDatabases", context, null, 2); + String query = "select * from c where c.id = '" + cosmosAsyncDatabase.getId() + "'"; client.queryDatabases(query, new CosmosQueryRequestOptions()).byPage().single().block(); Mockito.verify(tracerProvider, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "queryDatabases", context, null, 3); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void cosmosAsyncDatabase() { - TracerProvider tracer = Mockito.spy(new TracerProvider(getMockTracer())); - ReflectionUtils.setTracerProvider(client, tracer); + Tracer mockTracer = Mockito.mock(Tracer.class); + TracerProvider tracerProvider = Mockito.spy(new TracerProvider(getMockTracer(mockTracer))); + ReflectionUtils.setTracerProvider(client, tracerProvider); + TracerProviderCapture tracerProviderCapture = new TracerProviderCapture(); + Mockito.doAnswer(tracerProviderCapture).when(tracerProvider).startSpan(Matchers.anyString(), + Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncDatabase.createContainerIfNotExists(cosmosAsyncContainer.getId(), - "/pk", 5000).block(); - Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), + "/pk", 5000).block(); + Context context = tracerProviderCapture.getResult(); + Mockito.verify(tracerProvider, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "createContainerIfNotExists." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 1); try { cosmosAsyncDatabase.readThroughput().block(); @@ -81,145 +103,213 @@ public void cosmosAsyncDatabase() { //do nothing } - Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "readThroughput." + cosmosAsyncDatabase.getId(), context, + cosmosAsyncDatabase.getId(), 2); cosmosAsyncDatabase.readAllUsers().byPage().single().block(); - Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "readAllUsers." + cosmosAsyncDatabase.getId(), context, + cosmosAsyncDatabase.getId(), 3); cosmosAsyncDatabase.readAllContainers().byPage().single().block(); - Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "readAllContainers." + cosmosAsyncDatabase.getId(), context, + cosmosAsyncDatabase.getId(), 4); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void cosmosAsyncContainer() { - TracerProvider tracer = Mockito.spy(new TracerProvider(getMockTracer())); - ReflectionUtils.setTracerProvider(client, tracer); + Tracer mockTracer = Mockito.mock(Tracer.class); + TracerProvider tracerProvider = Mockito.spy(new TracerProvider(getMockTracer(mockTracer))); + ReflectionUtils.setTracerProvider(client, tracerProvider); + TracerProviderCapture tracerProviderCapture = new TracerProviderCapture(); + Mockito.doAnswer(tracerProviderCapture).when(tracerProvider).startSpan(Matchers.anyString(), + Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.read().block(); - Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), + Context context = tracerProviderCapture.getResult(); + Mockito.verify(tracerProvider, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "readContainer." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 1); try { cosmosAsyncContainer.readThroughput().block(); } catch (CosmosException ex) { //do nothing } - Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "readThroughput." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 2); InternalObjectNode item = new InternalObjectNode(); item.setId(ITEM_ID); cosmosAsyncContainer.createItem(item).block(); - Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "createItem." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 3); cosmosAsyncContainer.upsertItem(item, new CosmosItemRequestOptions()).block(); - Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "upsertItem." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 4); cosmosAsyncContainer.readItem(ITEM_ID, PartitionKey.NONE, InternalObjectNode.class).block(); - Mockito.verify(tracer, Mockito.times(5)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(5)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "readItem." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 5); cosmosAsyncContainer.deleteItem(ITEM_ID, PartitionKey.NONE).block(); - Mockito.verify(tracer, Mockito.times(6)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(6)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "deleteItem." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 6); + cosmosAsyncContainer.readAllItems(new CosmosQueryRequestOptions(), CosmosItemRequestOptions.class).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(7)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(7)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "readAllItems." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 7); String query = "select * from c where c.id = '" + ITEM_ID + "'"; cosmosAsyncContainer.queryItems(query, new CosmosQueryRequestOptions(), CosmosItemRequestOptions.class).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(8)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(8)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "queryItems." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 8); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void cosmosAsyncScripts() { - TracerProvider tracer = Mockito.spy(new TracerProvider(getMockTracer())); - ReflectionUtils.setTracerProvider(client, tracer); + Tracer mockTracer = Mockito.mock(Tracer.class); + TracerProvider tracerProvider = Mockito.spy(new TracerProvider(getMockTracer(mockTracer))); + ReflectionUtils.setTracerProvider(client, tracerProvider); + TracerProviderCapture tracerProviderCapture = new TracerProviderCapture(); + Mockito.doAnswer(tracerProviderCapture).when(tracerProvider).startSpan(Matchers.anyString(), + Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.getScripts().readAllStoredProcedures(new CosmosQueryRequestOptions()).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), + Context context = tracerProviderCapture.getResult(); + + Mockito.verify(tracerProvider, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "readAllStoredProcedures." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 1); cosmosAsyncContainer.getScripts().readAllTriggers(new CosmosQueryRequestOptions()).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "readAllTriggers." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 2); cosmosAsyncContainer.getScripts().readAllUserDefinedFunctions(new CosmosQueryRequestOptions()).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "readAllUserDefinedFunctions." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 3); CosmosUserDefinedFunctionProperties cosmosUserDefinedFunctionProperties = getCosmosUserDefinedFunctionProperties(); CosmosUserDefinedFunctionProperties resultUdf = cosmosAsyncContainer.getScripts().createUserDefinedFunction(cosmosUserDefinedFunctionProperties).block().getProperties(); - Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "createUserDefinedFunction." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 4); cosmosAsyncContainer.getScripts().getUserDefinedFunction(cosmosUserDefinedFunctionProperties.getId()).read().block(); - Mockito.verify(tracer, Mockito.times(5)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(5)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "readUDF." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 5); cosmosUserDefinedFunctionProperties.setBody("function() {var x = 15;}"); cosmosAsyncContainer.getScripts().getUserDefinedFunction(resultUdf.getId()).replace(resultUdf).block(); - Mockito.verify(tracer, Mockito.times(6)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(6)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "replaceUDF." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 6); cosmosAsyncContainer.getScripts().readAllUserDefinedFunctions(new CosmosQueryRequestOptions()).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(7)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(7)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.getScripts().getUserDefinedFunction(cosmosUserDefinedFunctionProperties.getId()).delete().block(); - Mockito.verify(tracer, Mockito.times(8)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(8)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "deleteUDF." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 8); CosmosTriggerProperties cosmosTriggerProperties = getCosmosTriggerProperties(); CosmosTriggerProperties resultTrigger = cosmosAsyncContainer.getScripts().createTrigger(cosmosTriggerProperties).block().getProperties(); - Mockito.verify(tracer, Mockito.times(9)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(9)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "createTrigger." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 9); cosmosAsyncContainer.getScripts().getTrigger(cosmosTriggerProperties.getId()).read().block(); - Mockito.verify(tracer, Mockito.times(10)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(10)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "readTrigger." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 10); cosmosAsyncContainer.getScripts().getTrigger(cosmosTriggerProperties.getId()).replace(resultTrigger).block(); - Mockito.verify(tracer, Mockito.times(11)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(11)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "replaceTrigger." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 11); + cosmosAsyncContainer.getScripts().readAllTriggers(new CosmosQueryRequestOptions()).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(12)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(12)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); cosmosAsyncContainer.getScripts().getTrigger(cosmosTriggerProperties.getId()).delete().block(); - Mockito.verify(tracer, Mockito.times(13)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(13)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "deleteTrigger." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 13); CosmosStoredProcedureProperties procedureProperties = getCosmosStoredProcedureProperties(); CosmosStoredProcedureProperties resultSproc = cosmosAsyncContainer.getScripts().createStoredProcedure(procedureProperties).block().getProperties(); - Mockito.verify(tracer, Mockito.times(14)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(14)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "createStoredProcedure." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 14); cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).read().block(); - Mockito.verify(tracer, Mockito.times(15)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(15)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "readStoredProcedure." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 15); + cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).replace(resultSproc).block(); - Mockito.verify(tracer, Mockito.times(16)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(16)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "replaceStoredProcedure." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 16); cosmosAsyncContainer.getScripts().readAllStoredProcedures(new CosmosQueryRequestOptions()).byPage().single().block(); cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).delete().block(); - Mockito.verify(tracer, Mockito.times(18)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(18)).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "deleteStoredProcedure." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), 18); } @AfterClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT) @@ -228,27 +318,58 @@ public void afterClass() { } private static CosmosUserDefinedFunctionProperties getCosmosUserDefinedFunctionProperties() { - CosmosUserDefinedFunctionProperties udf = new CosmosUserDefinedFunctionProperties(UUID.randomUUID().toString(), "function() {var x = 10;}"); + CosmosUserDefinedFunctionProperties udf = + new CosmosUserDefinedFunctionProperties(UUID.randomUUID().toString(), "function() {var x = 10;}"); return udf; } private static CosmosTriggerProperties getCosmosTriggerProperties() { - CosmosTriggerProperties trigger = new CosmosTriggerProperties(UUID.randomUUID().toString(), "function() {var x = 10;}"); + CosmosTriggerProperties trigger = new CosmosTriggerProperties(UUID.randomUUID().toString(), "function() {var " + + "x = 10;}"); trigger.setTriggerOperation(TriggerOperation.CREATE); trigger.setTriggerType(TriggerType.PRE); return trigger; } private static CosmosStoredProcedureProperties getCosmosStoredProcedureProperties() { - CosmosStoredProcedureProperties storedProcedureDef = new CosmosStoredProcedureProperties(UUID.randomUUID().toString(), "function() {var x = 10;}"); + CosmosStoredProcedureProperties storedProcedureDef = + new CosmosStoredProcedureProperties(UUID.randomUUID().toString(), "function() {var x = 10;}"); return storedProcedureDef; } - private List getMockTracer() { - Tracer tracer = Mockito.mock(Tracer.class); + private List getMockTracer(Tracer tracer) { List tracerList = new ArrayList<>(); tracerList.add(tracer); Mockito.when(tracer.start(Matchers.anyString(), Matchers.any(Context.class))).thenReturn(Context.NONE); return tracerList; } + + private void verifyTracerAttributes(Tracer mockTracer, String methodName, Context context, String databaseName, + int numberOfTimesCalledWithinTest) { + if (databaseName != null) { + Mockito.verify(mockTracer, Mockito.times(numberOfTimesCalledWithinTest)).setAttribute(TracerProvider.DB_INSTANCE, + databaseName, context); + } + Mockito.verify(mockTracer, Mockito.times(numberOfTimesCalledWithinTest)).setAttribute(TracerProvider.DB_TYPE, + TracerProvider.DB_TYPE_VALUE, context); + Mockito.verify(mockTracer, Mockito.times(numberOfTimesCalledWithinTest)).setAttribute(TracerProvider.DB_URL, + TestConfigurations.HOST, + context); + Mockito.verify(mockTracer, Mockito.times(1)).setAttribute(TracerProvider.DB_STATEMENT, methodName, context); + + } + + private class TracerProviderCapture implements Answer { + private Context result = null; + + public Context getResult() { + return result; + } + + @Override + public Context answer(InvocationOnMock invocationOnMock) throws Throwable { + result = (Context) invocationOnMock.callRealMethod(); + return result; + } + } } From fc1296f69012aaf51b1b229ed1e36625cec14bc6 Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Mon, 13 Jul 2020 13:09:16 -0400 Subject: [PATCH 39/40] resolving comments --- .../java/com/azure/cosmos/CosmosAsyncContainer.java | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index b967eb8b73d1..cd53a71718c8 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -783,7 +783,9 @@ private Mono> readItemInternal( } Mono read(CosmosContainerRequestOptions options, Context context) { - Mono responseMono = readInternal(options); + Mono responseMono = database.getDocClientWrapper().readCollection(getLink(), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosContainerResponse(response)).single(); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, this.readContainerSpanName, @@ -791,12 +793,6 @@ Mono read(CosmosContainerRequestOptions options, Contex database.getClient().getServiceEndpoint()); } - private Mono readInternal(CosmosContainerRequestOptions options) { - return database.getDocClientWrapper().readCollection(getLink(), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosContainerResponse(response)).single(); - } - private Mono deleteInternal(CosmosContainerRequestOptions options, Context context) { Mono responseMono = deleteInternal(options); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, From 0a7aab07efdce2088ef79cd27ad62576d6064a25 Mon Sep 17 00:00:00 2001 From: Naveen Kumar Singh Date: Tue, 14 Jul 2020 14:33:13 -0400 Subject: [PATCH 40/40] resolving comments --- .../com/azure/cosmos/CosmosAsyncClient.java | 10 +- .../azure/cosmos/CosmosAsyncContainer.java | 83 ++--- .../com/azure/cosmos/CosmosAsyncDatabase.java | 50 +--- .../azure/cosmos/CosmosAsyncPermission.java | 29 +- .../cosmos/CosmosAsyncStoredProcedure.java | 53 ++-- .../com/azure/cosmos/CosmosAsyncTrigger.java | 36 +-- .../com/azure/cosmos/CosmosAsyncUser.java | 58 ++-- .../CosmosAsyncUserDefinedFunction.java | 38 +-- .../com/azure/cosmos/CosmosTracerTest.java | 283 +++++++++++++----- 9 files changed, 314 insertions(+), 326 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index f8d694da9239..2a311b6dccdb 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java @@ -479,17 +479,13 @@ private Mono createDatabaseIfNotExistsInternal(CosmosAsy private Mono createDatabaseInternal(Database database, CosmosDatabaseRequestOptions options, Context context) { String spanName = "createDatabase." + database.getId(); - Mono responseMono = createDatabaseInternal(database, options); + Mono responseMono = asyncDocumentClient.createDatabase(database, ModelBridgeInternal.toRequestOptions(options)) + .map(databaseResourceResponse -> ModelBridgeInternal.createCosmosDatabaseResponse(databaseResourceResponse)) + .single(); return tracerProvider.traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), this.serviceEndpoint); } - - private Mono createDatabaseInternal(Database database, CosmosDatabaseRequestOptions options) { - return asyncDocumentClient.createDatabase(database, ModelBridgeInternal.toRequestOptions(options)) - .map(databaseResourceResponse -> ModelBridgeInternal.createCosmosDatabaseResponse(databaseResourceResponse)) - .single(); - } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index cd53a71718c8..b14175f7618c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -700,7 +700,11 @@ private Mono> deleteItemInternal( String itemId, RequestOptions requestOptions, Context context) { - Mono> responseMono = deleteItemInternal(itemId, requestOptions); + Mono> responseMono = this.getDatabase() + .getDocClientWrapper() + .deleteDocument(getItemLink(itemId), requestOptions) + .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponseWithObjectType(response)) + .single(); return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, context, this.deleteItemSpanName, @@ -708,78 +712,50 @@ private Mono> deleteItemInternal( database.getClient().getServiceEndpoint()); } - private Mono> deleteItemInternal( - String itemId, - RequestOptions requestOptions) { - return this.getDatabase() - .getDocClientWrapper() - .deleteDocument(getItemLink(itemId), requestOptions) - .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponseWithObjectType(response)) - .single(); - } - private Mono> replaceItemInternal( Class itemType, String itemId, Document doc, CosmosItemRequestOptions options, Context context) { - Mono> responseMono = replaceItemInternal(itemType, itemId, doc, options); - return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, - context, this.replaceItemSpanName, database.getId(), database.getClient().getServiceEndpoint()); - } - - private Mono> replaceItemInternal( - Class itemType, - String itemId, - Document doc, - CosmosItemRequestOptions options) { - return this.getDatabase() + Mono> responseMono = this.getDatabase() .getDocClientWrapper() .replaceDocument(getItemLink(itemId), doc, ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) .single(); - } - - private Mono> upsertItemInternal(T item, CosmosItemRequestOptions options, Context context) { - Mono> responseMono = upsertItemInternal(item, options); return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, - context, - this.upsertItemSpanName, - database.getId(), - database.getClient().getServiceEndpoint()); + context, this.replaceItemSpanName, database.getId(), database.getClient().getServiceEndpoint()); } - private Mono> upsertItemInternal(T item, CosmosItemRequestOptions options) { + private Mono> upsertItemInternal(T item, CosmosItemRequestOptions options, Context context) { @SuppressWarnings("unchecked") Class itemType = (Class) item.getClass(); - return this.getDatabase().getDocClientWrapper() + Mono> responseMono = this.getDatabase().getDocClientWrapper() .upsertDocument(this.getLink(), item, ModelBridgeInternal.toRequestOptions(options), true) .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) .single(); - } - - private Mono> readItemInternal( - String itemId, - RequestOptions requestOptions, Class itemType, - Context context) { - Mono> responseMono = readItemInternal(itemId, requestOptions, itemType); return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, context, - this.readItemSpanName, + this.upsertItemSpanName, database.getId(), database.getClient().getServiceEndpoint()); } private Mono> readItemInternal( String itemId, - RequestOptions requestOptions, Class itemType) { - return this.getDatabase().getDocClientWrapper() + RequestOptions requestOptions, Class itemType, + Context context) { + Mono> responseMono = this.getDatabase().getDocClientWrapper() .readDocument(getItemLink(itemId), requestOptions) .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) .single(); + return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, + context, + this.readItemSpanName, + database.getId(), + database.getClient().getServiceEndpoint()); } Mono read(CosmosContainerRequestOptions options, Context context) { @@ -794,7 +770,9 @@ Mono read(CosmosContainerRequestOptions options, Contex } private Mono deleteInternal(CosmosContainerRequestOptions options, Context context) { - Mono responseMono = deleteInternal(options); + Mono responseMono = database.getDocClientWrapper().deleteCollection(getLink(), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosContainerResponse(response)).single(); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, this.deleteContainerSpanName, @@ -802,16 +780,13 @@ private Mono deleteInternal(CosmosContainerRequestOptio database.getClient().getServiceEndpoint()); } - private Mono deleteInternal(CosmosContainerRequestOptions options) { - return database.getDocClientWrapper().deleteCollection(getLink(), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosContainerResponse(response)).single(); - } - private Mono replaceInternal(CosmosContainerProperties containerProperties, CosmosContainerRequestOptions options, Context context) { - Mono responseMono = replaceInternal(containerProperties, options); + Mono responseMono = database.getDocClientWrapper() + .replaceCollection(ModelBridgeInternal.getV2Collection(containerProperties), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosContainerResponse(response)).single(); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, this.replaceContainerSpanName, @@ -819,14 +794,6 @@ private Mono replaceInternal(CosmosContainerProperties database.getClient().getServiceEndpoint()); } - private Mono replaceInternal(CosmosContainerProperties containerProperties, - CosmosContainerRequestOptions options) { - return database.getDocClientWrapper() - .replaceCollection(ModelBridgeInternal.getV2Collection(containerProperties), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosContainerResponse(response)).single(); - } - private Mono readThroughputInternal(Context context) { Context nestedContext = context.addData(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.COSMOS_CALL_DEPTH_VAL); Mono responseMono = readThroughputInternal(this.read(new CosmosContainerRequestOptions(), diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java index 933847f78210..dc956b9318a4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java @@ -756,78 +756,56 @@ private Mono createContainerInternal( CosmosContainerRequestOptions options, Context context) { String spanName = "createContainer." + containerProperties.getId(); - Mono responseMono = createContainerInternal(containerProperties, options); + Mono responseMono = getDocClientWrapper() + .createCollection(this.getLink(), ModelBridgeInternal.getV2Collection(containerProperties), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosContainerResponse(response)).single(); return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, getId(), getClient().getServiceEndpoint()); } - private Mono createContainerInternal( - CosmosContainerProperties containerProperties, - CosmosContainerRequestOptions options) { - return getDocClientWrapper() - .createCollection(this.getLink(), ModelBridgeInternal.getV2Collection(containerProperties), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosContainerResponse(response)).single(); - } - Mono readInternal(CosmosDatabaseRequestOptions options, Context context) { String spanName = "readDatabase." + this.getId(); - Mono responseMono = readInternal(options); + Mono responseMono = getDocClientWrapper().readDatabase(getLink(), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosDatabaseResponse(response)).single(); return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, getId(), getClient().getServiceEndpoint()); } - private Mono readInternal(CosmosDatabaseRequestOptions options) { - return getDocClientWrapper().readDatabase(getLink(), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosDatabaseResponse(response)).single(); - } - private Mono deleteInternal(CosmosDatabaseRequestOptions options, Context context) { String spanName = "deleteDatabase." + this.getId(); - Mono responseMono = deleteInternal(options); + Mono responseMono = getDocClientWrapper().deleteDatabase(getLink(), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosDatabaseResponse(response)).single(); return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, getId(), getClient().getServiceEndpoint()); } - private Mono deleteInternal(CosmosDatabaseRequestOptions options) { - return getDocClientWrapper().deleteDatabase(getLink(), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosDatabaseResponse(response)).single(); - } - private Mono createUserInternal(CosmosUserProperties userProperties, Context context) { String spanName = "createUser." + this.getId(); - Mono responseMono = createUserInternal(userProperties); + Mono responseMono = getDocClientWrapper().createUser(this.getLink(), ModelBridgeInternal.getV2User(userProperties), null) + .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, getId(), getClient().getServiceEndpoint()); } - private Mono createUserInternal(CosmosUserProperties userProperties) { - return getDocClientWrapper().createUser(this.getLink(), ModelBridgeInternal.getV2User(userProperties), null) - .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); - } - private Mono upsertUserInternal(CosmosUserProperties userProperties, Context context) { String spanName = "upsertUser." + this.getId(); - Mono responseMono = upsertUserInternal(userProperties); + Mono responseMono = getDocClientWrapper().upsertUser(this.getLink(), ModelBridgeInternal.getV2User(userProperties), null) + .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, getId(), getClient().getServiceEndpoint()); } - private Mono upsertUserInternal(CosmosUserProperties userProperties) { - return getDocClientWrapper().upsertUser(this.getLink(), ModelBridgeInternal.getV2User(userProperties), null) - .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); - } - private Mono replaceThroughputInternal(ThroughputProperties throughputProperties, Context context){ String spanName = "replaceThroughput." + this.getId(); Context nestedContext = context.addData(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.COSMOS_CALL_DEPTH_VAL); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java index ead21b5b8cb5..5628f5a48a6b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java @@ -141,42 +141,31 @@ private Mono replaceInternal(CosmosPermissionPropertie Context context) { String spanName = "replacePermission." + cosmosUser.getId(); - Mono responseMono = replaceInternal(permissionProperties, options); - return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, - cosmosUser.getDatabase().getId(), - cosmosUser.getDatabase().getClient().getServiceEndpoint()); - } - - private Mono replaceInternal(CosmosPermissionProperties permissionProperties, - CosmosPermissionRequestOptions options) { CosmosAsyncDatabase databaseContext = cosmosUser.getDatabase(); - return cosmosUser.getDatabase() + Mono responseMono = cosmosUser.getDatabase() .getDocClientWrapper() .replacePermission(ModelBridgeInternal.getPermission(permissionProperties, databaseContext.getId()), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosPermissionResponse(response)) .single(); - } - - private Mono deleteInternal(CosmosPermissionRequestOptions options, - Context context) { - - String spanName = "deletePermission." + cosmosUser.getId(); - Mono responseMono = deleteInternal(options); return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosUser.getDatabase().getId(), cosmosUser.getDatabase().getClient().getServiceEndpoint()); } - private Mono deleteInternal(CosmosPermissionRequestOptions options) { + private Mono deleteInternal(CosmosPermissionRequestOptions options, + Context context) { - return cosmosUser.getDatabase() + String spanName = "deletePermission." + cosmosUser.getId(); + Mono responseMono = cosmosUser.getDatabase() .getDocClientWrapper() .deletePermission(getLink(), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosPermissionResponse(response)) .single(); + return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, + spanName, + cosmosUser.getDatabase().getId(), + cosmosUser.getDatabase().getClient().getServiceEndpoint()); } - } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java index a79396f27195..788aa4befa9e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java @@ -176,28 +176,14 @@ String getLink() { private Mono readInternal(CosmosStoredProcedureRequestOptions options, Context context) { - String spanName = "readStoredProcedure." + cosmosContainer.getId(); - Mono responseMono = readInternal(options); - return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, - spanName, - cosmosContainer.getDatabase().getId(), - cosmosContainer.getDatabase().getClient().getServiceEndpoint()); - } - - private Mono readInternal(CosmosStoredProcedureRequestOptions options) { if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } - return cosmosContainer.getDatabase().getDocClientWrapper().readStoredProcedure(getLink(), + + String spanName = "readStoredProcedure." + cosmosContainer.getId(); + Mono responseMono = cosmosContainer.getDatabase().getDocClientWrapper().readStoredProcedure(getLink(), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosStoredProcedureResponse(response)).single(); - } - - private Mono deleteInternal(CosmosStoredProcedureRequestOptions options, - Context context) { - String spanName = "deleteStoredProcedure." + cosmosContainer.getId(); - Mono responseMono = deleteInternal(options); return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, @@ -205,59 +191,56 @@ private Mono deleteInternal(CosmosStoredProcedure cosmosContainer.getDatabase().getClient().getServiceEndpoint()); } - private Mono deleteInternal(CosmosStoredProcedureRequestOptions options) { + private Mono deleteInternal(CosmosStoredProcedureRequestOptions options, + Context context) { if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } - return cosmosContainer.getDatabase() + String spanName = "deleteStoredProcedure." + cosmosContainer.getId(); + Mono responseMono = cosmosContainer.getDatabase() .getDocClientWrapper() .deleteStoredProcedure(getLink(), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosStoredProcedureResponse(response)) .single(); + return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, + context, + spanName, + cosmosContainer.getDatabase().getId(), + cosmosContainer.getDatabase().getClient().getServiceEndpoint()); } private Mono executeInternal(List procedureParams, CosmosStoredProcedureRequestOptions options, Context context) { - String spanName = "executeStoredProcedure." + cosmosContainer.getId(); - Mono responseMono = executeInternal(procedureParams, options); - return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); - } - - private Mono executeInternal(List procedureParams, - CosmosStoredProcedureRequestOptions options) { if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } - return cosmosContainer.getDatabase() + String spanName = "executeStoredProcedure." + cosmosContainer.getId(); + Mono responseMono = cosmosContainer.getDatabase() .getDocClientWrapper() .executeStoredProcedure(getLink(), ModelBridgeInternal.toRequestOptions(options), procedureParams) .map(response -> ModelBridgeInternal.createCosmosStoredProcedureResponse(response)) .single(); + return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); } private Mono replaceInternal(CosmosStoredProcedureProperties storedProcedureSettings, CosmosStoredProcedureRequestOptions options, Context context) { - String spanName = "replaceStoredProcedure." + cosmosContainer.getId(); - Mono responseMono = replaceInternal(storedProcedureSettings, options); - return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); - } - - private Mono replaceInternal(CosmosStoredProcedureProperties storedProcedureSettings, - CosmosStoredProcedureRequestOptions options) { if (options == null) { options = new CosmosStoredProcedureRequestOptions(); } - return cosmosContainer.getDatabase() + String spanName = "replaceStoredProcedure." + cosmosContainer.getId(); + Mono responseMono = cosmosContainer.getDatabase() .getDocClientWrapper() .replaceStoredProcedure(new StoredProcedure(ModelBridgeInternal.toJsonFromJsonSerializable( ModelBridgeInternal.getResource(storedProcedureSettings))), ModelBridgeInternal.toRequestOptions(options)) .map(response -> ModelBridgeInternal.createCosmosStoredProcedureResponse(response)) .single(); + return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java index e3d627b49542..cba93102874d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java @@ -105,25 +105,11 @@ String getLink() { private Mono readInternal(Context context) { String spanName = "readTrigger." + container.getId(); - Mono responseMono = readInternal(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, - spanName, - container.getDatabase().getId(), - container.getDatabase().getClient().getServiceEndpoint()); - } - - private Mono readInternal() { - return container.getDatabase() + Mono responseMono = container.getDatabase() .getDocClientWrapper() .readTrigger(getLink(), null) .map(response -> ModelBridgeInternal.createCosmosTriggerResponse(response)) .single(); - } - - private Mono replaceInternal(CosmosTriggerProperties triggerSettings, Context context) { - String spanName = "replaceTrigger." + container.getId(); - Mono responseMono = replaceInternal(triggerSettings); return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, @@ -131,18 +117,14 @@ private Mono replaceInternal(CosmosTriggerProperties trig container.getDatabase().getClient().getServiceEndpoint()); } - private Mono replaceInternal(CosmosTriggerProperties triggerSettings) { - return container.getDatabase() + private Mono replaceInternal(CosmosTriggerProperties triggerSettings, Context context) { + String spanName = "replaceTrigger." + container.getId(); + Mono responseMono = container.getDatabase() .getDocClientWrapper() .replaceTrigger(new Trigger(ModelBridgeInternal.toJsonFromJsonSerializable( ModelBridgeInternal.getResource(triggerSettings))), null) .map(response -> ModelBridgeInternal.createCosmosTriggerResponse(response)) .single(); - } - - private Mono deleteInternal(Context context) { - String spanName = "deleteTrigger." + container.getId(); - Mono responseMono = deleteInternal(); return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, @@ -150,11 +132,17 @@ private Mono deleteInternal(Context context) { container.getDatabase().getClient().getServiceEndpoint()); } - private Mono deleteInternal() { - return container.getDatabase() + private Mono deleteInternal(Context context) { + String spanName = "deleteTrigger." + container.getId(); + Mono responseMono = container.getDatabase() .getDocClientWrapper() .deleteTrigger(getLink(), null) .map(response -> ModelBridgeInternal.createCosmosTriggerResponse(response)) .single(); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, + context, + spanName, + container.getDatabase().getId(), + container.getDatabase().getClient().getServiceEndpoint()); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java index 9a449665540f..ea879a5ff899 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java @@ -242,88 +242,64 @@ CosmosAsyncDatabase getDatabase() { private Mono readInternal(Context context) { String spanName = "readUser." + getId(); - Mono responseMono = readInternal(); + Mono responseMono = this.database.getDocClientWrapper() + .readUser(getLink(), null) + .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } - private Mono readInternal() { - return this.database.getDocClientWrapper() - .readUser(getLink(), null) - .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); - } - private Mono replaceInternal(CosmosUserProperties userSettings, Context context) { String spanName = "replaceUser." + getId(); - Mono responseMono = replaceInternal(userSettings); + Mono responseMono = this.database.getDocClientWrapper() + .replaceUser(ModelBridgeInternal.getV2User(userSettings), null) + .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } - private Mono replaceInternal(CosmosUserProperties userSettings) { - return this.database.getDocClientWrapper() - .replaceUser(ModelBridgeInternal.getV2User(userSettings), null) - .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); - } - private Mono deleteInternal(Context context) { String spanName = "deleteUser." + getId(); - Mono responseMono = deleteInternal(); + Mono responseMono = this.database.getDocClientWrapper() + .deleteUser(getLink(), null) + .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } - private Mono deleteInternal() { - return this.database.getDocClientWrapper() - .deleteUser(getLink(), null) - .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); - } - private Mono createPermissionInternal( Permission permission, CosmosPermissionRequestOptions options, Context context) { String spanName = "createPermission." + getId(); - Mono responseMono = createPermissionInternal(permission, options); + Mono responseMono = database.getDocClientWrapper() + .createPermission(getLink(), permission, ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosPermissionResponse(response)) + .single(); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } - private Mono createPermissionInternal( - Permission permission, - CosmosPermissionRequestOptions options) { - return database.getDocClientWrapper() - .createPermission(getLink(), permission, ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosPermissionResponse(response)) - .single(); - } - private Mono upsertPermissionInternal( Permission permission, CosmosPermissionRequestOptions options, Context context) { String spanName = "upsertPermission." + getId(); - Mono responseMono = upsertPermissionInternal(permission, options); + Mono responseMono = database.getDocClientWrapper() + .upsertPermission(getLink(), permission, ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosPermissionResponse(response)) + .single(); return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, database.getId(), database.getClient().getServiceEndpoint()); } - - private Mono upsertPermissionInternal( - Permission permission, - CosmosPermissionRequestOptions options) { - return database.getDocClientWrapper() - .upsertPermission(getLink(), permission, ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosPermissionResponse(response)) - .single(); - } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java index f0a35931713a..457f912bc497 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java @@ -110,8 +110,9 @@ String getLink() { } private Mono readInternal(Context context) { - String spanName = "readUDF." + container.getId(); - Mono responseMono = readInternal(); + String spanName = "readUserDefinedFunction." + container.getId(); + Mono responseMono = container.getDatabase().getDocClientWrapper().readUserDefinedFunction(getLink(), null) + .map(response -> ModelBridgeInternal.createCosmosUserDefinedFunctionResponse(response)).single(); return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, @@ -119,34 +120,15 @@ private Mono readInternal(Context context) { container.getDatabase().getClient().getServiceEndpoint()); } - private Mono readInternal() { - return container.getDatabase().getDocClientWrapper().readUserDefinedFunction(getLink(), null) - .map(response -> ModelBridgeInternal.createCosmosUserDefinedFunctionResponse(response)).single(); - } - private Mono replaceInternal(CosmosUserDefinedFunctionProperties udfSettings, Context context) { - String spanName = "replaceUDF." + container.getId(); - Mono responseMono = replaceInternal(udfSettings); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, - spanName, - container.getDatabase().getId(), - container.getDatabase().getClient().getServiceEndpoint()); - } - - private Mono replaceInternal(CosmosUserDefinedFunctionProperties udfSettings) { - return container.getDatabase() + String spanName = "replaceUserDefinedFunction." + container.getId(); + Mono responseMono = container.getDatabase() .getDocClientWrapper() .replaceUserDefinedFunction(new UserDefinedFunction(ModelBridgeInternal.toJsonFromJsonSerializable( ModelBridgeInternal.getResource(udfSettings))), null) .map(response -> ModelBridgeInternal.createCosmosUserDefinedFunctionResponse(response)) .single(); - } - - private Mono deleteInternal(Context context) { - String spanName = "deleteUDF." + container.getId(); - Mono responseMono = deleteInternal(); return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, @@ -154,11 +136,17 @@ private Mono deleteInternal(Context context) container.getDatabase().getClient().getServiceEndpoint()); } - private Mono deleteInternal() { - return container.getDatabase() + private Mono deleteInternal(Context context) { + String spanName = "deleteUserDefinedFunction." + container.getId(); + Mono responseMono = container.getDatabase() .getDocClientWrapper() .deleteUserDefinedFunction(this.getLink(), null) .map(response -> ModelBridgeInternal.createCosmosUserDefinedFunctionResponse(response)) .single(); + return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, + context, + spanName, + container.getDatabase().getId(), + container.getDatabase().getClient().getServiceEndpoint()); } } diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java index 29ef071f0867..1f3cd8b49ea3 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java @@ -4,6 +4,7 @@ import com.azure.core.util.Context; import com.azure.core.util.tracing.Tracer; +import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.InternalObjectNode; import com.azure.cosmos.implementation.LifeCycleUtils; import com.azure.cosmos.implementation.TestConfigurations; @@ -14,7 +15,10 @@ import com.azure.cosmos.models.CosmosStoredProcedureProperties; import com.azure.cosmos.models.CosmosTriggerProperties; import com.azure.cosmos.models.CosmosUserDefinedFunctionProperties; +import com.azure.cosmos.models.CosmosUserProperties; +import com.azure.cosmos.models.FeedResponse; import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.models.ThroughputProperties; import com.azure.cosmos.models.TriggerOperation; import com.azure.cosmos.models.TriggerType; import com.azure.cosmos.rx.TestSuiteBase; @@ -30,6 +34,9 @@ import java.util.List; import java.util.UUID; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.fail; + public class CosmosTracerTest extends TestSuiteBase { private static final String ITEM_ID = "tracerDoc"; CosmosAsyncClient client; @@ -53,30 +60,36 @@ public void cosmosAsyncClient() { Tracer mockTracer = Mockito.mock(Tracer.class); TracerProvider tracerProvider = Mockito.spy(new TracerProvider(getMockTracer(mockTracer))); ReflectionUtils.setTracerProvider(client, tracerProvider); + int traceApiCounter = 1; TracerProviderCapture tracerProviderCapture = new TracerProviderCapture(); Mockito.doAnswer(tracerProviderCapture).when(tracerProvider).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); - client.createDatabaseIfNotExists(cosmosAsyncDatabase.getId()).block(); + client.createDatabaseIfNotExists(cosmosAsyncDatabase.getId(), ThroughputProperties.createManualThroughput(5000)).block(); Context context = tracerProviderCapture.getResult(); - Mockito.verify(tracerProvider, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "createDatabaseIfNotExists." + cosmosAsyncDatabase.getId(), context, - cosmosAsyncDatabase.getId(), 1); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; client.readAllDatabases(new CosmosQueryRequestOptions()).byPage().single().block(); - Mockito.verify(tracerProvider, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); - verifyTracerAttributes(mockTracer, "readAllDatabases", context, null, 2); + verifyTracerAttributes(mockTracer, "readAllDatabases", context, null, traceApiCounter, null); + traceApiCounter++; String query = "select * from c where c.id = '" + cosmosAsyncDatabase.getId() + "'"; client.queryDatabases(query, new CosmosQueryRequestOptions()).byPage().single().block(); - Mockito.verify(tracerProvider, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); - verifyTracerAttributes(mockTracer, "queryDatabases", context, null, 3); + verifyTracerAttributes(mockTracer, "queryDatabases", context, null, traceApiCounter, null); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) @@ -88,37 +101,46 @@ public void cosmosAsyncDatabase() { Mockito.doAnswer(tracerProviderCapture).when(tracerProvider).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + int traceApiCounter = 1; cosmosAsyncDatabase.createContainerIfNotExists(cosmosAsyncContainer.getId(), "/pk", 5000).block(); Context context = tracerProviderCapture.getResult(); - Mockito.verify(tracerProvider, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "createContainerIfNotExists." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 1); - - try { - cosmosAsyncDatabase.readThroughput().block(); - } catch (CosmosException ex) { - //do nothing - } - - Mockito.verify(tracerProvider, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), - Matchers.anyString(), Matchers.any(Context.class)); - verifyTracerAttributes(mockTracer, "readThroughput." + cosmosAsyncDatabase.getId(), context, - cosmosAsyncDatabase.getId(), 2); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; cosmosAsyncDatabase.readAllUsers().byPage().single().block(); - Mockito.verify(tracerProvider, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "readAllUsers." + cosmosAsyncDatabase.getId(), context, - cosmosAsyncDatabase.getId(), 3); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; cosmosAsyncDatabase.readAllContainers().byPage().single().block(); - Mockito.verify(tracerProvider, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "readAllContainers." + cosmosAsyncDatabase.getId(), context, - cosmosAsyncDatabase.getId(), 4); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; + + String errorType = null; + try { + cosmosAsyncDatabase.readThroughput().block(); + } catch (CosmosException ex) { + errorType = ex.getClass().getName(); + } + + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "readThroughput." + cosmosAsyncDatabase.getId(), context, + cosmosAsyncDatabase.getId(), traceApiCounter, errorType); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) @@ -130,64 +152,81 @@ public void cosmosAsyncContainer() { Mockito.doAnswer(tracerProviderCapture).when(tracerProvider).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + int traceApiCounter = 1; cosmosAsyncContainer.read().block(); Context context = tracerProviderCapture.getResult(); - Mockito.verify(tracerProvider, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "readContainer." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 1); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; try { cosmosAsyncContainer.readThroughput().block(); } catch (CosmosException ex) { //do nothing } - Mockito.verify(tracerProvider, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "readThroughput." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 2); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; InternalObjectNode item = new InternalObjectNode(); item.setId(ITEM_ID); cosmosAsyncContainer.createItem(item).block(); - Mockito.verify(tracerProvider, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "createItem." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 3); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; cosmosAsyncContainer.upsertItem(item, new CosmosItemRequestOptions()).block(); - Mockito.verify(tracerProvider, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "upsertItem." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 4); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; - cosmosAsyncContainer.readItem(ITEM_ID, PartitionKey.NONE, - InternalObjectNode.class).block(); - Mockito.verify(tracerProvider, Mockito.times(5)).startSpan(Matchers.anyString(), Matchers.anyString(), + InternalObjectNode node = cosmosAsyncContainer.readItem(ITEM_ID, PartitionKey.NONE, + InternalObjectNode.class).block().getItem(); + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "readItem." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 5); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; cosmosAsyncContainer.deleteItem(ITEM_ID, PartitionKey.NONE).block(); - Mockito.verify(tracerProvider, Mockito.times(6)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "deleteItem." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 6); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; cosmosAsyncContainer.readAllItems(new CosmosQueryRequestOptions(), CosmosItemRequestOptions.class).byPage().single().block(); - Mockito.verify(tracerProvider, Mockito.times(7)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "readAllItems." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 7); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; String query = "select * from c where c.id = '" + ITEM_ID + "'"; cosmosAsyncContainer.queryItems(query, new CosmosQueryRequestOptions(), CosmosItemRequestOptions.class).byPage().single().block(); - Mockito.verify(tracerProvider, Mockito.times(8)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "queryItems." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 8); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; } @Test(groups = {"emulator"}, timeOut = TIMEOUT) @@ -199,117 +238,191 @@ public void cosmosAsyncScripts() { Mockito.doAnswer(tracerProviderCapture).when(tracerProvider).startSpan(Matchers.anyString(), Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + int traceApiCounter = 1; cosmosAsyncContainer.getScripts().readAllStoredProcedures(new CosmosQueryRequestOptions()).byPage().single().block(); Context context = tracerProviderCapture.getResult(); - Mockito.verify(tracerProvider, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "readAllStoredProcedures." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 1); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; cosmosAsyncContainer.getScripts().readAllTriggers(new CosmosQueryRequestOptions()).byPage().single().block(); - Mockito.verify(tracerProvider, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "readAllTriggers." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 2); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; cosmosAsyncContainer.getScripts().readAllUserDefinedFunctions(new CosmosQueryRequestOptions()).byPage().single().block(); - Mockito.verify(tracerProvider, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "readAllUserDefinedFunctions." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 3); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; CosmosUserDefinedFunctionProperties cosmosUserDefinedFunctionProperties = getCosmosUserDefinedFunctionProperties(); CosmosUserDefinedFunctionProperties resultUdf = cosmosAsyncContainer.getScripts().createUserDefinedFunction(cosmosUserDefinedFunctionProperties).block().getProperties(); - Mockito.verify(tracerProvider, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "createUserDefinedFunction." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 4); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; cosmosAsyncContainer.getScripts().getUserDefinedFunction(cosmosUserDefinedFunctionProperties.getId()).read().block(); - Mockito.verify(tracerProvider, Mockito.times(5)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); - verifyTracerAttributes(mockTracer, "readUDF." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 5); + verifyTracerAttributes(mockTracer, "readUserDefinedFunction." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; cosmosUserDefinedFunctionProperties.setBody("function() {var x = 15;}"); cosmosAsyncContainer.getScripts().getUserDefinedFunction(resultUdf.getId()).replace(resultUdf).block(); - Mockito.verify(tracerProvider, Mockito.times(6)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); - verifyTracerAttributes(mockTracer, "replaceUDF." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 6); + verifyTracerAttributes(mockTracer, "replaceUserDefinedFunction." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; cosmosAsyncContainer.getScripts().readAllUserDefinedFunctions(new CosmosQueryRequestOptions()).byPage().single().block(); - Mockito.verify(tracerProvider, Mockito.times(7)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + traceApiCounter++; cosmosAsyncContainer.getScripts().getUserDefinedFunction(cosmosUserDefinedFunctionProperties.getId()).delete().block(); - Mockito.verify(tracerProvider, Mockito.times(8)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); - verifyTracerAttributes(mockTracer, "deleteUDF." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 8); + verifyTracerAttributes(mockTracer, "deleteUserDefinedFunction." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; CosmosTriggerProperties cosmosTriggerProperties = getCosmosTriggerProperties(); CosmosTriggerProperties resultTrigger = cosmosAsyncContainer.getScripts().createTrigger(cosmosTriggerProperties).block().getProperties(); - Mockito.verify(tracerProvider, Mockito.times(9)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "createTrigger." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 9); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; cosmosAsyncContainer.getScripts().getTrigger(cosmosTriggerProperties.getId()).read().block(); - Mockito.verify(tracerProvider, Mockito.times(10)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "readTrigger." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 10); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; cosmosAsyncContainer.getScripts().getTrigger(cosmosTriggerProperties.getId()).replace(resultTrigger).block(); - Mockito.verify(tracerProvider, Mockito.times(11)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "replaceTrigger." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 11); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; cosmosAsyncContainer.getScripts().readAllTriggers(new CosmosQueryRequestOptions()).byPage().single().block(); - Mockito.verify(tracerProvider, Mockito.times(12)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); + traceApiCounter++; cosmosAsyncContainer.getScripts().getTrigger(cosmosTriggerProperties.getId()).delete().block(); - Mockito.verify(tracerProvider, Mockito.times(13)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "deleteTrigger." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 13); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; CosmosStoredProcedureProperties procedureProperties = getCosmosStoredProcedureProperties(); CosmosStoredProcedureProperties resultSproc = cosmosAsyncContainer.getScripts().createStoredProcedure(procedureProperties).block().getProperties(); - Mockito.verify(tracerProvider, Mockito.times(14)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "createStoredProcedure." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 14); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).read().block(); - Mockito.verify(tracerProvider, Mockito.times(15)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "readStoredProcedure." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 15); - + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).replace(resultSproc).block(); - Mockito.verify(tracerProvider, Mockito.times(16)).startSpan(Matchers.anyString(), Matchers.anyString(), + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "replaceStoredProcedure." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 16); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; cosmosAsyncContainer.getScripts().readAllStoredProcedures(new CosmosQueryRequestOptions()).byPage().single().block(); cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).delete().block(); - Mockito.verify(tracerProvider, Mockito.times(18)).startSpan(Matchers.anyString(), Matchers.anyString(), + traceApiCounter++; + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), Matchers.anyString(), Matchers.any(Context.class)); verifyTracerAttributes(mockTracer, "deleteStoredProcedure." + cosmosAsyncContainer.getId(), context, - cosmosAsyncDatabase.getId(), 18); + cosmosAsyncDatabase.getId(), traceApiCounter, null); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void tracerExceptionSpan() { + Tracer mockTracer = Mockito.mock(Tracer.class); + TracerProvider tracerProvider = Mockito.spy(new TracerProvider(getMockTracer(mockTracer))); + ReflectionUtils.setTracerProvider(client, tracerProvider); + int traceApiCounter = 1; + + TracerProviderCapture tracerProviderCapture = new TracerProviderCapture(); + Mockito.doAnswer(tracerProviderCapture).when(tracerProvider).startSpan(Matchers.anyString(), + Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); + + InternalObjectNode item = new InternalObjectNode(); + item.setId("testDoc"); + cosmosAsyncContainer.createItem(item).block(); + Context context = tracerProviderCapture.getResult(); + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "createItem." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), traceApiCounter, null); + traceApiCounter++; + + String errorType = null; + try { + PartitionKey partitionKey = new PartitionKey("wrongPk"); + cosmosAsyncContainer.readItem("testDoc", partitionKey, null, InternalObjectNode.class).block(); + fail("readItem should fail due to wrong pk"); + } catch (CosmosException ex) { + assertThat(ex.getStatusCode()).isEqualTo(HttpConstants.StatusCodes.NOTFOUND); + errorType = ex.getClass().getName(); + } + Mockito.verify(tracerProvider, Mockito.times(traceApiCounter)).startSpan(Matchers.anyString(), + Matchers.anyString(), + Matchers.anyString(), Matchers.any(Context.class)); + verifyTracerAttributes(mockTracer, "readItem." + cosmosAsyncContainer.getId(), context, + cosmosAsyncDatabase.getId(), traceApiCounter + , errorType); } @AfterClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT) @@ -345,7 +458,7 @@ private List getMockTracer(Tracer tracer) { } private void verifyTracerAttributes(Tracer mockTracer, String methodName, Context context, String databaseName, - int numberOfTimesCalledWithinTest) { + int numberOfTimesCalledWithinTest, String errorType) { if (databaseName != null) { Mockito.verify(mockTracer, Mockito.times(numberOfTimesCalledWithinTest)).setAttribute(TracerProvider.DB_INSTANCE, databaseName, context); @@ -356,7 +469,17 @@ private void verifyTracerAttributes(Tracer mockTracer, String methodName, Contex TestConfigurations.HOST, context); Mockito.verify(mockTracer, Mockito.times(1)).setAttribute(TracerProvider.DB_STATEMENT, methodName, context); - + if (errorType == null) { + Mockito.verify(mockTracer, Mockito.times(0)).setAttribute(Mockito.eq(TracerProvider.ERROR_MSG) + , Matchers.anyString(), Mockito.eq(context)); + Mockito.verify(mockTracer, Mockito.times(0)).setAttribute(Mockito.eq(TracerProvider.ERROR_TYPE) + , Matchers.anyString(), Mockito.eq(context)); + } else { + Mockito.verify(mockTracer, Mockito.times(1)).setAttribute(Mockito.eq(TracerProvider.ERROR_TYPE) + , Mockito.eq(errorType), Mockito.eq(context)); + Mockito.verify(mockTracer, Mockito.times(1)).setAttribute(Mockito.eq(TracerProvider.ERROR_MSG) + , Matchers.anyString(), Mockito.eq(context)); + } } private class TracerProviderCapture implements Answer {