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/26] 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 14909afd49ec4..c56da8757b2fa 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 e28517af0d707..1bd2dc54cab82 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 f84535a1ee775..676e1dd103dd3 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 270f6b51e2d40..e21feb07a508b 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 0000000000000..0ffc8f1a4ad2a --- /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 06c610563564a..3fe48a84f5368 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 0000000000000..5e5651827b356 --- /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/26] 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 5e5651827b356..a50072377f609 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/26] 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 1bd2dc54cab82..9197b60739321 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 f33665cdbc546..2dbe77aeabcc9 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 676e1dd103dd3..5614a9cf8ae20 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 7c066bbf3a1af..0e4e0bedd5dc6 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 1692ceaf0b213..7bf12cdaad145 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 c4407b68bfee2..eb6340bba9bae 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 0ffc8f1a4ad2a..f3381af9504f7 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 ca9f20352f3af..1527d08267b9e 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/26] 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 29c9ba150e587..c36e195c1e866 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/26] 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 fdad4fd4f0eb6..00daf223c0ab4 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 2dbe77aeabcc9..88d0114ec93e0 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 5614a9cf8ae20..0e430db5cbce9 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 0e4e0bedd5dc6..55afad42f2656 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 9efa1778e3b8f..2b84be5b9d8a0 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 456b522f26aa5..be0e6d4039eb8 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 3823f1067a0b2..27d517bd1d7e0 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/26] 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 00daf223c0ab4..ccb4a72386ecf 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 88d0114ec93e0..929857fbfeb01 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 0e430db5cbce9..2df8033ba2e2b 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 55afad42f2656..00b1578e18b49 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 2b84be5b9d8a0..102445d056261 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 be0e6d4039eb8..498d7b5d85884 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 27d517bd1d7e0..3c9eae92d9caf 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 f3381af9504f7..25bcb7b394bc8 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 cba344b67ef1d..6a926894171c8 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 0000000000000..963d05c49ce93 --- /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 9dbdc312ffcd8..5a9063584e374 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/26] 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 9c8f93ca2ced9..445feb8ede6c0 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 ccb4a72386ecf..3ec556efe5b8f 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 929857fbfeb01..29cd1edbc909e 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 2df8033ba2e2b..bcfb8415b06d5 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 00b1578e18b49..c4618f20016ff 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 102445d056261..9a1a95fed2fce 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 498d7b5d85884..5ac774388ee4d 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 a231fa5d328f7..528eb57cab2b2 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 3c9eae92d9caf..82a7a0802869e 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 eb6340bba9bae..1f0a3904db478 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 25bcb7b394bc8..691b1b522c53f 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 1527d08267b9e..388d99b61f8aa 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 5500ddcc3e81c..b29c0d76586a3 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 963d05c49ce93..7ac4460a5881e 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/26] 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 388d99b61f8aa..0ce63cca60dee 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/26] 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 82a7a0802869e..0d3ed2d6575a2 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 691b1b522c53f..376099c2d6974 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 0ce63cca60dee..998cb6d7c42c0 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 7ac4460a5881e..3bff6207d2330 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/26] 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 3ec556efe5b8f..3b1ff821d391d 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 29cd1edbc909e..2c9df7932ae9f 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 bcfb8415b06d5..68d9c5a2562e2 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 f4798285b5aa6..5075c2ffc293e 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 c4618f20016ff..e9fd16f1141c1 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 9a1a95fed2fce..39e064790fa2f 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 5ac774388ee4d..e0efc005db0d9 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 528eb57cab2b2..f51e3e61af9c0 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 0d3ed2d6575a2..6137dfb09468c 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 1f0a3904db478..28b4054180dcb 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 376099c2d6974..1733bef533ba4 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 998cb6d7c42c0..f1a8e61f60420 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 3bff6207d2330..f53375fcb012f 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/26] 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 2c9df7932ae9f..842cc54c8df33 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 e9fd16f1141c1..06f47bc15c764 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 1733bef533ba4..078c297daf29a 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/26] 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 f803ddf91dcb9..a7dd95664e075 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 68d9c5a2562e2..3dd11d2438575 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 9724f652fe5ab..398b08cf65dee 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/26] 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 a7dd95664e075..e03c64535981b 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 c02ae8fb04028..db819a4f97c2b 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 3dd11d2438575..d8ee89cb4bf17 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 5075c2ffc293e..694303473c9d0 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 06f47bc15c764..ae5249e03bfda 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 39e064790fa2f..8acc71aefd154 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 e0efc005db0d9..47ebc197cb5c2 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 f51e3e61af9c0..c9125995f1127 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 6137dfb09468c..e16ee192c4e82 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 398b08cf65dee..2790b37dc3ace 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 078c297daf29a..37a74912594b9 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/26] 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 d8ee89cb4bf17..fb09250cd6726 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 2790b37dc3ace..0435020dc3e8b 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 37a74912594b9..2661004a9e924 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/26] 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 e03c64535981b..bbe11ede3be0e 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/26] 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 2661004a9e924..2a08d5e2fb8d9 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/26] 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 763d5827f7051..af40f300d83e9 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 f53375fcb012f..5178f2099dec4 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/26] 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 2b58b0fb2d99d..e6c4d05a50d45 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/26] 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 8f20fba7d01eb..65283817fa575 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/26] 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 2fa9ff8b01c0e..a38000b6edaa7 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/26] 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 fddd4e502b990..42c87add6fd1c 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 43ae1e1c81435..6069057f99d6a 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/26] 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 fcb4f63dbf4c5..29ecd116d22c6 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 8d735b5c28495..5bf4aee03b5b4 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 a99658973234c..4abba56c2bd67 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 f00ad63fd615d..3d5474bda9f8e 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 6ed36f6b0888a..34adcd145f827 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/26] 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 59d6294cbcf79..473c45d6d25d8 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 9e7a336e782f4..96f1edaaa7ed1 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 d52028cc3ceb7..e08efdd4d8c48 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 6eb23ab98a743..6f09a2f43f61c 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 96b8060e7ada3..128937f4f7620 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 9b0f2ec70c835..7d4e222449484 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 71650e825e69e..e6d4433de9891 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 eef66c14e3ffe..666ddad0a6468 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 43442f117d22d..f2ea9c442a9aa 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 ff524c1783e70..45b6803fdaac8 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/26] 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 e08efdd4d8c48..81055b0dc932e 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 6f09a2f43f61c..1da2affd09160 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 0c8a6716e4325..344cef63c391f 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/26] 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 344cef63c391f..677e87db4e2f6 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/26] 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 588a958e95b66..ae8dee16f2b31 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