From 29b0fc4e50743e044dd3ae4295b862e738dbd5cf Mon Sep 17 00:00:00 2001 From: Kushagra Thapar Date: Tue, 16 Jun 2020 12:53:34 -0700 Subject: [PATCH] Revert "Tracer support in cosmosdb (#10265)" This reverts commit 55c23c6c4fe118fca384f0fd10a648569558659f. --- sdk/cosmos/azure-cosmos/pom.xml | 13 +- .../com/azure/cosmos/CosmosAsyncClient.java | 164 ++----- .../com/azure/cosmos/CosmosAsyncConflict.java | 45 +- .../azure/cosmos/CosmosAsyncContainer.java | 437 ++++------------- .../com/azure/cosmos/CosmosAsyncDatabase.java | 450 +++++------------- .../azure/cosmos/CosmosAsyncPermission.java | 104 +--- .../com/azure/cosmos/CosmosAsyncScripts.java | 215 ++------- .../cosmos/CosmosAsyncStoredProcedure.java | 138 ++---- .../com/azure/cosmos/CosmosAsyncTrigger.java | 89 +--- .../com/azure/cosmos/CosmosAsyncUser.java | 151 +----- .../CosmosAsyncUserDefinedFunction.java | 84 +--- .../CosmosPagedFluxOptions.java | 46 -- .../cosmos/implementation/HttpConstants.java | 1 - .../cosmos/implementation/TracerProvider.java | 157 ------ .../azure/cosmos/util/CosmosPagedFlux.java | 46 +- .../azure/cosmos/util/UtilBridgeInternal.java | 5 +- .../src/main/java/module-info.java | 1 - .../com/azure/cosmos/CosmosTracerTest.java | 245 ---------- .../directconnectivity/ReflectionUtils.java | 5 - .../rx/ReadFeedExceptionHandlingTest.java | 10 +- .../com/azure/cosmos/rx/TestSuiteBase.java | 8 +- 21 files changed, 427 insertions(+), 1987 deletions(-) delete mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java delete mode 100644 sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java diff --git a/sdk/cosmos/azure-cosmos/pom.xml b/sdk/cosmos/azure-cosmos/pom.xml index a7af309071aa..23d2adad218c 100644 --- a/sdk/cosmos/azure-cosmos/pom.xml +++ b/sdk/cosmos/azure-cosmos/pom.xml @@ -128,18 +128,7 @@ Licensed under the MIT License. - - com.azure - azure-core-tracing-opentelemetry - test - 1.0.0-beta.5 - - - io.netty - netty-tcnative-boringssl-static - - - + com.fasterxml.jackson.module jackson-module-afterburner diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index 473c45d6d25d..bbe989db8e80 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,19 +4,16 @@ import com.azure.core.annotation.ServiceClient; import com.azure.core.credential.AzureKeyCredential; -import com.azure.core.util.Context; -import com.azure.core.util.tracing.Tracer; import com.azure.cosmos.implementation.AsyncDocumentClient; import com.azure.cosmos.implementation.Configs; import com.azure.cosmos.implementation.ConnectionPolicy; import com.azure.cosmos.implementation.CosmosAuthorizationTokenResolver; import com.azure.cosmos.implementation.Database; import com.azure.cosmos.implementation.HttpConstants; -import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdMetrics; +import com.azure.cosmos.models.CosmosDatabaseResponse; import com.azure.cosmos.models.CosmosDatabaseProperties; import com.azure.cosmos.models.CosmosDatabaseRequestOptions; -import com.azure.cosmos.models.CosmosDatabaseResponse; import com.azure.cosmos.models.CosmosPermissionProperties; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.ModelBridgeInternal; @@ -30,9 +27,7 @@ import java.io.Closeable; import java.util.List; -import java.util.ServiceLoader; -import static com.azure.core.util.FluxUtil.withContext; import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; /** @@ -56,7 +51,6 @@ public final class CosmosAsyncClient implements Closeable { private final AzureKeyCredential credential; private final boolean sessionCapturingOverride; private final boolean enableTransportClientSharing; - private final TracerProvider tracerProvider; private final boolean contentResponseOnWriteEnabled; CosmosAsyncClient(CosmosClientBuilder builder) { @@ -71,7 +65,6 @@ public final class CosmosAsyncClient implements Closeable { this.sessionCapturingOverride = builder.isSessionCapturingOverrideEnabled(); this.enableTransportClientSharing = builder.isConnectionSharingAcrossClientsEnabled(); this.contentResponseOnWriteEnabled = builder.isContentResponseOnWriteEnabled(); - this.tracerProvider = new TracerProvider(ServiceLoader.load(Tracer.class)); this.asyncDocumentClient = new AsyncDocumentClient.Builder() .withServiceEndpoint(this.serviceEndpoint) .withMasterKeyOrResourceToken(this.keyOrResourceToken) @@ -201,14 +194,8 @@ boolean isContentResponseOnWriteEnabled() { * @return a {@link Mono} containing the cosmos database response with the created or existing database or * an error. */ - public Mono createDatabaseIfNotExists(CosmosDatabaseProperties databaseProperties) { - if(!getTracerProvider().isEnabled()) { - CosmosAsyncDatabase database = getDatabase(databaseProperties.getId()); - return createDatabaseIfNotExistsInternal(database.read(), database, null, null); - } - - return withContext(context -> createDatabaseIfNotExistsInternal(getDatabase(databaseProperties.getId()), - null, context)); + Mono createDatabaseIfNotExists(CosmosDatabaseProperties databaseProperties) { + return createDatabaseIfNotExistsInternal(getDatabase(databaseProperties.getId())); } /** @@ -222,12 +209,21 @@ public Mono createDatabaseIfNotExists(CosmosDatabaseProp * an error. */ public Mono createDatabaseIfNotExists(String id) { - if(!getTracerProvider().isEnabled()) { - CosmosAsyncDatabase database = getDatabase(id); - return createDatabaseIfNotExistsInternal(database.read(), database, null,null); - } + return createDatabaseIfNotExistsInternal(getDatabase(id)); + } - return withContext(context -> createDatabaseIfNotExistsInternal(getDatabase(id), null, context)); + private Mono createDatabaseIfNotExistsInternal(CosmosAsyncDatabase database) { + return database.read().onErrorResume(exception -> { + final Throwable unwrappedException = Exceptions.unwrap(exception); + if (unwrappedException instanceof CosmosException) { + final CosmosException cosmosException = (CosmosException) unwrappedException; + if (cosmosException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) { + return createDatabase(new CosmosDatabaseProperties(database.getId()), + new CosmosDatabaseRequestOptions()); + } + } + return Mono.error(unwrappedException); + }); } /** @@ -244,13 +240,19 @@ public Mono createDatabaseIfNotExists(String id) { * @return the mono. */ public Mono createDatabaseIfNotExists(String id, ThroughputProperties throughputProperties) { - if(!getTracerProvider().isEnabled()) { - CosmosAsyncDatabase database = getDatabase(id); - return createDatabaseIfNotExistsInternal(database.read(), database, throughputProperties, null); - } - - return withContext(context -> createDatabaseIfNotExistsInternal(getDatabase(id), - throughputProperties, context)); + 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); + }); } /** @@ -272,12 +274,9 @@ public Mono createDatabase(CosmosDatabaseProperties data } Database wrappedDatabase = new Database(); wrappedDatabase.setId(databaseProperties.getId()); - if(!getTracerProvider().isEnabled()) { - return createDatabaseInternal(wrappedDatabase, options); - } - - final CosmosDatabaseRequestOptions requestOptions = options; - return withContext(context -> createDatabaseInternal(wrappedDatabase, requestOptions, context)); + return asyncDocumentClient.createDatabase(wrappedDatabase, ModelBridgeInternal.toRequestOptions(options)) + .map(databaseResourceResponse -> ModelBridgeInternal.createCosmosDatabaseResponse(databaseResourceResponse)) + .single(); } /** @@ -332,13 +331,9 @@ public Mono createDatabase(CosmosDatabaseProperties data ModelBridgeInternal.setThroughputProperties(options, throughputProperties); Database wrappedDatabase = new Database(); wrappedDatabase.setId(databaseProperties.getId()); - if (!getTracerProvider().isEnabled()) { - return createDatabaseInternal(wrappedDatabase, options); - } - - - final CosmosDatabaseRequestOptions requestOptions = options; - return withContext(context -> createDatabaseInternal(wrappedDatabase, requestOptions, context)); + return asyncDocumentClient.createDatabase(wrappedDatabase, ModelBridgeInternal.toRequestOptions(options)) + .map(databaseResourceResponse -> ModelBridgeInternal.createCosmosDatabaseResponse(databaseResourceResponse)) + .single(); } /** @@ -402,15 +397,13 @@ public Mono createDatabase(String id, ThroughputProperti */ CosmosPagedFlux readAllDatabases(CosmosQueryRequestOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "readAllDatabases"; - pagedFluxOptions.setTracerInformation(this.tracerProvider, spanName, this.serviceEndpoint, null); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().readDatabases(options) .map(response -> BridgeInternal.createFeedResponse( ModelBridgeInternal.getCosmosDatabasePropertiesFromV2Results(response.getResults()), response.getResponseHeaders())); - }, this.tracerProvider.isEnabled()); + }); } /** @@ -439,7 +432,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, CosmosQueryRequestOptions options) { - return queryDatabasesInternal(new SqlQuerySpec(query), options); + return queryDatabases(new SqlQuerySpec(query), options); } /** @@ -454,7 +447,13 @@ public CosmosPagedFlux queryDatabases(String query, Co * @return a {@link CosmosPagedFlux} containing one or several feed response pages of read databases or an error. */ public CosmosPagedFlux queryDatabases(SqlQuerySpec querySpec, CosmosQueryRequestOptions options) { - return queryDatabasesInternal(querySpec, options); + return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); + return getDocClientWrapper().queryDatabases(querySpec, options) + .map(response -> BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosDatabasePropertiesFromV2Results(response.getResults()), + response.getResponseHeaders())); + }); } /** @@ -474,77 +473,4 @@ public CosmosAsyncDatabase getDatabase(String id) { public void close() { asyncDocumentClient.close(); } - - TracerProvider getTracerProvider(){ - return this.tracerProvider; - } - - private CosmosPagedFlux queryDatabasesInternal(SqlQuerySpec querySpec, CosmosQueryRequestOptions options){ - return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "queryDatabases"; - pagedFluxOptions.setTracerInformation(this.tracerProvider, spanName, this.serviceEndpoint, null); - setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); - return getDocClientWrapper().queryDatabases(querySpec, options) - .map(response -> BridgeInternal.createFeedResponse( - ModelBridgeInternal.getCosmosDatabasePropertiesFromV2Results(response.getResults()), - response.getResponseHeaders())); - }, this.tracerProvider.isEnabled()); - } - - - private Mono createDatabaseIfNotExistsInternal(CosmosAsyncDatabase database, - ThroughputProperties throughputProperties, Context context) { - String spanName = "createDatabaseIfNotExists." + database.getId(); - Context nestedContext = context.addData(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.COSMOS_CALL_DEPTH_VAL); - Mono 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, ThroughputProperties throughputProperties, Context context) { - return responseMono.onErrorResume(exception -> { - final Throwable unwrappedException = Exceptions.unwrap(exception); - if (unwrappedException instanceof CosmosException) { - final CosmosException cosmosException = (CosmosException) unwrappedException; - if (cosmosException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) { - CosmosDatabaseRequestOptions requestOptions = new CosmosDatabaseRequestOptions(); - if(throughputProperties != null) { - ModelBridgeInternal.setThroughputProperties(requestOptions, throughputProperties); - } - - if (context != null) { - Database wrappedDatabase = new Database(); - wrappedDatabase.setId(database.getId()); - return createDatabaseInternal(wrappedDatabase, - requestOptions, context); - } - - return createDatabase(new CosmosDatabaseProperties(database.getId()), - requestOptions); - } - } - return Mono.error(unwrappedException); - }); - } - - - private Mono createDatabaseInternal(Database database, CosmosDatabaseRequestOptions options, - Context context) { - String spanName = "createDatabase." + database.getId(); - 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.createCosmosDatabaseResponse(databaseResourceResponse)) - .single(); - } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncConflict.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncConflict.java index 96f1edaaa7ed..9e7a336e782f 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,7 +2,6 @@ // 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; @@ -10,8 +9,6 @@ import com.azure.cosmos.models.ModelBridgeInternal; import reactor.core.publisher.Mono; -import static com.azure.core.util.FluxUtil.withContext; - /** * Read and delete conflicts */ @@ -67,11 +64,9 @@ public Mono read(CosmosConflictRequestOptions options) { options = new CosmosConflictRequestOptions(); } RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); - if (!this.container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return readInternal(requestOptions); - } + return this.container.getDatabase().getDocClientWrapper().readConflict(getLink(), requestOptions) + .map(response -> ModelBridgeInternal.createCosmosConflictResponse(response)).single(); - return withContext(context -> readInternal(requestOptions, context)); } /** @@ -90,11 +85,8 @@ public Mono delete(CosmosConflictRequestOptions options) options = new CosmosConflictRequestOptions(); } RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); - if (!this.container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return deleteInternal(requestOptions); - } - - return withContext(context -> deleteInternal(requestOptions, context)); + return this.container.getDatabase().getDocClientWrapper().deleteConflict(getLink(), requestOptions) + .map(response -> ModelBridgeInternal.createCosmosConflictResponse(response)).single(); } String getURIPathSegment() { @@ -114,33 +106,4 @@ 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 72961b5426e0..7e65650ff500 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,7 +2,6 @@ // Licensed under the MIT License. package com.azure.cosmos; -import com.azure.core.util.Context; import com.azure.cosmos.implementation.Constants; import com.azure.cosmos.implementation.CosmosItemProperties; import com.azure.cosmos.implementation.Document; @@ -10,15 +9,14 @@ 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.implementation.Utils; import com.azure.cosmos.implementation.query.QueryInfo; +import com.azure.cosmos.models.CosmosContainerResponse; +import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosConflictProperties; import com.azure.cosmos.models.CosmosContainerProperties; import com.azure.cosmos.models.CosmosContainerRequestOptions; -import com.azure.cosmos.models.CosmosContainerResponse; import com.azure.cosmos.models.CosmosItemRequestOptions; -import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.FeedResponse; import com.azure.cosmos.models.ModelBridgeInternal; @@ -34,7 +32,6 @@ import java.util.List; import java.util.stream.Collectors; -import static com.azure.core.util.FluxUtil.withContext; import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; /** @@ -46,40 +43,12 @@ public class CosmosAsyncContainer { private final CosmosAsyncDatabase database; private final String id; private final String link; - private final String replaceContainerSpanName; - private final String deleteContainerSpanName; - private final String replaceThroughputSpanName; - private final String readThroughputSpanName; - private final String readContainerSpanName; - private final String readItemSpanName; - private final String upsertItemSpanName; - private final String deleteItemSpanName; - private final String replaceItemSpanName; - private final String createItemSpanName; - private final String readAllItemsSpanName; - private final String queryItemsSpanName; - private final String readAllConflictsSpanName; - private final String queryConflictsSpanName; private CosmosAsyncScripts scripts; CosmosAsyncContainer(String id, CosmosAsyncDatabase database) { this.id = id; this.database = database; this.link = getParentLink() + "/" + getURIPathSegment() + "/" + getId(); - this.replaceContainerSpanName = "replaceContainer." + this.id; - this.deleteContainerSpanName = "deleteContainer." + this.id; - this.replaceThroughputSpanName = "replaceThroughput." + this.id; - this.readThroughputSpanName = "readThroughput." + this.id; - this.readContainerSpanName = "readContainer." + this.id; - this.readItemSpanName = "readItem." + this.id; - this.upsertItemSpanName = "upsertItem." + this.id; - this.deleteItemSpanName = "deleteItem." + this.id; - this.replaceItemSpanName = "replaceItem." + this.id; - this.createItemSpanName = "createItem." + this.id; - this.readAllItemsSpanName = "readAllItems." + this.id; - this.queryItemsSpanName = "queryItems." + this.id; - this.readAllConflictsSpanName = "readAllConflicts." + this.id; - this.queryConflictsSpanName = "queryConflicts." + this.id; } /** @@ -120,13 +89,8 @@ public Mono read(CosmosContainerRequestOptions options) if (options == null) { options = new CosmosContainerRequestOptions(); } - - final CosmosContainerRequestOptions requestOptions = options; - if(!database.getClient().getTracerProvider().isEnabled()){ - return readInternal(options); - } - - return withContext(context -> read(requestOptions, context)); + return database.getDocClientWrapper().readCollection(getLink(), ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosContainerResponse(response)).single(); } /** @@ -144,13 +108,8 @@ public Mono delete(CosmosContainerRequestOptions option if (options == null) { options = new CosmosContainerRequestOptions(); } - - if (!database.getClient().getTracerProvider().isEnabled()) { - return deleteInternal(options); - } - - final CosmosContainerRequestOptions requestOptions = options; - return withContext(context -> deleteInternal(requestOptions, context)); + return database.getDocClientWrapper().deleteCollection(getLink(), ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosContainerResponse(response)).single(); } /** @@ -202,13 +161,9 @@ public Mono replace( if (options == null) { options = new CosmosContainerRequestOptions(); } - - if(!database.getClient().getTracerProvider().isEnabled()) { - return replaceInternal(containerProperties, options); - } - - final CosmosContainerRequestOptions requestOptions = options; - return withContext(context -> replaceInternal(containerProperties, requestOptions, context)); + return database.getDocClientWrapper() + .replaceCollection(ModelBridgeInternal.getV2Collection(containerProperties), ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosContainerResponse(response)).single(); } /* CosmosAsyncItem operations */ @@ -253,6 +208,7 @@ public Mono> createItem( return createItem(item, options); } + /** * Creates a Cosmos item. * @@ -265,35 +221,16 @@ public Mono> createItem(T item, CosmosItemRequestOptio if (options == null) { options = new CosmosItemRequestOptions(); } - - if (!database.getClient().getTracerProvider().isEnabled()) { - return createItemInternal(item, options); - } - - final CosmosItemRequestOptions requestOptions = options; - return withContext(context -> createItemInternal(item, requestOptions, context)); - } - - private Mono> createItemInternal(T item, CosmosItemRequestOptions options, Context context) { - Mono> responseMono = createItemInternal(item, options); - return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, - context, - this.createItemSpanName, - database.getId(), - database.getClient().getServiceEndpoint()); - } - - private Mono> createItemInternal(T item, CosmosItemRequestOptions options) { @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) + .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) + .single(); } /** @@ -327,13 +264,14 @@ public Mono> upsertItem(T item, CosmosItemRequestOptio if (options == null) { options = new CosmosItemRequestOptions(); } - - if(!database.getClient().getTracerProvider().isEnabled()) { - return upsertItemInternal(item, options); - } - - final CosmosItemRequestOptions requestOptions = options; - return withContext(context -> upsertItemInternal(item, requestOptions, context)); + @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(); } /** @@ -367,12 +305,10 @@ CosmosPagedFlux readAllItems(Class classType) { */ CosmosPagedFlux readAllItems(CosmosQueryRequestOptions options, Class classType) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), this.readAllItemsSpanName, - this.getDatabase().getClient().getServiceEndpoint(), database.getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDatabase().getDocClientWrapper().readDocuments(getLink(), options).map( response -> prepareFeedResponse(response, classType)); - }, this.getDatabase().getClient().getTracerProvider().isEnabled()); + }); } /** @@ -389,7 +325,7 @@ CosmosPagedFlux readAllItems(CosmosQueryRequestOptions options, Class * error. */ public CosmosPagedFlux queryItems(String query, Class classType) { - return queryItemsInternal(new SqlQuerySpec(query), new CosmosQueryRequestOptions(), classType); + return queryItems(new SqlQuerySpec(query), classType); } /** @@ -407,7 +343,7 @@ public CosmosPagedFlux queryItems(String query, Class classType) { * error. */ public CosmosPagedFlux queryItems(String query, CosmosQueryRequestOptions options, Class classType) { - return queryItemsInternal(new SqlQuerySpec(query), options, classType); + return queryItems(new SqlQuerySpec(query), options, classType); } /** @@ -424,7 +360,7 @@ public CosmosPagedFlux queryItems(String query, CosmosQueryRequestOptions * error. */ public CosmosPagedFlux queryItems(SqlQuerySpec querySpec, Class classType) { - return queryItemsInternal(querySpec, new CosmosQueryRequestOptions(), classType); + return queryItems(querySpec, new CosmosQueryRequestOptions(), classType); } /** @@ -446,17 +382,14 @@ public CosmosPagedFlux queryItems(SqlQuerySpec querySpec, CosmosQueryRequ } private CosmosPagedFlux queryItemsInternal( - SqlQuerySpec sqlQuerySpec, CosmosQueryRequestOptions cosmosQueryRequestOptions, Class classType) { + SqlQuerySpec sqlQuerySpec, CosmosQueryRequestOptions cosmosQueryRequestOptions, Class classType) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = this.queryItemsSpanName; - pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), spanName, - this.getDatabase().getClient().getServiceEndpoint(), database.getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, cosmosQueryRequestOptions); return getDatabase().getDocClientWrapper() .queryDocuments(CosmosAsyncContainer.this.getLink(), sqlQuerySpec, cosmosQueryRequestOptions) .map(response -> prepareFeedResponse(response, classType)); - }, this.getDatabase().getClient().getTracerProvider().isEnabled()); + }); } private FeedResponse prepareFeedResponse(FeedResponse response, Class classType) { @@ -521,14 +454,12 @@ public Mono> readItem( if (options == null) { options = new CosmosItemRequestOptions(); } - ModelBridgeInternal.setPartitionKey(options, partitionKey); RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); - if(!database.getClient().getTracerProvider().isEnabled()) { - return readItemInternal(itemId, requestOptions, itemType); - } - - return withContext(context -> readItemInternal(itemId, requestOptions, itemType, context)); + return this.getDatabase().getDocClientWrapper() + .readDocument(getItemLink(itemId), requestOptions) + .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) + .single(); } /** @@ -570,12 +501,11 @@ public Mono> replaceItem( ModelBridgeInternal.setPartitionKey(options, partitionKey); @SuppressWarnings("unchecked") Class itemType = (Class) item.getClass(); - if(!database.getClient().getTracerProvider().isEnabled()){ - return replaceItemInternal(itemType, itemId, doc, options); - } - - final CosmosItemRequestOptions requestOptions = options; - return withContext(context -> replaceItemInternal(itemType, itemId, doc, requestOptions,context)); + return this.getDatabase() + .getDocClientWrapper() + .replaceDocument(getItemLink(itemId), doc, ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) + .single(); } /** @@ -611,11 +541,11 @@ public Mono> deleteItem( } ModelBridgeInternal.setPartitionKey(options, partitionKey); RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options); - if(!database.getClient().getTracerProvider().isEnabled()) { - return deleteItemInternal(itemId, requestOptions); - } - - return withContext(context -> deleteItemInternal(itemId, requestOptions, context)); + return this.getDatabase() + .getDocClientWrapper() + .deleteDocument(getItemLink(itemId), requestOptions) + .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponseWithObjectType(response)) + .single(); } private String getItemLink(String itemId) { @@ -651,14 +581,12 @@ public CosmosAsyncScripts getScripts() { */ public CosmosPagedFlux readAllConflicts(CosmosQueryRequestOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), this.readAllConflictsSpanName, - this.getDatabase().getClient().getServiceEndpoint(), database.getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper().readConflicts(getLink(), options) .map(response -> BridgeInternal.createFeedResponse( ModelBridgeInternal.getCosmosConflictPropertiesFromV2Results(response.getResults()), response.getResponseHeaders())); - }, this.getDatabase().getClient().getTracerProvider().isEnabled()); + }); } /** @@ -682,14 +610,12 @@ public CosmosPagedFlux queryConflicts(String query) { */ public CosmosPagedFlux queryConflicts(String query, CosmosQueryRequestOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), this.queryConflictsSpanName, - this.getDatabase().getClient().getServiceEndpoint(), database.getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return database.getDocClientWrapper().queryConflicts(getLink(), query, options) .map(response -> BridgeInternal.createFeedResponse( ModelBridgeInternal.getCosmosConflictPropertiesFromV2Results(response.getResults()), response.getResponseHeaders())); - }, this.getDatabase().getClient().getTracerProvider().isEnabled()); + }); } /** @@ -703,17 +629,35 @@ public CosmosAsyncConflict getConflict(String id) { } /** - * Replace the throughput . + * Replace the throughput provisioned for the current container. * - * @param throughputProperties the throughput properties - * @return the mono containing throughput response + * @param throughputProperties the throughput properties. + * @return the mono containing throughput response. */ public Mono replaceThroughput(ThroughputProperties throughputProperties) { - if (!this.database.getClient().getTracerProvider().isEnabled()) { - return replaceThroughputInternal(this.read(), throughputProperties); - } - - return withContext(context -> replaceThroughputInternal(throughputProperties, context)); + return this.read() + .flatMap(response -> this.database.getDocClientWrapper() + .queryOffers(database.getOfferQuerySpecFromResourceId(response.getProperties() + .getResourceId()) + , new CosmosQueryRequestOptions()) + .single() + .flatMap(offerFeedResponse -> { + if (offerFeedResponse.getResults().isEmpty()) { + return Mono.error(BridgeInternal + .createCosmosException( + 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.database.getDocClientWrapper() + .replaceOffer(updatedOffer) + .single(); + }).map(ModelBridgeInternal::createThroughputRespose)); } /** @@ -722,13 +666,30 @@ public Mono replaceThroughput(ThroughputProperties throughpu * @return the mono containing throughput response. */ public Mono readThroughput() { - if (!this.database.getClient().getTracerProvider().isEnabled()) { - return readThroughputInternal(this.read()); - } - - return withContext(context -> readThroughputInternal(context)); + return this.read() + .flatMap(response -> this.database.getDocClientWrapper() + .queryOffers(database.getOfferQuerySpecFromResourceId(response.getProperties() + .getResourceId()) + , new CosmosQueryRequestOptions()) + .single() + .flatMap(offerFeedResponse -> { + if (offerFeedResponse.getResults().isEmpty()) { + return Mono.error(BridgeInternal + .createCosmosException( + HttpConstants.StatusCodes.BADREQUEST, + "No offers found for the resource " + + this.getId())); + } + return this.database.getDocClientWrapper() + .readOffer(offerFeedResponse.getResults() + .get(0) + .getSelfLink()) + .single(); + }) + .map(ModelBridgeInternal::createThroughputRespose)); } + /** * Gets the parent {@link CosmosAsyncDatabase} for the current container. * @@ -749,214 +710,4 @@ String getParentLink() { String getLink() { return this.link; } - - private Mono> deleteItemInternal( - String itemId, - RequestOptions requestOptions, - Context context) { - Mono> responseMono = deleteItemInternal(itemId, requestOptions); - return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, - context, - this.deleteItemSpanName, - 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(); - } - - private Mono> replaceItemInternal( - Class itemType, - String itemId, - Document doc, - CosmosItemRequestOptions options, - Context context) { - Mono> responseMono = replaceItemInternal(itemType, itemId, doc, options); - return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, - context, this.replaceItemSpanName, database.getId(), database.getClient().getServiceEndpoint()); - } - - private Mono> replaceItemInternal( - Class itemType, - String itemId, - Document doc, - CosmosItemRequestOptions options) { - return this.getDatabase() - .getDocClientWrapper() - .replaceDocument(getItemLink(itemId), doc, ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) - .single(); - } - - private Mono> upsertItemInternal(T item, CosmosItemRequestOptions options, Context context) { - Mono> responseMono = upsertItemInternal(item, options); - return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, - context, - this.upsertItemSpanName, - database.getId(), - database.getClient().getServiceEndpoint()); - } - - private Mono> upsertItemInternal(T item, CosmosItemRequestOptions options) { - @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(); - } - - private Mono> readItemInternal( - String itemId, - RequestOptions requestOptions, Class itemType, - Context context) { - Mono> responseMono = readItemInternal(itemId, requestOptions, itemType); - return database.getClient().getTracerProvider().traceEnabledCosmosItemResponsePublisher(responseMono, - context, - this.readItemSpanName, - database.getId(), - database.getClient().getServiceEndpoint()); - } - - private Mono> readItemInternal( - String itemId, - RequestOptions requestOptions, Class itemType) { - return this.getDatabase().getDocClientWrapper() - .readDocument(getItemLink(itemId), requestOptions) - .map(response -> ModelBridgeInternal.createCosmosAsyncItemResponse(response, itemType)) - .single(); - } - - Mono read(CosmosContainerRequestOptions options, Context context) { - Mono responseMono = readInternal(options); - return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, - this.readContainerSpanName, - database.getId(), - database.getClient().getServiceEndpoint()); - } - - private Mono readInternal(CosmosContainerRequestOptions options) { - return database.getDocClientWrapper().readCollection(getLink(), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosContainerResponse(response)).single(); - } - - private Mono deleteInternal(CosmosContainerRequestOptions options, Context context) { - Mono responseMono = deleteInternal(options); - return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, - this.deleteContainerSpanName, - database.getId(), - database.getClient().getServiceEndpoint()); - } - - private Mono deleteInternal(CosmosContainerRequestOptions options) { - return database.getDocClientWrapper().deleteCollection(getLink(), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosContainerResponse(response)).single(); - } - - private Mono replaceInternal(CosmosContainerProperties containerProperties, - CosmosContainerRequestOptions options, - Context context) { - Mono responseMono = replaceInternal(containerProperties, options); - return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, - this.replaceContainerSpanName, - 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.createCosmosContainerResponse(response)).single(); - } - - private Mono readThroughputInternal(Context context) { - Context nestedContext = context.addData(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.COSMOS_CALL_DEPTH_VAL); - Mono responseMono = readThroughputInternal(this.read(new CosmosContainerRequestOptions(), - nestedContext)); - return this.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, - this.readThroughputSpanName, - database.getId(), - database.getClient().getServiceEndpoint()); - } - - private Mono readThroughputInternal(Mono responseMono) { - return responseMono - .flatMap(response -> this.database.getDocClientWrapper() - .queryOffers(database.getOfferQuerySpecFromResourceId(response.getProperties() - .getResourceId()) - , new CosmosQueryRequestOptions()) - .single() - .flatMap(offerFeedResponse -> { - if (offerFeedResponse.getResults().isEmpty()) { - return Mono.error(BridgeInternal - .createCosmosException( - HttpConstants.StatusCodes.BADREQUEST, - "No offers found for the resource " - + this.getId())); - } - return this.database.getDocClientWrapper() - .readOffer(offerFeedResponse.getResults() - .get(0) - .getSelfLink()) - .single(); - }) - .map(ModelBridgeInternal::createThroughputRespose)); - } - - private Mono replaceThroughputInternal(ThroughputProperties throughputProperties, - Context context) { - Context nestedContext = context.addData(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.COSMOS_CALL_DEPTH_VAL); - Mono responseMono = - replaceThroughputInternal(this.read(new CosmosContainerRequestOptions(), nestedContext), - throughputProperties); - return this.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, - this.replaceThroughputSpanName, - database.getId(), - database.getClient().getServiceEndpoint()); - } - - private Mono replaceThroughputInternal(Mono responseMono, - ThroughputProperties throughputProperties) { - return responseMono - .flatMap(response -> this.database.getDocClientWrapper() - .queryOffers(database.getOfferQuerySpecFromResourceId(response.getProperties() - .getResourceId()) - , new CosmosQueryRequestOptions()) - .single() - .flatMap(offerFeedResponse -> { - if (offerFeedResponse.getResults().isEmpty()) { - return Mono.error(BridgeInternal - .createCosmosException( - 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.database.getDocClientWrapper() - .replaceOffer(updatedOffer) - .single(); - }).map(ModelBridgeInternal::createThroughputRespose)); - } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java index 1da2affd0916..7d495d715b45 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,24 +2,22 @@ // Licensed under the MIT License. package com.azure.cosmos; -import com.azure.core.util.Context; 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.CosmosContainerResponse; +import com.azure.cosmos.models.CosmosDatabaseResponse; +import com.azure.cosmos.models.CosmosUserResponse; import com.azure.cosmos.models.CosmosContainerProperties; import com.azure.cosmos.models.CosmosContainerRequestOptions; -import com.azure.cosmos.models.CosmosContainerResponse; import com.azure.cosmos.models.CosmosDatabaseRequestOptions; -import com.azure.cosmos.models.CosmosDatabaseResponse; -import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.CosmosUserProperties; -import com.azure.cosmos.models.CosmosUserResponse; +import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.ModelBridgeInternal; import com.azure.cosmos.models.SqlParameter; import com.azure.cosmos.models.SqlQuerySpec; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.models.ThroughputProperties; import com.azure.cosmos.models.ThroughputResponse; import com.azure.cosmos.util.CosmosPagedFlux; @@ -30,7 +28,6 @@ import java.util.Collections; import java.util.List; -import static com.azure.core.util.FluxUtil.withContext; import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount; /** @@ -85,12 +82,8 @@ public Mono read(CosmosDatabaseRequestOptions options) { if (options == null) { options = new CosmosDatabaseRequestOptions(); } - - if (!client.getTracerProvider().isEnabled()) { - return readInternal(options); - } - final CosmosDatabaseRequestOptions requestOptions = options; - return withContext(context -> readInternal(requestOptions, context)); + return getDocClientWrapper().readDatabase(getLink(), ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosDatabaseResponse(response)).single(); } /** @@ -120,12 +113,8 @@ public Mono delete(CosmosDatabaseRequestOptions options) if (options == null) { options = new CosmosDatabaseRequestOptions(); } - if (!client.getTracerProvider().isEnabled()) { - return deleteInternal(options); - } - - final CosmosDatabaseRequestOptions requestOptions = options; - return withContext(context -> deleteInternal(requestOptions, context)); + return getDocClientWrapper().deleteDatabase(getLink(), ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosDatabaseResponse(response)).single(); } /* CosmosAsyncContainer operations */ @@ -208,13 +197,10 @@ public Mono createContainer( if (options == null) { options = new CosmosContainerRequestOptions(); } - - if (!client.getTracerProvider().isEnabled()) { - return createContainerInternal(containerProperties, options); - } - - final CosmosContainerRequestOptions requestOptions = options; - return withContext(context -> createContainerInternal(containerProperties, requestOptions, context)); + return getDocClientWrapper() + .createCollection(this.getLink(), ModelBridgeInternal.getV2Collection(containerProperties), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosContainerResponse(response)).single(); } /** @@ -292,12 +278,7 @@ public Mono createContainer(String id, String partition public Mono createContainerIfNotExists( CosmosContainerProperties containerProperties) { CosmosAsyncContainer container = getContainer(containerProperties.getId()); - if (!client.getTracerProvider().isEnabled()) { - return createContainerIfNotExistsInternal(container.read(), containerProperties, null, null); - } - - return withContext(context -> createContainerIfNotExistsInternal(containerProperties, container, null, - context)); + return createContainerIfNotExistsInternal(containerProperties, container, null); } /** @@ -322,12 +303,7 @@ Mono createContainerIfNotExists( CosmosContainerRequestOptions options = new CosmosContainerRequestOptions(); ModelBridgeInternal.setThroughputProperties(options, ThroughputProperties.createManualThroughput(throughput)); CosmosAsyncContainer container = getContainer(containerProperties.getId()); - if (!client.getTracerProvider().isEnabled()) { - return createContainerIfNotExistsInternal(container.read(), containerProperties, options, null); - } - - return withContext(context -> createContainerIfNotExistsInternal(containerProperties, container, options, - context)); + return createContainerIfNotExistsInternal(containerProperties, container, options); } /** @@ -352,12 +328,7 @@ public Mono createContainerIfNotExists( CosmosContainerRequestOptions options = new CosmosContainerRequestOptions(); ModelBridgeInternal.setThroughputProperties(options, throughputProperties); CosmosAsyncContainer container = getContainer(containerProperties.getId()); - if (!client.getTracerProvider().isEnabled()) { - return createContainerIfNotExistsInternal(container.read(), containerProperties, options, null); - } - - return withContext(context -> createContainerIfNotExistsInternal(containerProperties, container, options, - context)); + return createContainerIfNotExistsInternal(containerProperties, container, options); } /** @@ -374,21 +345,16 @@ public Mono createContainerIfNotExists( */ public Mono createContainerIfNotExists(String id, String partitionKeyPath) { CosmosAsyncContainer container = getContainer(id); - if (!client.getTracerProvider().isEnabled()) { - return createContainerIfNotExistsInternal(container.read(), new CosmosContainerProperties(id, - partitionKeyPath), null, null); - } - - return withContext(context -> createContainerIfNotExistsInternal(new CosmosContainerProperties(id, - partitionKeyPath), container, null, - context)); + return createContainerIfNotExistsInternal(new CosmosContainerProperties(id, partitionKeyPath), + container, + null); } /** * Creates a Cosmos container if it does not exist on the service. *

- * The throughput properties will only be used if the specified container - * does not exist and therefor a new container will be created. + * The throughput setting will only be used if the specified container + * does not exist and a new container will be created. * * After subscription the operation will be performed. The {@link Mono} upon * successful completion will contain a cosmos container response with the @@ -396,29 +362,25 @@ public Mono createContainerIfNotExists(String id, Strin * * @param id the cosmos container id. * @param partitionKeyPath the partition key path. - * @param throughputProperties the throughput properties for the container. + * @param throughput the throughput for the container. * @return a {@link Mono} containing the cosmos container response with the * created container or an error. */ - public Mono createContainerIfNotExists( + Mono createContainerIfNotExists( String id, String partitionKeyPath, - ThroughputProperties throughputProperties) { + int throughput) { CosmosContainerRequestOptions options = new CosmosContainerRequestOptions(); - ModelBridgeInternal.setThroughputProperties(options, throughputProperties); + ModelBridgeInternal.setThroughputProperties(options, ThroughputProperties.createManualThroughput(throughput)); CosmosAsyncContainer container = getContainer(id); - if (!client.getTracerProvider().isEnabled()) { - return createContainerIfNotExistsInternal(container.read(), new CosmosContainerProperties(id, partitionKeyPath), options, null); - } - - return withContext(context -> createContainerIfNotExistsInternal(new CosmosContainerProperties(id, - partitionKeyPath), container, options, context)); + return createContainerIfNotExistsInternal(new CosmosContainerProperties(id, partitionKeyPath), container, + options); } /** * Creates a Cosmos container if it does not exist on the service. *

- * The throughput setting will only be used if the specified container - * does not exist and a new container will be created. + * The throughput properties will only be used if the specified container + * does not exist and therefor a new container will be created. * * After subscription the operation will be performed. The {@link Mono} upon * successful completion will contain a cosmos container response with the @@ -426,24 +388,33 @@ public Mono createContainerIfNotExists( * * @param id the cosmos container id. * @param partitionKeyPath the partition key path. - * @param throughput the throughput for the container. + * @param throughputProperties the throughput properties for the container. * @return a {@link Mono} containing the cosmos container response with the * created container or an error. */ - Mono createContainerIfNotExists( + public Mono createContainerIfNotExists( String id, String partitionKeyPath, - int throughput) { + ThroughputProperties throughputProperties) { CosmosContainerRequestOptions options = new CosmosContainerRequestOptions(); - ModelBridgeInternal.setThroughputProperties(options, ThroughputProperties.createManualThroughput(throughput)); + ModelBridgeInternal.setThroughputProperties(options, throughputProperties); CosmosAsyncContainer container = getContainer(id); - if (!client.getTracerProvider().isEnabled()) { - return createContainerIfNotExistsInternal(container.read(), new CosmosContainerProperties(id, - partitionKeyPath), - options, null); - } + return createContainerIfNotExistsInternal(new CosmosContainerProperties(id, partitionKeyPath), container, + options); + } - return withContext(context -> createContainerIfNotExistsInternal(new CosmosContainerProperties(id, - partitionKeyPath), container, options, context)); + private Mono createContainerIfNotExistsInternal( + CosmosContainerProperties containerProperties, CosmosAsyncContainer container, + CosmosContainerRequestOptions options) { + return container.read(options).onErrorResume(exception -> { + final Throwable unwrappedException = Exceptions.unwrap(exception); + if (unwrappedException instanceof CosmosException) { + final CosmosException cosmosException = (CosmosException) unwrappedException; + if (cosmosException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) { + return createContainer(containerProperties, options); + } + } + return Mono.error(unwrappedException); + }); } /** @@ -459,15 +430,12 @@ Mono createContainerIfNotExists( */ public CosmosPagedFlux readAllContainers(CosmosQueryRequestOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "readAllContainers." + this.getId(); - pagedFluxOptions.setTracerInformation(this.getClient().getTracerProvider(), spanName, - this.getClient().getServiceEndpoint(), getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().readCollections(getLink(), options) .map(response -> BridgeInternal.createFeedResponse( ModelBridgeInternal.getCosmosContainerPropertiesFromV2Results(response.getResults()), response.getResponseHeaders())); - }, this.getClient().getTracerProvider().isEnabled()); + }); } /** @@ -496,7 +464,7 @@ public CosmosPagedFlux readAllContainers() { * obtained containers or an error. */ public CosmosPagedFlux queryContainers(String query) { - return queryContainersInternal(new SqlQuerySpec(query), new CosmosQueryRequestOptions()); + return queryContainers(new SqlQuerySpec(query)); } /** @@ -512,7 +480,7 @@ public CosmosPagedFlux queryContainers(String query) * obtained containers or an error. */ public CosmosPagedFlux queryContainers(String query, CosmosQueryRequestOptions options) { - return queryContainersInternal(new SqlQuerySpec(query), options); + return queryContainers(new SqlQuerySpec(query), options); } /** @@ -527,7 +495,7 @@ public CosmosPagedFlux queryContainers(String query, * obtained containers or an error. */ public CosmosPagedFlux queryContainers(SqlQuerySpec querySpec) { - return queryContainersInternal(querySpec, new CosmosQueryRequestOptions()); + return queryContainers(querySpec, new CosmosQueryRequestOptions()); } /** @@ -542,9 +510,14 @@ 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 - , CosmosQueryRequestOptions options) { - return queryContainersInternal(querySpec, options); + public CosmosPagedFlux queryContainers(SqlQuerySpec querySpec, CosmosQueryRequestOptions options) { + return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); + return getDocClientWrapper().queryCollections(getLink(), querySpec, options) + .map(response -> BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosContainerPropertiesFromV2Results(response.getResults()), + response.getResponseHeaders())); + }); } /** @@ -557,6 +530,8 @@ public CosmosAsyncContainer getContainer(String id) { return new CosmosAsyncContainer(id, this); } + /** User operations **/ + /** * Creates a user After subscription the operation will be performed. The * {@link Mono} upon successful completion will contain a single resource @@ -568,12 +543,11 @@ public CosmosAsyncContainer getContainer(String id) { * created cosmos user or an error. */ public Mono createUser(CosmosUserProperties userProperties) { - if (!client.getTracerProvider().isEnabled()) { - return createUserInternal(userProperties); - } - return withContext(context -> createUserInternal(userProperties, context)); + return getDocClientWrapper().createUser(this.getLink(), ModelBridgeInternal.getV2User(userProperties), null) + .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); } + /** * Upsert a user. Upsert will create a new user if it doesn't exist, or replace * the existing one if it does. After subscription the operation will be @@ -586,11 +560,8 @@ public Mono createUser(CosmosUserProperties userProperties) * upserted user or an error. */ public Mono upsertUser(CosmosUserProperties userProperties) { - if (!client.getTracerProvider().isEnabled()) { - return upsertUserInternal(userProperties); - } - - return withContext(context -> upsertUserInternal(userProperties, context)); + return getDocClientWrapper().upsertUser(this.getLink(), ModelBridgeInternal.getV2User(userProperties), null) + .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); } /** @@ -620,15 +591,12 @@ public CosmosPagedFlux readAllUsers() { */ CosmosPagedFlux readAllUsers(CosmosQueryRequestOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "readAllUsers." + this.getId(); - pagedFluxOptions.setTracerInformation(this.getClient().getTracerProvider(), spanName, - this.getClient().getServiceEndpoint(), getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDocClientWrapper().readUsers(getLink(), options) .map(response -> BridgeInternal.createFeedResponse( ModelBridgeInternal.getCosmosUserPropertiesFromV2Results(response.getResults()), response .getResponseHeaders())); - }, this.getClient().getTracerProvider().isEnabled()); + }); } /** @@ -659,7 +627,7 @@ public CosmosPagedFlux queryUsers(String query) { * obtained users or an error. */ public CosmosPagedFlux queryUsers(String query, CosmosQueryRequestOptions options) { - return queryUsersInternal(new SqlQuerySpec(query), options); + return queryUsers(new SqlQuerySpec(query), options); } /** @@ -674,7 +642,7 @@ public CosmosPagedFlux queryUsers(String query, CosmosQuer * obtained users or an error. */ public CosmosPagedFlux queryUsers(SqlQuerySpec querySpec) { - return queryUsersInternal(querySpec, new CosmosQueryRequestOptions()); + return queryUsers(querySpec, new CosmosQueryRequestOptions()); } /** @@ -690,7 +658,13 @@ public CosmosPagedFlux queryUsers(SqlQuerySpec querySpec) * obtained users or an error. */ public CosmosPagedFlux queryUsers(SqlQuerySpec querySpec, CosmosQueryRequestOptions options) { - return queryUsersInternal(querySpec, options); + return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); + return getDocClientWrapper().queryUsers(getLink(), querySpec, options) + .map(response -> BridgeInternal.createFeedResponseWithQueryMetrics( + ModelBridgeInternal.getCosmosUserPropertiesFromV2Results(response.getResults()), response.getResponseHeaders(), + ModelBridgeInternal.queryMetrics(response))); + }); } /** @@ -711,11 +685,30 @@ public CosmosAsyncUser getUser(String id) { * @return the mono. */ public Mono replaceThroughput(ThroughputProperties throughputProperties) { - if(!this.client.getTracerProvider().isEnabled()) { - return replaceThroughputInternal(this.read(), throughputProperties); - } - - return withContext(context -> replaceThroughputInternal(throughputProperties, context)); + return this.read() + .flatMap(response -> this.getDocClientWrapper() + .queryOffers(getOfferQuerySpecFromResourceId(response.getProperties().getResourceId()), + new CosmosQueryRequestOptions()) + .single() + .flatMap(offerFeedResponse -> { + if (offerFeedResponse.getResults().isEmpty()) { + return Mono.error(BridgeInternal + .createCosmosException( + 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)); } /** @@ -724,11 +717,26 @@ public Mono replaceThroughput(ThroughputProperties throughpu * @return the mono containing throughput response. */ public Mono readThroughput() { - if(!this.client.getTracerProvider().isEnabled()) { - return readThroughputInternal(this.read()); - } - - return withContext(context -> readThroughputInternal(context)); + return this.read() + .flatMap(response -> getDocClientWrapper() + .queryOffers(getOfferQuerySpecFromResourceId(response.getProperties().getResourceId()), + new CosmosQueryRequestOptions()) + .single() + .flatMap(offerFeedResponse -> { + if (offerFeedResponse.getResults().isEmpty()) { + return Mono.error(BridgeInternal + .createCosmosException( + 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) { @@ -760,218 +768,4 @@ String getLink() { return this.link; } - private CosmosPagedFlux queryContainersInternal(SqlQuerySpec querySpec - , CosmosQueryRequestOptions options) { - return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String 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())); - }, this.getClient().getTracerProvider().isEnabled()); - } - - private CosmosPagedFlux queryUsersInternal(SqlQuerySpec querySpec, CosmosQueryRequestOptions options) { - return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String 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))); - }, this.getClient().getTracerProvider().isEnabled()); - } - - private Mono createContainerIfNotExistsInternal( - CosmosContainerProperties containerProperties, - CosmosAsyncContainer container, - CosmosContainerRequestOptions options, - Context context) { - String spanName = "createContainerIfNotExistsInternal." + containerProperties.getId(); - Context nestedContext = context.addData(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.COSMOS_CALL_DEPTH_VAL); - if (options == null) { - options = new CosmosContainerRequestOptions(); - } - - Mono responseMono = createContainerIfNotExistsInternal(container.read(options, nestedContext), containerProperties, options, nestedContext); - return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, - getId(), - getClient().getServiceEndpoint()); - } - - private Mono createContainerIfNotExistsInternal( - Mono responseMono, - CosmosContainerProperties containerProperties, - CosmosContainerRequestOptions options, - Context context) { - return responseMono.onErrorResume(exception -> { - final Throwable unwrappedException = Exceptions.unwrap(exception); - if (unwrappedException instanceof CosmosException) { - final CosmosException cosmosException = (CosmosException) unwrappedException; - if (cosmosException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) { - if(context != null) { - return createContainerInternal(containerProperties, options, context); - } - - return createContainer(containerProperties, options); - } - } - return Mono.error(unwrappedException); - }); - } - - private Mono createContainerInternal( - CosmosContainerProperties containerProperties, - CosmosContainerRequestOptions options, - Context context) { - String spanName = "createContainer." + containerProperties.getId(); - 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.createCosmosContainerResponse(response)).single(); - } - - 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 readInternal(CosmosDatabaseRequestOptions options) { - return getDocClientWrapper().readDatabase(getLink(), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosDatabaseResponse(response)).single(); - } - - private Mono deleteInternal(CosmosDatabaseRequestOptions options, Context context) { - String spanName = "deleteDatabase." + this.getId(); - Mono responseMono = deleteInternal(options); - return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, - getId(), - getClient().getServiceEndpoint()); - } - - private Mono deleteInternal(CosmosDatabaseRequestOptions options) { - return getDocClientWrapper().deleteDatabase(getLink(), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosDatabaseResponse(response)).single(); - } - - private Mono createUserInternal(CosmosUserProperties userProperties, Context context) { - String spanName = "createUser." + this.getId(); - Mono responseMono = createUserInternal(userProperties); - return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, - getId(), - getClient().getServiceEndpoint()); - } - - private Mono createUserInternal(CosmosUserProperties userProperties) { - return getDocClientWrapper().createUser(this.getLink(), ModelBridgeInternal.getV2User(userProperties), null) - .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); - } - - private Mono upsertUserInternal(CosmosUserProperties userProperties, Context context) { - String spanName = "upsertUser." + this.getId(); - Mono responseMono = upsertUserInternal(userProperties); - return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, getId(), getClient().getServiceEndpoint()); - } - - private Mono upsertUserInternal(CosmosUserProperties userProperties) { - return getDocClientWrapper().upsertUser(this.getLink(), ModelBridgeInternal.getV2User(userProperties), null) - .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); - } - - private Mono replaceThroughputInternal(ThroughputProperties throughputProperties, Context context){ - String spanName = "replaceThroughput." + this.getId(); - Context nestedContext = context.addData(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.COSMOS_CALL_DEPTH_VAL); - Mono responseMono = replaceThroughputInternal(this.readInternal(new CosmosDatabaseRequestOptions(), nestedContext), throughputProperties); - return this.client.getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, - spanName, - getId(), - getClient().getServiceEndpoint()); - } - - private Mono replaceThroughputInternal(Mono responseMono, ThroughputProperties throughputProperties) { - return responseMono - .flatMap(response -> this.getDocClientWrapper() - .queryOffers(getOfferQuerySpecFromResourceId(response.getProperties().getResourceId()), - new CosmosQueryRequestOptions()) - .single() - .flatMap(offerFeedResponse -> { - if (offerFeedResponse.getResults().isEmpty()) { - return Mono.error(BridgeInternal - .createCosmosException( - 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)); - } - - 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().traceEnabledCosmosResponsePublisher(responseMono, - context, - spanName, - getId(), - getClient().getServiceEndpoint()); - } - - private Mono readThroughputInternal(Mono responseMono) { - return responseMono - .flatMap(response -> getDocClientWrapper() - .queryOffers(getOfferQuerySpecFromResourceId(response.getProperties().getResourceId()), - new CosmosQueryRequestOptions()) - .single() - .flatMap(offerFeedResponse -> { - if (offerFeedResponse.getResults().isEmpty()) { - return Mono.error(BridgeInternal - .createCosmosException( - HttpConstants.StatusCodes.BADREQUEST, - "No offers found for the " + - "resource " + this.getId())); - } - return getDocClientWrapper() - .readOffer(offerFeedResponse.getResults() - .get(0) - .getSelfLink()) - .single(); - }) - .map(ModelBridgeInternal::createThroughputRespose)); - } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncPermission.java index 128937f4f762..f01a52766820 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,7 +2,6 @@ // 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.models.CosmosPermissionResponse; import com.azure.cosmos.models.CosmosPermissionProperties; @@ -10,8 +9,6 @@ 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 */ @@ -59,13 +56,11 @@ public Mono read(CosmosPermissionRequestOptions option if (options == null) { options = new CosmosPermissionRequestOptions(); } - - if (!cosmosUser.getDatabase().getClient().getTracerProvider().isEnabled()) { - return readInternal(options); - } - - final CosmosPermissionRequestOptions requestOptions = options; - return withContext(context -> readInternal(requestOptions, context)); + return cosmosUser.getDatabase() + .getDocClientWrapper() + .readPermission(getLink(), ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosPermissionResponse(response)) + .single(); } /** @@ -84,13 +79,13 @@ public Mono replace(CosmosPermissionProperties permiss if (options == null) { options = new CosmosPermissionRequestOptions(); } - - if (!cosmosUser.getDatabase().getClient().getTracerProvider().isEnabled()) { - return replaceInternal(permissionProperties, options); - } - - final CosmosPermissionRequestOptions requestOptions = options; - return withContext(context -> replaceInternal(permissionProperties, requestOptions, context)); + CosmosAsyncDatabase databaseContext = cosmosUser.getDatabase(); + return databaseContext + .getDocClientWrapper() + .replacePermission(ModelBridgeInternal.getPermission(permissionProperties, databaseContext.getId()), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosPermissionResponse(response)) + .single(); } /** @@ -107,13 +102,11 @@ public Mono delete(CosmosPermissionRequestOptions opti if (options == null) { options = new CosmosPermissionRequestOptions(); } - - if (!cosmosUser.getDatabase().getClient().getTracerProvider().isEnabled()) { - return deleteInternal(options); - } - - final CosmosPermissionRequestOptions requestOptions = options; - return withContext(context -> deleteInternal(requestOptions, context)); + return cosmosUser.getDatabase() + .getDocClientWrapper() + .deletePermission(getLink(), ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosPermissionResponse(response)) + .single(); } String getURIPathSegment() { @@ -133,67 +126,4 @@ String getLink() { builder.append(getId()); return builder.toString(); } - - private Mono readInternal(CosmosPermissionRequestOptions options, Context context) { - - String spanName = "readPermission." + cosmosUser.getId(); - Mono responseMono = readInternal(options); - return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, - cosmosUser.getDatabase().getId(), - cosmosUser.getDatabase().getClient().getServiceEndpoint()); - } - - private Mono readInternal(CosmosPermissionRequestOptions options) { - - return cosmosUser.getDatabase() - .getDocClientWrapper() - .readPermission(getLink(), ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosPermissionResponse(response)) - .single(); - } - - private Mono replaceInternal(CosmosPermissionProperties permissionProperties, - CosmosPermissionRequestOptions options, - Context context) { - - String spanName = "replacePermission." + cosmosUser.getId(); - Mono responseMono = replaceInternal(permissionProperties, options); - return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, - cosmosUser.getDatabase().getId(), - cosmosUser.getDatabase().getClient().getServiceEndpoint()); - } - - private Mono replaceInternal(CosmosPermissionProperties permissionProperties, - CosmosPermissionRequestOptions options) { - CosmosAsyncDatabase databaseContext = cosmosUser.getDatabase(); - return cosmosUser.getDatabase() - .getDocClientWrapper() - .replacePermission(ModelBridgeInternal.getPermission(permissionProperties, databaseContext.getId()), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosPermissionResponse(response)) - .single(); - } - - private Mono deleteInternal(CosmosPermissionRequestOptions options, - Context context) { - - String spanName = "deletePermission." + cosmosUser.getId(); - Mono responseMono = deleteInternal(options); - return cosmosUser.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, - cosmosUser.getDatabase().getId(), - cosmosUser.getDatabase().getClient().getServiceEndpoint()); - } - - private Mono deleteInternal(CosmosPermissionRequestOptions options) { - - return cosmosUser.getDatabase() - .getDocClientWrapper() - .deletePermission(getLink(), ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosPermissionResponse(response)) - .single(); - } - } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncScripts.java index 7d4e22244948..7ea5249b8b70 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,25 +2,23 @@ // Licensed under the MIT License. package com.azure.cosmos; -import com.azure.core.util.Context; import com.azure.cosmos.implementation.StoredProcedure; import com.azure.cosmos.implementation.Trigger; import com.azure.cosmos.implementation.UserDefinedFunction; -import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.CosmosStoredProcedureResponse; +import com.azure.cosmos.models.CosmosTriggerResponse; +import com.azure.cosmos.models.CosmosUserDefinedFunctionResponse; import com.azure.cosmos.models.CosmosStoredProcedureProperties; import com.azure.cosmos.models.CosmosStoredProcedureRequestOptions; -import com.azure.cosmos.models.CosmosStoredProcedureResponse; import com.azure.cosmos.models.CosmosTriggerProperties; -import com.azure.cosmos.models.CosmosTriggerResponse; import com.azure.cosmos.models.CosmosUserDefinedFunctionProperties; -import com.azure.cosmos.models.CosmosUserDefinedFunctionResponse; +import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.ModelBridgeInternal; import com.azure.cosmos.models.SqlQuerySpec; 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; /** @@ -73,12 +71,10 @@ public Mono createStoredProcedure( StoredProcedure sProc = new StoredProcedure(); sProc.setId(properties.getId()); sProc.setBody(properties.getBody()); - if (!container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return createStoredProcedureInternal(sProc, options); - } - - final CosmosStoredProcedureRequestOptions requestOptions = options; - return withContext(context -> createStoredProcedureInternal(sProc, requestOptions, context)); + return database.getDocClientWrapper() + .createStoredProcedure(container.getLink(), sProc, ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosStoredProcedureResponse(response)) + .single(); } /** @@ -112,18 +108,13 @@ public CosmosPagedFlux readAllStoredProcedures( */ CosmosPagedFlux readAllStoredProcedures(CosmosQueryRequestOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "readAllStoredProcedures." + 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() .readStoredProcedures(container.getLink(), options) .map(response -> BridgeInternal.createFeedResponse( ModelBridgeInternal.getCosmosStoredProcedurePropertiesFromV2Results(response.getResults()), response.getResponseHeaders())); - }, this.container.getDatabase().getClient().getTracerProvider().isEnabled()); + }); } /** @@ -141,8 +132,8 @@ CosmosPagedFlux readAllStoredProcedures(CosmosQ */ public CosmosPagedFlux queryStoredProcedures( String query, - CosmosQueryRequestOptions options) { - return queryStoredProceduresInternal(new SqlQuerySpec(query), options); + CosmosQueryRequestOptions options) { + return queryStoredProcedures(new SqlQuerySpec(query), options); } /** @@ -161,7 +152,14 @@ public CosmosPagedFlux queryStoredProcedures( public CosmosPagedFlux queryStoredProcedures( SqlQuerySpec querySpec, CosmosQueryRequestOptions options) { - return queryStoredProceduresInternal(querySpec, options); + return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); + return database.getDocClientWrapper() + .queryStoredProcedures(container.getLink(), querySpec, options) + .map(response -> BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosStoredProcedurePropertiesFromV2Results(response.getResults()), + response.getResponseHeaders())); + }); } /** @@ -174,6 +172,7 @@ public CosmosAsyncStoredProcedure getStoredProcedure(String id) { return new CosmosAsyncStoredProcedure(id, this.container); } + /* UDF Operations */ /** @@ -192,11 +191,10 @@ public Mono createUserDefinedFunction( UserDefinedFunction udf = new UserDefinedFunction(); udf.setId(properties.getId()); udf.setBody(properties.getBody()); - if (!container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return createUserDefinedFunctionInternal(udf); - } - return withContext(context -> createUserDefinedFunctionInternal(udf, context)); + return database.getDocClientWrapper() + .createUserDefinedFunction(container.getLink(), udf, null) + .map(response -> ModelBridgeInternal.createCosmosUserDefinedFunctionResponse(response)).single(); } /** @@ -228,18 +226,13 @@ public CosmosPagedFlux readAllUserDefinedFu */ CosmosPagedFlux readAllUserDefinedFunctions(CosmosQueryRequestOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "readAllUserDefinedFunctions." + 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() - .readUserDefinedFunctions(container.getLink(), options) - .map(response -> BridgeInternal.createFeedResponse( - ModelBridgeInternal.getCosmosUserDefinedFunctionPropertiesFromV2Results(response.getResults()), - response.getResponseHeaders())); - }, this.container.getDatabase().getClient().getTracerProvider().isEnabled()); + .readUserDefinedFunctions(container.getLink(), options) + .map(response -> BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosUserDefinedFunctionPropertiesFromV2Results(response.getResults()), + response.getResponseHeaders())); + }); } /** @@ -279,7 +272,14 @@ public CosmosPagedFlux queryUserDefinedFunc public CosmosPagedFlux queryUserDefinedFunctions( SqlQuerySpec querySpec, CosmosQueryRequestOptions options) { - return queryUserDefinedFunctionsInternal(querySpec, options); + return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); + return database.getDocClientWrapper() + .queryUserDefinedFunctions(container.getLink(), querySpec, options) + .map(response -> BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosUserDefinedFunctionPropertiesFromV2Results(response.getResults()), + response.getResponseHeaders())); + }); } /** @@ -305,11 +305,12 @@ public CosmosAsyncUserDefinedFunction getUserDefinedFunction(String id) { * @return an {@link Mono} containing the single resource response with the created trigger or an error. */ public Mono createTrigger(CosmosTriggerProperties properties) { - if (!container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return createTriggerInternal(properties); - } + Trigger trigger = new Trigger(ModelBridgeInternal.toJsonFromJsonSerializable(ModelBridgeInternal.getResource(properties))); - return withContext(context -> createTriggerInternal(properties, context)); + return database.getDocClientWrapper() + .createTrigger(container.getLink(), trigger, null) + .map(response -> ModelBridgeInternal.createCosmosTriggerResponse(response)) + .single(); } /** @@ -343,18 +344,13 @@ public CosmosPagedFlux readAllTriggers() { */ CosmosPagedFlux readAllTriggers(CosmosQueryRequestOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "readAllTriggers." + 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() .readTriggers(container.getLink(), options) .map(response -> BridgeInternal.createFeedResponse( ModelBridgeInternal.getCosmosTriggerPropertiesFromV2Results(response.getResults()), response.getResponseHeaders())); - }, this.container.getDatabase().getClient().getTracerProvider().isEnabled()); + }); } /** @@ -370,7 +366,7 @@ CosmosPagedFlux readAllTriggers(CosmosQueryRequestOptio * error. */ public CosmosPagedFlux queryTriggers(String query, CosmosQueryRequestOptions options) { - return queryTriggersInternal(false, new SqlQuerySpec(query), options); + return queryTriggers(new SqlQuerySpec(query), options); } /** @@ -388,7 +384,14 @@ public CosmosPagedFlux queryTriggers(String query, Cosm public CosmosPagedFlux queryTriggers( SqlQuerySpec querySpec, CosmosQueryRequestOptions options) { - return queryTriggersInternal(true, querySpec, options); + return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { + setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); + return database.getDocClientWrapper() + .queryTriggers(container.getLink(), querySpec, options) + .map(response -> BridgeInternal.createFeedResponse( + ModelBridgeInternal.getCosmosTriggerPropertiesFromV2Results(response.getResults()), + response.getResponseHeaders())); + }); } /** @@ -401,120 +404,4 @@ public CosmosAsyncTrigger getTrigger(String id) { return new CosmosAsyncTrigger(id, this.container); } - private CosmosPagedFlux queryStoredProceduresInternal( - SqlQuerySpec querySpec, - CosmosQueryRequestOptions options) { - return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String 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() - .queryStoredProcedures(container.getLink(), querySpec, options) - .map(response -> BridgeInternal.createFeedResponse( - ModelBridgeInternal.getCosmosStoredProcedurePropertiesFromV2Results(response.getResults()), - response.getResponseHeaders())); - }, this.container.getDatabase().getClient().getTracerProvider().isEnabled()); - } - - private CosmosPagedFlux queryUserDefinedFunctionsInternal( - SqlQuerySpec querySpec, - CosmosQueryRequestOptions options) { - return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String 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())); - }, this.container.getDatabase().getClient().getTracerProvider().isEnabled()); - } - - private CosmosPagedFlux queryTriggersInternal( - boolean isParameterised, - SqlQuerySpec querySpec, - CosmosQueryRequestOptions 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())); - }, this.container.getDatabase().getClient().getTracerProvider().isEnabled()); - } - - private Mono createStoredProcedureInternal(StoredProcedure sProc, - CosmosStoredProcedureRequestOptions options, - 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()); - } - - private Mono createStoredProcedureInternal(StoredProcedure sProc, - CosmosStoredProcedureRequestOptions options) { - return database.getDocClientWrapper() - .createStoredProcedure(container.getLink(), sProc, ModelBridgeInternal.toRequestOptions(options)).map(response -> ModelBridgeInternal.createCosmosStoredProcedureResponse(response)) - .single(); - } - - private Mono createUserDefinedFunctionInternal( - UserDefinedFunction udf, - 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()); - } - - private Mono createUserDefinedFunctionInternal( - UserDefinedFunction udf) { - return database.getDocClientWrapper() - .createUserDefinedFunction(container.getLink(), udf, null).map(response -> ModelBridgeInternal.createCosmosUserDefinedFunctionResponse(response)).single(); - } - - 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(ModelBridgeInternal.getResource(properties))); - return database.getDocClientWrapper() - .createTrigger(container.getLink(), trigger, null) - .map(response -> ModelBridgeInternal.createCosmosTriggerResponse(response)) - .single(); - } - } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncStoredProcedure.java index e6d4433de989..8e0a51768818 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,7 +2,6 @@ // 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.StoredProcedure; import com.azure.cosmos.models.CosmosStoredProcedureResponse; @@ -11,7 +10,6 @@ import com.azure.cosmos.models.ModelBridgeInternal; import reactor.core.publisher.Mono; -import static com.azure.core.util.FluxUtil.withContext; import java.util.List; /** @@ -74,11 +72,12 @@ public Mono read() { * @return an {@link Mono} containing the single resource response with the read stored procedure or an error. */ public Mono read(CosmosStoredProcedureRequestOptions options) { - if (!cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled()) { - return readInternal(options); + if (options == null) { + options = new CosmosStoredProcedureRequestOptions(); } - - return withContext(context -> readInternal(options, context)); + return cosmosContainer.getDatabase().getDocClientWrapper().readStoredProcedure(getLink(), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosStoredProcedureResponse(response)).single(); } /** @@ -107,11 +106,14 @@ public Mono delete() { * @return an {@link Mono} containing the single resource response for the deleted stored procedure or an error. */ public Mono delete(CosmosStoredProcedureRequestOptions options) { - if (!cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled()) { - return deleteInternal(options); + if (options == null) { + options = new CosmosStoredProcedureRequestOptions(); } - - return withContext(context -> deleteInternal(options, context)); + return cosmosContainer.getDatabase() + .getDocClientWrapper() + .deleteStoredProcedure(getLink(), ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosStoredProcedureResponse(response)) + .single(); } /** @@ -127,12 +129,15 @@ public Mono delete(CosmosStoredProcedureRequestOp * @return an {@link Mono} containing the single resource response with the stored procedure response or an error. */ public Mono execute(List procedureParams, - CosmosStoredProcedureRequestOptions options) { - if (!cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled()) { - return executeInternal(procedureParams, options); + CosmosStoredProcedureRequestOptions options) { + if (options == null) { + options = new CosmosStoredProcedureRequestOptions(); } - - return withContext(context -> executeInternal(procedureParams, options, context)); + return cosmosContainer.getDatabase() + .getDocClientWrapper() + .executeStoredProcedure(getLink(), ModelBridgeInternal.toRequestOptions(options), procedureParams) + .map(response -> ModelBridgeInternal.createCosmosStoredProcedureResponse(response)) + .single(); } /** @@ -163,13 +168,17 @@ public Mono replace(CosmosStoredProcedureProperti * @return an {@link Mono} containing the single resource response with the replaced stored procedure or an error. */ public Mono replace(CosmosStoredProcedureProperties storedProcedureProperties, - CosmosStoredProcedureRequestOptions options) { - if (!cosmosContainer.getDatabase().getClient().getTracerProvider().isEnabled()) { - return replaceInternal(storedProcedureProperties, options); + CosmosStoredProcedureRequestOptions options) { + if (options == null) { + options = new CosmosStoredProcedureRequestOptions(); } - - return withContext(context -> replaceInternal(storedProcedureProperties, options, - context)); + return cosmosContainer.getDatabase() + .getDocClientWrapper() + .replaceStoredProcedure(new StoredProcedure(ModelBridgeInternal.toJsonFromJsonSerializable( + ModelBridgeInternal.getResource(storedProcedureProperties))), + ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosStoredProcedureResponse(response)) + .single(); } String getURIPathSegment() { @@ -189,91 +198,4 @@ String getLink() { builder.append(getId()); return builder.toString(); } - - private Mono readInternal(CosmosStoredProcedureRequestOptions options, - Context context) { - String spanName = "readStoredProcedure." + cosmosContainer.getId(); - Mono responseMono = readInternal(options); - return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, - spanName, - cosmosContainer.getDatabase().getId(), - cosmosContainer.getDatabase().getClient().getServiceEndpoint()); - } - - private Mono readInternal(CosmosStoredProcedureRequestOptions options) { - if (options == null) { - options = new CosmosStoredProcedureRequestOptions(); - } - return cosmosContainer.getDatabase().getDocClientWrapper().readStoredProcedure(getLink(), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosStoredProcedureResponse(response)).single(); - } - - private Mono deleteInternal(CosmosStoredProcedureRequestOptions options, - Context context) { - String spanName = "deleteStoredProcedure." + cosmosContainer.getId(); - Mono responseMono = deleteInternal(options); - return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, - spanName, - cosmosContainer.getDatabase().getId(), - cosmosContainer.getDatabase().getClient().getServiceEndpoint()); - } - - private Mono deleteInternal(CosmosStoredProcedureRequestOptions options) { - if (options == null) { - options = new CosmosStoredProcedureRequestOptions(); - } - - return cosmosContainer.getDatabase() - .getDocClientWrapper() - .deleteStoredProcedure(getLink(), ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosStoredProcedureResponse(response)) - .single(); - } - - private Mono executeInternal(List procedureParams, - CosmosStoredProcedureRequestOptions options, - Context context) { - String spanName = "executeStoredProcedure." + cosmosContainer.getId(); - Mono responseMono = executeInternal(procedureParams, options); - return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); - } - - private Mono executeInternal(List procedureParams, - CosmosStoredProcedureRequestOptions options) { - if (options == null) { - options = new CosmosStoredProcedureRequestOptions(); - } - - return cosmosContainer.getDatabase() - .getDocClientWrapper() - .executeStoredProcedure(getLink(), ModelBridgeInternal.toRequestOptions(options), procedureParams) - .map(response -> ModelBridgeInternal.createCosmosStoredProcedureResponse(response)) - .single(); - } - - private Mono replaceInternal(CosmosStoredProcedureProperties storedProcedureSettings, - CosmosStoredProcedureRequestOptions options, - Context context) { - String spanName = "replaceStoredProcedure." + cosmosContainer.getId(); - Mono responseMono = replaceInternal(storedProcedureSettings, options); - return this.cosmosContainer.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, spanName, cosmosContainer.getDatabase().getId(), cosmosContainer.getDatabase().getClient().getServiceEndpoint()); - } - - private Mono replaceInternal(CosmosStoredProcedureProperties storedProcedureSettings, - CosmosStoredProcedureRequestOptions options) { - if (options == null) { - options = new CosmosStoredProcedureRequestOptions(); - } - - return cosmosContainer.getDatabase() - .getDocClientWrapper() - .replaceStoredProcedure(new StoredProcedure(ModelBridgeInternal.toJsonFromJsonSerializable( - ModelBridgeInternal.getResource(storedProcedureSettings))), - ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosStoredProcedureResponse(response)) - .single(); - } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncTrigger.java index 666ddad0a646..ae06eb230079 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,7 +2,6 @@ // 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.Trigger; import com.azure.cosmos.models.CosmosTriggerResponse; @@ -10,8 +9,6 @@ import com.azure.cosmos.models.ModelBridgeInternal; import reactor.core.publisher.Mono; -import static com.azure.core.util.FluxUtil.withContext; - /** * The type Cosmos async trigger. This contains methods to operate on a cosmos trigger asynchronously */ @@ -55,13 +52,14 @@ CosmosAsyncTrigger setId(String id) { * @return an {@link Mono} containing the single resource response for the read cosmos trigger or an error. */ public Mono read() { - if (!container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return readInternal(); - } - - return withContext(context -> readInternal(context)); + return container.getDatabase() + .getDocClientWrapper() + .readTrigger(getLink(), null) + .map(response -> ModelBridgeInternal.createCosmosTriggerResponse(response)) + .single(); } + /** * Replaces a cosmos trigger. *

@@ -73,11 +71,12 @@ public Mono read() { * @return an {@link Mono} containing the single resource response with the replaced cosmos trigger or an error. */ public Mono replace(CosmosTriggerProperties triggerProperties) { - if (!container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return replaceInternal(triggerProperties); - } - - return withContext(context -> replaceInternal(triggerProperties, context)); + return container.getDatabase() + .getDocClientWrapper() + .replaceTrigger(new Trigger(ModelBridgeInternal.toJsonFromJsonSerializable( + ModelBridgeInternal.getResource(triggerProperties))), null) + .map(response -> ModelBridgeInternal.createCosmosTriggerResponse(response)) + .single(); } /** @@ -90,11 +89,11 @@ public Mono replace(CosmosTriggerProperties triggerProper * @return an {@link Mono} containing the single resource response for the deleted cosmos trigger or an error. */ public Mono delete() { - if (!container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return deleteInternal(); - } - - return withContext(context -> deleteInternal(context)); + return container.getDatabase() + .getDocClientWrapper() + .deleteTrigger(getLink(), null) + .map(response -> ModelBridgeInternal.createCosmosTriggerResponse(response)) + .single(); } String getURIPathSegment() { @@ -115,58 +114,4 @@ String getLink() { return builder.toString(); } - private Mono readInternal(Context context) { - String spanName = "readTrigger." + container.getId(); - Mono responseMono = readInternal(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, - spanName, - container.getDatabase().getId(), - container.getDatabase().getClient().getServiceEndpoint()); - } - - private Mono readInternal() { - return container.getDatabase() - .getDocClientWrapper() - .readTrigger(getLink(), null) - .map(response -> ModelBridgeInternal.createCosmosTriggerResponse(response)) - .single(); - } - - private Mono replaceInternal(CosmosTriggerProperties triggerSettings, Context context) { - String spanName = "replaceTrigger." + container.getId(); - Mono responseMono = replaceInternal(triggerSettings); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, - spanName, - container.getDatabase().getId(), - container.getDatabase().getClient().getServiceEndpoint()); - } - - private Mono replaceInternal(CosmosTriggerProperties triggerSettings) { - return container.getDatabase() - .getDocClientWrapper() - .replaceTrigger(new Trigger(ModelBridgeInternal.toJsonFromJsonSerializable( - ModelBridgeInternal.getResource(triggerSettings))), null) - .map(response -> ModelBridgeInternal.createCosmosTriggerResponse(response)) - .single(); - } - - private Mono deleteInternal(Context context) { - String spanName = "deleteTrigger." + container.getId(); - Mono responseMono = deleteInternal(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, - spanName, - container.getDatabase().getId(), - container.getDatabase().getClient().getServiceEndpoint()); - } - - private Mono deleteInternal() { - return container.getDatabase() - .getDocClientWrapper() - .deleteTrigger(getLink(), null) - .map(response -> ModelBridgeInternal.createCosmosTriggerResponse(response)) - .single(); - } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUser.java index f2ea9c442a9a..d23e239e2969 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,9 +3,7 @@ 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.models.CosmosPermissionResponse; import com.azure.cosmos.models.CosmosUserResponse; import com.azure.cosmos.models.CosmosPermissionProperties; @@ -13,11 +11,11 @@ import com.azure.cosmos.models.CosmosUserProperties; import com.azure.cosmos.models.CosmosQueryRequestOptions; 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; /** @@ -58,11 +56,9 @@ CosmosAsyncUser setId(String id) { * @return a {@link Mono} containing the single resource response with the read user or an error. */ public Mono read() { - if (!database.getClient().getTracerProvider().isEnabled()) { - return readInternal(); - } - - return withContext(context -> readInternal(context)); + return this.database.getDocClientWrapper() + .readUser(getLink(), null) + .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); } /** @@ -72,11 +68,9 @@ public Mono read() { * @return a {@link Mono} containing the single resource response with the replaced user or an error. */ public Mono replace(CosmosUserProperties userProperties) { - if (!database.getClient().getTracerProvider().isEnabled()) { - return replaceInternal(userProperties); - } - - return withContext(context -> replaceInternal(userProperties, context)); + return this.database.getDocClientWrapper() + .replaceUser(ModelBridgeInternal.getV2User(userProperties), null) + .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); } /** @@ -85,11 +79,9 @@ public Mono replace(CosmosUserProperties userProperties) { * @return a {@link Mono} containing the single resource response with the deleted user or an error. */ public Mono delete() { - if (!database.getClient().getTracerProvider().isEnabled()) { - return deleteInternal(); - } - - return withContext(context -> deleteInternal(context)); + return this.database.getDocClientWrapper() + .deleteUser(getLink(), null) + .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); } /** @@ -109,14 +101,11 @@ public Mono createPermission( if (options == null) { options = new CosmosPermissionRequestOptions(); } - Permission permission = ModelBridgeInternal.getPermission(permissionProperties, database.getId()); - if (!database.getClient().getTracerProvider().isEnabled()) { - return createPermissionInternal(permission, options); - } - - final CosmosPermissionRequestOptions requesOptions = options; - return withContext(context -> createPermissionInternal(permission, requesOptions, context)); + return database.getDocClientWrapper() + .createPermission(getLink(), permission, ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosPermissionResponse(response)) + .single(); } /** @@ -137,13 +126,10 @@ public Mono upsertPermission( if (options == null) { options = new CosmosPermissionRequestOptions(); } - - if (!database.getClient().getTracerProvider().isEnabled()) { - return upsertPermissionInternal(permission, options); - } - - final CosmosPermissionRequestOptions requestOptions = options; - return withContext(context -> upsertPermissionInternal(permission, requestOptions, context)); + return database.getDocClientWrapper() + .upsertPermission(getLink(), permission, ModelBridgeInternal.toRequestOptions(options)) + .map(response -> ModelBridgeInternal.createCosmosPermissionResponse(response)) + .single(); } @@ -174,18 +160,13 @@ public CosmosPagedFlux readAllPermissions() { */ CosmosPagedFlux readAllPermissions(CosmosQueryRequestOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "readAllPermissions." + this.getId(); - pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), - spanName, - this.getDatabase().getClient().getServiceEndpoint(), - this.getDatabase().getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDatabase().getDocClientWrapper() .readPermissions(getLink(), options) .map(response -> BridgeInternal.createFeedResponse( ModelBridgeInternal.getCosmosPermissionPropertiesFromResults(response.getResults()), response.getResponseHeaders())); - }, this.getDatabase().getClient().getTracerProvider().isEnabled()); + }); } /** @@ -217,18 +198,13 @@ public CosmosPagedFlux queryPermissions(String query */ public CosmosPagedFlux queryPermissions(String query, CosmosQueryRequestOptions options) { return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - String spanName = "queryPermissions." + this.getId(); - pagedFluxOptions.setTracerInformation(this.getDatabase().getClient().getTracerProvider(), - spanName, - this.getDatabase().getClient().getServiceEndpoint(), - this.getDatabase().getId()); setContinuationTokenAndMaxItemCount(pagedFluxOptions, options); return getDatabase().getDocClientWrapper() .queryPermissions(getLink(), query, options) .map(response -> BridgeInternal.createFeedResponse( ModelBridgeInternal.getCosmosPermissionPropertiesFromResults(response.getResults()), response.getResponseHeaders())); - }, this.getDatabase().getClient().getTracerProvider().isEnabled()); + }); } /** @@ -267,91 +243,4 @@ String getLink() { CosmosAsyncDatabase getDatabase() { return database; } - - 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()); - } - - private Mono readInternal() { - return this.database.getDocClientWrapper() - .readUser(getLink(), null) - .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); - } - - private Mono replaceInternal(CosmosUserProperties userSettings, Context context) { - String spanName = "replaceUser." + getId(); - Mono responseMono = replaceInternal(userSettings); - return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, - database.getId(), - database.getClient().getServiceEndpoint()); - } - - private Mono replaceInternal(CosmosUserProperties userSettings) { - return this.database.getDocClientWrapper() - .replaceUser(ModelBridgeInternal.getV2User(userSettings), null) - .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); - } - - private Mono deleteInternal(Context context) { - String spanName = "deleteUser." + getId(); - Mono responseMono = deleteInternal(); - return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, - database.getId(), - database.getClient().getServiceEndpoint()); - } - - private Mono deleteInternal() { - return this.database.getDocClientWrapper() - .deleteUser(getLink(), null) - .map(response -> ModelBridgeInternal.createCosmosUserResponse(response)).single(); - } - - private Mono createPermissionInternal( - Permission permission, - CosmosPermissionRequestOptions options, - Context context) { - String spanName = "createPermission." + getId(); - Mono responseMono = createPermissionInternal(permission, options); - return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, - database.getId(), - database.getClient().getServiceEndpoint()); - } - - private Mono createPermissionInternal( - Permission permission, - CosmosPermissionRequestOptions options) { - return database.getDocClientWrapper() - .createPermission(getLink(), permission, ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosPermissionResponse(response)) - .single(); - } - - private Mono upsertPermissionInternal( - Permission permission, - CosmosPermissionRequestOptions options, - Context context) { - String spanName = "upsertPermission." + getId(); - Mono responseMono = upsertPermissionInternal(permission, options); - return database.getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context, - spanName, - database.getId(), - database.getClient().getServiceEndpoint()); - } - - private Mono upsertPermissionInternal( - Permission permission, - CosmosPermissionRequestOptions options) { - return database.getDocClientWrapper() - .upsertPermission(getLink(), permission, ModelBridgeInternal.toRequestOptions(options)) - .map(response -> ModelBridgeInternal.createCosmosPermissionResponse(response)) - .single(); - } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncUserDefinedFunction.java index 45b6803fdaac..7710818dd0fb 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,7 +2,6 @@ // 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.UserDefinedFunction; import com.azure.cosmos.models.CosmosUserDefinedFunctionResponse; @@ -10,8 +9,6 @@ import com.azure.cosmos.models.ModelBridgeInternal; import reactor.core.publisher.Mono; -import static com.azure.core.util.FluxUtil.withContext; - /** * The type Cosmos async user defined function. */ @@ -57,11 +54,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() { - if(!container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return readInternal(); - } - - return withContext(context -> readInternal(context)); + return container.getDatabase().getDocClientWrapper().readUserDefinedFunction(getLink(), null) + .map(response -> ModelBridgeInternal.createCosmosUserDefinedFunctionResponse(response)).single(); } /** @@ -77,11 +71,12 @@ public Mono read() { * or an error. */ public Mono replace(CosmosUserDefinedFunctionProperties udfSettings) { - if(!container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return replaceInternal(udfSettings); - } - - return withContext(context -> replaceInternal(udfSettings, context)); + return container.getDatabase() + .getDocClientWrapper() + .replaceUserDefinedFunction(new UserDefinedFunction(ModelBridgeInternal.toJsonFromJsonSerializable( + ModelBridgeInternal.getResource(udfSettings))), null) + .map(response -> ModelBridgeInternal.createCosmosUserDefinedFunctionResponse(response)) + .single(); } /** @@ -96,11 +91,11 @@ public Mono replace(CosmosUserDefinedFunction * an error. */ public Mono delete() { - if(!container.getDatabase().getClient().getTracerProvider().isEnabled()) { - return deleteInternal(); - } - - return withContext(context -> deleteInternal(context)); + return container.getDatabase() + .getDocClientWrapper() + .deleteUserDefinedFunction(this.getLink(), null) + .map(response -> ModelBridgeInternal.createCosmosUserDefinedFunctionResponse(response)) + .single(); } String getURIPathSegment() { @@ -120,57 +115,4 @@ String getLink() { builder.append(getId()); return builder.toString(); } - - 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()); - } - - private Mono readInternal() { - return container.getDatabase().getDocClientWrapper().readUserDefinedFunction(getLink(), null) - .map(response -> ModelBridgeInternal.createCosmosUserDefinedFunctionResponse(response)).single(); - } - - private Mono replaceInternal(CosmosUserDefinedFunctionProperties udfSettings, - Context context) { - String spanName = "replaceUDF." + container.getId(); - Mono responseMono = replaceInternal(udfSettings); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, - spanName, - container.getDatabase().getId(), - container.getDatabase().getClient().getServiceEndpoint()); - } - - private Mono replaceInternal(CosmosUserDefinedFunctionProperties udfSettings) { - return container.getDatabase() - .getDocClientWrapper() - .replaceUserDefinedFunction(new UserDefinedFunction(ModelBridgeInternal.toJsonFromJsonSerializable( - ModelBridgeInternal.getResource(udfSettings))), null) - .map(response -> ModelBridgeInternal.createCosmosUserDefinedFunctionResponse(response)) - .single(); - } - - private Mono deleteInternal(Context context) { - String spanName = "deleteUDF." + container.getId(); - Mono responseMono = deleteInternal(); - return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, - context, - spanName, - container.getDatabase().getId(), - container.getDatabase().getClient().getServiceEndpoint()); - } - - private Mono deleteInternal() { - return container.getDatabase() - .getDocClientWrapper() - .deleteUserDefinedFunction(this.getLink(), null) - .map(response -> ModelBridgeInternal.createCosmosUserDefinedFunctionResponse(response)) - .single(); - } } 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 28b4054180dc..c4407b68bfee 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,8 +5,6 @@ import com.azure.cosmos.util.CosmosPagedFlux; -import java.util.Map; - /** * Specifies paging options for Cosmos Paged Flux implementation. * @see CosmosPagedFlux @@ -15,11 +13,6 @@ public class CosmosPagedFluxOptions { private String requestContinuation; private Integer maxItemCount; - private TracerProvider tracerProvider; - private String tracerSpanName; - private String databaseId; - private String serviceEndpoint; - public CosmosPagedFluxOptions() {} @@ -64,43 +57,4 @@ public CosmosPagedFluxOptions setMaxItemCount(Integer maxItemCount) { this.maxItemCount = maxItemCount; return this; } - - /** - * Gets the tracer provider - * @return tracerProvider - */ - public TracerProvider getTracerProvider() { - return this.tracerProvider; - } - - /** - * Gets the tracer span name - * @return tracerSpanName - */ - public String getTracerSpanName() { - return 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/HttpConstants.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java index 70df4dfc7a0f..b14fce7394f8 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java @@ -265,7 +265,6 @@ public static class Versions { } public static class StatusCodes { - public static final int OK = 200; public static final int NOT_MODIFIED = 304; // Client error public static final int MINIMUM_STATUSCODE_AS_ERROR_GATEWAY = 400; 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 deleted file mode 100644 index 677e87db4e2f..000000000000 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java +++ /dev/null @@ -1,157 +0,0 @@ -// 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.tracing.Tracer; -import com.azure.cosmos.CosmosException; -import com.azure.cosmos.models.CosmosItemResponse; -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.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 Tracer tracer; - 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"; - public static final String DB_STATEMENT = "db.statement"; - public static final String ERROR_MSG = "error.msg"; - public static final String ERROR_TYPE = "error.type"; - public static final String ERROR_STACK = "error.stack"; - public static final String COSMOS_CALL_DEPTH = "cosmosCallDepth"; - public static final String COSMOS_CALL_DEPTH_VAL = "nested"; - public static final int ERROR_CODE = 0; - public static final String RESOURCE_PROVIDER_NAME = "Microsoft.DocumentDB"; - - public TracerProvider(Iterable tracers) { - Objects.requireNonNull(tracers, "'tracers' cannot be null."); - if (tracers.iterator().hasNext()) { - tracer = tracers.iterator().next(); - } - } - - public boolean isEnabled() { - return tracer != null; - } - - /** - * 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, 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(TracerProvider.DB_TYPE, DB_TYPE_VALUE, local); - tracer.setAttribute(TracerProvider.DB_URL, endpoint, local); - tracer.setAttribute(TracerProvider.DB_STATEMENT, methodName, local); - 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, int statusCode) { - Objects.requireNonNull(context, "'context' cannot be null."); - Objects.requireNonNull(signal, "'signal' cannot be null."); - - switch (signal.getType()) { - case ON_COMPLETE: - end(statusCode, null, context); - break; - case ON_ERROR: - Throwable throwable = null; - if (signal.hasError()) { - // The last status available is on error, this contains the thrown error. - throwable = signal.getThrowable(); - - if (throwable instanceof CosmosException) { - CosmosException exception = (CosmosException) throwable; - statusCode = exception.getStatusCode(); - } - } - end(statusCode, throwable, context); - break; - default: - // ON_SUBSCRIBE and ON_NEXT don't have the information to end the span so just return. - break; - } - } - - public > Mono traceEnabledCosmosResponsePublisher(Mono resultPublisher, - Context context, - String spanName, - String databaseId, - String endpoint) { - return traceEnabledPublisher(resultPublisher, context, spanName,databaseId, endpoint, - (T response) -> response.getStatusCode()); - } - - public Mono> traceEnabledCosmosItemResponsePublisher(Mono> resultPublisher, - Context context, - String spanName, - String databaseId, - String endpoint) { - return traceEnabledPublisher(resultPublisher, context, spanName,databaseId, endpoint, - CosmosItemResponse::getStatusCode); - } - - public Mono traceEnabledPublisher(Mono resultPublisher, - Context context, - String spanName, - String databaseId, - String endpoint, - Function statusCodeFunc) { - final AtomicReference parentContext = new AtomicReference<>(Context.NONE); - Optional callDepth = context.getData(COSMOS_CALL_DEPTH); - final boolean isNestedCall = callDepth.isPresent(); - return resultPublisher - .doOnSubscribe(ignoredValue -> { - if (!isNestedCall) { - parentContext.set(this.startSpan(spanName, databaseId, endpoint, - context)); - } - }).doOnSuccess(response -> { - if (!isNestedCall) { - this.endSpan(parentContext.get(), Signal.complete(), statusCodeFunc.apply(response)); - } - }).doOnError(throwable -> { - if (!isNestedCall) { - this.endSpan(parentContext.get(), Signal.error(throwable), ERROR_CODE); - } - }); - } - - 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); - } - 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 77167bc5970d..ca9f20352f3a 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,19 +3,13 @@ package com.azure.cosmos.util; -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.cosmos.implementation.CosmosPagedFluxOptions; -import com.azure.cosmos.implementation.HttpConstants; -import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.models.FeedResponse; import reactor.core.CoreSubscriber; import reactor.core.publisher.Flux; -import reactor.core.publisher.Signal; -import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; /** @@ -35,44 +29,32 @@ public final class CosmosPagedFlux extends ContinuablePagedFlux> { private final Function>> optionsFluxFunction; - private final boolean isTracerEnabled; - CosmosPagedFlux(Function>> optionsFluxFunction, - boolean isTracerEnable) { + CosmosPagedFlux(Function>> optionsFluxFunction) { this.optionsFluxFunction = optionsFluxFunction; - this.isTracerEnabled = isTracerEnable; } @Override public Flux> byPage() { CosmosPagedFluxOptions cosmosPagedFluxOptions = new CosmosPagedFluxOptions(); - if (!this.isTracerEnabled) { - return this.optionsFluxFunction.apply(cosmosPagedFluxOptions); - } - return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)); + return this.optionsFluxFunction.apply(cosmosPagedFluxOptions); } @Override public Flux> byPage(String continuationToken) { CosmosPagedFluxOptions cosmosPagedFluxOptions = new CosmosPagedFluxOptions(); cosmosPagedFluxOptions.setRequestContinuation(continuationToken); - if (!this.isTracerEnabled) { - return this.optionsFluxFunction.apply(cosmosPagedFluxOptions); - } - return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)); + return this.optionsFluxFunction.apply(cosmosPagedFluxOptions); } @Override public Flux> byPage(int preferredPageSize) { CosmosPagedFluxOptions cosmosPagedFluxOptions = new CosmosPagedFluxOptions(); cosmosPagedFluxOptions.setMaxItemCount(preferredPageSize); - if (!this.isTracerEnabled) { - return this.optionsFluxFunction.apply(cosmosPagedFluxOptions); - } - return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)); + return this.optionsFluxFunction.apply(cosmosPagedFluxOptions); } @Override @@ -80,11 +62,8 @@ public Flux> byPage(String continuationToken, int preferredPageS CosmosPagedFluxOptions cosmosPagedFluxOptions = new CosmosPagedFluxOptions(); cosmosPagedFluxOptions.setRequestContinuation(continuationToken); cosmosPagedFluxOptions.setMaxItemCount(preferredPageSize); - if (!this.isTracerEnabled) { - return this.optionsFluxFunction.apply(cosmosPagedFluxOptions); - } - return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)); + return this.optionsFluxFunction.apply(cosmosPagedFluxOptions); } /** @@ -105,19 +84,4 @@ 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 -> { - parentContext.set(pagedFluxOptions.getTracerProvider().startSpan(pagedFluxOptions.getTracerSpanName(), - pagedFluxOptions.getDatabaseId(), pagedFluxOptions.getServiceEndpoint(), - context)); - }).doOnComplete(() -> { - pagedFluxOptions.getTracerProvider().endSpan(parentContext.get(), Signal.complete(), - HttpConstants.StatusCodes.OK); - }).doOnError(throwable -> { - pagedFluxOptions.getTracerProvider().endSpan(parentContext.get(), Signal.error(throwable), - TracerProvider.ERROR_CODE); - }); - } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/UtilBridgeInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/UtilBridgeInternal.java index 141797b27ed3..1282cd7a56de 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/UtilBridgeInternal.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/UtilBridgeInternal.java @@ -24,11 +24,10 @@ public final class UtilBridgeInternal { private UtilBridgeInternal() {} @Warning(value = INTERNAL_USE_ONLY_WARNING) - public static CosmosPagedFlux createCosmosPagedFlux(Function>> pagedFluxOptionsFluxFunction, boolean isTracerEnabled) { - return new CosmosPagedFlux<>(pagedFluxOptionsFluxFunction, isTracerEnabled); + public static CosmosPagedFlux createCosmosPagedFlux(Function>> pagedFluxOptionsFluxFunction) { + return new CosmosPagedFlux<>(pagedFluxOptionsFluxFunction); } - @Warning(value = INTERNAL_USE_ONLY_WARNING) public static CosmosPagedIterable createCosmosPagedIterable(ContinuablePagedFlux> pagedFlux) { return new CosmosPagedIterable<>(pagedFlux); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/module-info.java b/sdk/cosmos/azure-cosmos/src/main/java/module-info.java index 71c5f58fb735..553cca6bf71f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/module-info.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/module-info.java @@ -49,5 +49,4 @@ opens com.azure.cosmos.util 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 deleted file mode 100644 index d7ec6300e224..000000000000 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosTracerTest.java +++ /dev/null @@ -1,245 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. -package com.azure.cosmos; - -import com.azure.core.util.Context; -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.CosmosItemRequestOptions; -import com.azure.cosmos.models.CosmosQueryRequestOptions; -import com.azure.cosmos.models.CosmosStoredProcedureProperties; -import com.azure.cosmos.models.CosmosTriggerProperties; -import com.azure.cosmos.models.CosmosUserDefinedFunctionProperties; -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 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() { - client = new CosmosClientBuilder() - .endpoint(TestConfigurations.HOST) - .key(TestConfigurations.MASTER_KEY) - .directMode(DirectConnectionConfig.getDefaultConfig()) - .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); - - client.createDatabaseIfNotExists(cosmosAsyncDatabase.getId()).block(); - Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), - Matchers.anyString(), Matchers.any(Context.class)); - - client.readAllDatabases(new CosmosQueryRequestOptions()).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), - Matchers.anyString(), Matchers.any(Context.class)); - - String query = "select * from c where c.id = '" + cosmosAsyncDatabase.getId() + "'"; - client.queryDatabases(query, new CosmosQueryRequestOptions()).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(3)).startSpan(Matchers.anyString(), Matchers.anyString(), - Matchers.anyString(), Matchers.any(Context.class)); - } - - @Test(groups = {"emulator"}, timeOut = TIMEOUT) - public void cosmosAsyncDatabase() { - TracerProvider tracer = Mockito.spy(new TracerProvider(ServiceLoader.load(Tracer.class))); - ReflectionUtils.setTracerProvider(client, tracer); - - cosmosAsyncDatabase.createContainerIfNotExists(cosmosAsyncContainer.getId(), - "/pk", 5000).block(); - Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), - Matchers.anyString(), Matchers.any(Context.class)); - - try { - cosmosAsyncDatabase.readThroughput().block(); - } catch (CosmosException ex) { - //do nothing - } - - 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.anyString(), - Matchers.anyString(), Matchers.any(Context.class)); - - cosmosAsyncDatabase.readAllContainers().byPage().single().block(); - Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.anyString(), - Matchers.anyString(), Matchers.any(Context.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.anyString(), - Matchers.anyString(), Matchers.any(Context.class)); - - try { - cosmosAsyncContainer.readThroughput().block(); - } catch (CosmosException ex) { - //do nothing - } - 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.anyString(), - Matchers.anyString(), Matchers.any(Context.class)); - - cosmosAsyncContainer.upsertItem(properties, - new CosmosItemRequestOptions()).block(); - Mockito.verify(tracer, Mockito.times(4)).startSpan(Matchers.anyString(), Matchers.anyString(), - Matchers.anyString(), Matchers.any(Context.class)); - - cosmosAsyncContainer.readItem(ITEM_ID, PartitionKey.NONE, - CosmosItemProperties.class).block(); - 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.anyString(), - Matchers.anyString(), Matchers.any(Context.class)); - cosmosAsyncContainer.readAllItems(new CosmosQueryRequestOptions(), CosmosItemRequestOptions.class).byPage().single().block(); - 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 CosmosQueryRequestOptions(), CosmosItemRequestOptions.class).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(8)).startSpan(Matchers.anyString(), Matchers.anyString(), - Matchers.anyString(), Matchers.any(Context.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 CosmosQueryRequestOptions()).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(1)).startSpan(Matchers.anyString(), Matchers.anyString(), - Matchers.anyString(), Matchers.any(Context.class)); - - cosmosAsyncContainer.getScripts().readAllTriggers(new CosmosQueryRequestOptions()).byPage().single().block(); - Mockito.verify(tracer, Mockito.times(2)).startSpan(Matchers.anyString(), Matchers.anyString(), - Matchers.anyString(), Matchers.any(Context.class)); - - cosmosAsyncContainer.getScripts().readAllUserDefinedFunctions(new CosmosQueryRequestOptions()).byPage().single().block(); - 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.anyString(), - Matchers.anyString(), Matchers.any(Context.class)); - - cosmosAsyncContainer.getScripts().getUserDefinedFunction(cosmosUserDefinedFunctionProperties.getId()).read().block(); - 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.anyString(), - Matchers.anyString(), Matchers.any(Context.class)); - - cosmosAsyncContainer.getScripts().readAllUserDefinedFunctions(new CosmosQueryRequestOptions()).byPage().single().block(); - 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.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.anyString(), - Matchers.anyString(), Matchers.any(Context.class)); - - cosmosAsyncContainer.getScripts().getTrigger(cosmosTriggerProperties.getId()).read().block(); - 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.anyString(), - Matchers.anyString(), Matchers.any(Context.class)); - cosmosAsyncContainer.getScripts().readAllTriggers(new CosmosQueryRequestOptions()).byPage().single().block(); - 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.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.anyString(), - Matchers.anyString(), Matchers.any(Context.class)); - - cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).read().block(); - 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.anyString(), - Matchers.anyString(), Matchers.any(Context.class)); - - cosmosAsyncContainer.getScripts().readAllStoredProcedures(new CosmosQueryRequestOptions()).byPage().single().block(); - - cosmosAsyncContainer.getScripts().getStoredProcedure(procedureProperties.getId()).delete().block(); - Mockito.verify(tracer, Mockito.times(18)).startSpan(Matchers.anyString(), Matchers.anyString(), - Matchers.anyString(), Matchers.any(Context.class)); - } - - @AfterClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT) - public void afterClass() { - LifeCycleUtils.closeQuietly(client); - } - - private static CosmosUserDefinedFunctionProperties getCosmosUserDefinedFunctionProperties() { - CosmosUserDefinedFunctionProperties udf = new CosmosUserDefinedFunctionProperties(UUID.randomUUID().toString(), "function() {var x = 10;}"); - return udf; - } - - private static CosmosTriggerProperties getCosmosTriggerProperties() { - CosmosTriggerProperties trigger = new CosmosTriggerProperties(UUID.randomUUID().toString(), "function() {var x = 10;}"); - trigger.setTriggerOperation(TriggerOperation.CREATE); - trigger.setTriggerType(TriggerType.PRE); - return trigger; - } - - private static CosmosStoredProcedureProperties getCosmosStoredProcedureProperties() { - CosmosStoredProcedureProperties storedProcedureDef = new CosmosStoredProcedureProperties(UUID.randomUUID().toString(), "function() {var x = 10;}"); - 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 5a9063584e37..9dbdc312ffcd 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,7 +9,6 @@ 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; @@ -105,8 +104,4 @@ 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"); - } } diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/ReadFeedExceptionHandlingTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/ReadFeedExceptionHandlingTest.java index a812c50cecd7..b9424abd3d32 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,14 +2,12 @@ // 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.implementation.TracerProvider; +import com.azure.cosmos.util.CosmosPagedFlux; import com.azure.cosmos.models.CosmosDatabaseProperties; 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; @@ -20,7 +18,6 @@ import reactor.core.publisher.Flux; import java.util.ArrayList; -import java.util.ServiceLoader; import static org.assertj.core.api.Assertions.assertThat; @@ -49,10 +46,7 @@ public void readFeedException() throws Exception { .mergeWith(Flux.fromIterable(frps)); final CosmosAsyncClientWrapper mockedClientWrapper = Mockito.spy(new CosmosAsyncClientWrapper(client)); - Mockito.when(mockedClientWrapper.readAllDatabases()).thenReturn(UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> { - pagedFluxOptions.setTracerInformation(new TracerProvider(ServiceLoader.load(Tracer.class)), "testSpan", "testEndpoint,", "testDb"); - return response; - }, false)); + Mockito.when(mockedClientWrapper.readAllDatabases()).thenReturn(UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> response)); TestSubscriber> subscriber = new TestSubscriber<>(); mockedClientWrapper.readAllDatabases().byPage().subscribe(subscriber); assertThat(subscriber.valueCount()).isEqualTo(2); 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 3224cc519b6d..f6a1a5a1a2bd 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/TestSuiteBase.java @@ -484,7 +484,7 @@ public Flux> bulkInsert(CosmosAsyncContainer cosmosCon return Flux.merge(Flux.fromIterable(result), concurrencyLevel); } public List bulkInsertBlocking(CosmosAsyncContainer cosmosContainer, - List documentDefinitionList) { + List documentDefinitionList) { return bulkInsert(cosmosContainer, documentDefinitionList, DEFAULT_BULK_INSERT_CONCURRENCY_LEVEL) .publishOn(Schedulers.parallel()) .map(itemResponse -> itemResponse.getItem()) @@ -813,12 +813,12 @@ public void validateItemFailure( } public void validateQuerySuccess(Flux> flowable, - FeedResponseListValidator validator) { + FeedResponseListValidator validator) { validateQuerySuccess(flowable, validator, subscriberValidationTimeout); } public static void validateQuerySuccess(Flux> flowable, - FeedResponseListValidator validator, long timeout) { + FeedResponseListValidator validator, long timeout) { TestSubscriber> testSubscriber = new TestSubscriber<>(); @@ -834,7 +834,7 @@ public void validateQueryFailure(Flux> flowable, FailureVali } public static void validateQueryFailure(Flux> flowable, - FailureValidator validator, long timeout) { + FailureValidator validator, long timeout) { TestSubscriber> testSubscriber = new TestSubscriber<>();