From 0f9318b1c84f9ddf8fdfb8dc33b2fc03124c4b52 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 7 Jan 2025 19:05:43 +0000 Subject: [PATCH 01/11] Initial draft --- .../com/azure/cosmos/CosmosAsyncClient.java | 14 +++++++++++ .../azure/cosmos/CosmosAsyncContainer.java | 4 +++ .../com/azure/cosmos/CosmosAsyncDatabase.java | 11 +++++++- .../java/com/azure/cosmos/CosmosClient.java | 1 + .../com/azure/cosmos/CosmosClientBuilder.java | 25 +++++++++++++++++++ 5 files changed, 54 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index 48f631b4af3c..5c86d1ac1404 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 @@ -30,6 +30,7 @@ import com.azure.cosmos.implementation.clienttelemetry.TagName; import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdMetrics; import com.azure.cosmos.implementation.faultinjection.IFaultInjectorProvider; +import com.azure.cosmos.implementation.guava25.base.Function; import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; import com.azure.cosmos.models.CosmosAuthorizationTokenResolver; import com.azure.cosmos.models.CosmosClientTelemetryConfig; @@ -89,6 +90,9 @@ public final class CosmosAsyncClient implements Closeable { .CosmosClientTelemetryConfigHelper .getCosmosClientTelemetryConfigAccessor(); + private final static Function DEFAULT_CONTAINER_FACTORY = + (originalContainer) -> originalContainer; + private final AsyncDocumentClient asyncDocumentClient; private final String serviceEndpoint; private final ConnectionPolicy connectionPolicy; @@ -107,6 +111,7 @@ public final class CosmosAsyncClient implements Closeable { private final WriteRetryPolicy nonIdempotentWriteRetryPolicy; private final List requestPolicies; private final CosmosItemSerializer defaultCustomSerializer; + private final java.util.function.Function containerFactory; CosmosAsyncClient(CosmosClientBuilder builder) { // Async Cosmos client wrapper @@ -125,6 +130,11 @@ public final class CosmosAsyncClient implements Closeable { this.nonIdempotentWriteRetryPolicy = builder.getNonIdempotentWriteRetryPolicy(); this.requestPolicies = builder.getOperationPolicies(); this.defaultCustomSerializer = builder.getCustomItemSerializer(); + if (builder.containerCreationInterceptor() != null) { + this.containerFactory = builder.containerCreationInterceptor(); + } else { + this.containerFactory = DEFAULT_CONTAINER_FACTORY; + } CosmosEndToEndOperationLatencyPolicyConfig endToEndOperationLatencyPolicyConfig = builder.getEndToEndOperationConfig(); SessionRetryOptions sessionRetryOptions = builder.getSessionRetryOptions(); @@ -809,6 +819,10 @@ String getUserAgent() { return this.asyncDocumentClient.getUserAgent(); } + java.util.function.Function getContainerCreationInterceptor() { + return this.containerFactory; + } + /////////////////////////////////////////////////////////////////////////////////////////// // the following helper/accessor only helps to access this class outside of this package.// /////////////////////////////////////////////////////////////////////////////////////////// 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 188a9fd7e468..765c65187b3f 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 @@ -146,6 +146,10 @@ public class CosmosAsyncContainer { private CosmosAsyncScripts scripts; private IFaultInjectorProvider faultInjectorProvider; + protected CosmosAsyncContainer(CosmosAsyncContainer toBeWrappedContainer) { + this(toBeWrappedContainer.getId(), toBeWrappedContainer.getDatabase()); + } + CosmosAsyncContainer(String id, CosmosAsyncDatabase database) { this.id = id; this.database = database; 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 3ce76b08602f..1e808cbecd53 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 @@ -815,7 +815,16 @@ public CosmosPagedFlux queryContainers(SqlQuerySpec q * @return Cosmos Container */ public CosmosAsyncContainer getContainer(String id) { - return new CosmosAsyncContainer(id, this); + CosmosAsyncContainer asyncContainer = this + .client + .getContainerCreationInterceptor() + .apply(new CosmosAsyncContainer(id, this)); + + if (asyncContainer == null) { + throw new IllegalStateException( + "The implementation of teh custom container creation interceptor must not return null."); + } + return asyncContainer; } /** diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java index 9a8b2e7e7bb4..ba5adcb057e7 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java @@ -5,6 +5,7 @@ import com.azure.core.annotation.ServiceClient; import com.azure.cosmos.implementation.ImplementationBridgeHelpers; +import com.azure.cosmos.implementation.guava25.base.Function; import com.azure.cosmos.models.CosmosContainerIdentity; import com.azure.cosmos.models.CosmosDatabaseProperties; import com.azure.cosmos.models.CosmosDatabaseRequestOptions; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java index 3be810c46c06..211cf19590b7 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java @@ -39,6 +39,7 @@ import java.util.Locale; import java.util.Objects; import java.util.Set; +import java.util.function.Function; import java.util.function.Supplier; import static com.azure.cosmos.implementation.ImplementationBridgeHelpers.CosmosClientBuilderHelper; @@ -152,6 +153,8 @@ public class CosmosClientBuilder implements private boolean isRegionScopedSessionCapturingEnabled = false; private boolean serverCertValidationDisabled = false; + private Function containerFactory = null; + /** * Instantiates a new Cosmos client builder. */ @@ -175,6 +178,28 @@ CosmosClientMetadataCachesSnapshot metadataCaches() { return this.state; } + /** + * Gets the container creation interceptor. + * @return the function that should be invoked to allow wrapping containers - or null if no interceptor is defined. + */ + Function containerCreationInterceptor() { + return this.containerFactory; + } + + /** + * Sets a function that allows intercepting container creation - for example to wrap the original + * Cosmos(Async)Container in an extended custom class to add diagnostics, custom validations or behavior. + * @param factory - the factory method allowing to wrap the original container in a custom class. + * @return current {@link CosmosClientBuilder} + */ + public CosmosClientBuilder containerCreationInterceptor( + Function factory) { + + this.containerFactory = factory; + + return this; + } + /** * Sets a {@code boolean} flag to reduce the frequency of retries when the client * strives to meet Session Consistency guarantees for operations From b4e42f6e00c60d09eb8ae0f522016f5f23f4ae82 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 7 Jan 2025 23:55:51 +0000 Subject: [PATCH 02/11] Added unit test --- .../azure/cosmos/CosmosClientBuilderTest.java | 68 +++++++++++++++++++ .../com/azure/cosmos/CosmosClientBuilder.java | 2 +- 2 files changed, 69 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java index 30f509e2ef24..503669653eb7 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java @@ -11,6 +11,9 @@ import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils; import com.azure.cosmos.models.CosmosClientTelemetryConfig; +import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.SqlQuerySpec; +import com.azure.cosmos.util.CosmosPagedFlux; import org.testng.SkipException; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; @@ -268,4 +271,69 @@ public void validateSessionTokenCapturingForAccountDefaultConsistencyWithEnvVari System.clearProperty("COSMOS.SESSION_CAPTURING_TYPE"); } } + + @Test(groups = "unit") + public void validateContainerCreationInterceptor() { + CosmosClient clientWithoutInterceptor = new CosmosClientBuilder() + .endpoint(TestConfigurations.HOST) + .key(TestConfigurations.MASTER_KEY) + .userAgentSuffix("noInterceptor") + .buildClient(); + + CosmosClient clientWithInterceptor = new CosmosClientBuilder() + .endpoint(TestConfigurations.HOST) + .key(TestConfigurations.MASTER_KEY) + .userAgentSuffix("withInterceptor") + .containerCreationInterceptor(originalContainer -> new DisallowQueriesContainer(originalContainer)) + .buildClient(); + + CosmosAsyncClient asyncClientWithInterceptor = new CosmosClientBuilder() + .endpoint(TestConfigurations.HOST) + .key(TestConfigurations.MASTER_KEY) + .userAgentSuffix("withInterceptor") + .containerCreationInterceptor(originalContainer -> new DisallowQueriesContainer(originalContainer)) + .buildAsyncClient(); + + CosmosContainer normalContainer = clientWithoutInterceptor + .getDatabase("TestDB") + .getContainer("TestContainer"); + assertThat(normalContainer).isNotNull(); + assertThat(normalContainer.getClass()).isEqualTo(CosmosContainer.class); + assertThat(normalContainer.asyncContainer.getClass()).isEqualTo(CosmosAsyncContainer.class); + + CosmosContainer customSyncContainer = clientWithInterceptor + .getDatabase("TestDB") + .getContainer("TestContainer"); + assertThat(customSyncContainer).isNotNull(); + assertThat(customSyncContainer.getClass()).isEqualTo(CosmosContainer.class); + assertThat(customSyncContainer.asyncContainer.getClass()).isEqualTo(DisallowQueriesContainer.class); + + CosmosAsyncContainer customAsyncContainer = asyncClientWithInterceptor + .getDatabase("TestDB") + .getContainer("TestContainer"); + assertThat(customAsyncContainer).isNotNull(); + assertThat(customAsyncContainer.getClass()).isEqualTo(DisallowQueriesContainer.class); + } + + private static class DisallowQueriesContainer extends CosmosAsyncContainer { + + protected DisallowQueriesContainer(CosmosAsyncContainer toBeWrappedContainer) { + super(toBeWrappedContainer); + } + + @Override + public CosmosPagedFlux queryItems(String query, CosmosQueryRequestOptions options, Class classType) { + throw new IllegalStateException("No queries allowed. Use point reads or readMany instead."); + } + + @Override + public CosmosPagedFlux queryItems(SqlQuerySpec querySpec, Class classType) { + throw new IllegalStateException("No queries allowed. Use point reads or readMany instead."); + } + + @Override + public CosmosPagedFlux queryItems(String query, Class classType) { + throw new IllegalStateException("No queries allowed. Use point reads or readMany instead."); + } + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java index 211cf19590b7..dfd189726f4e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java @@ -188,7 +188,7 @@ Function containerCreationIntercepto /** * Sets a function that allows intercepting container creation - for example to wrap the original - * Cosmos(Async)Container in an extended custom class to add diagnostics, custom validations or behavior. + * CosmosAsyncContainer in an extended custom class to add diagnostics, custom validations or behavior. * @param factory - the factory method allowing to wrap the original container in a custom class. * @return current {@link CosmosClientBuilder} */ From 3e268cc249f3d5d66afdfdd5499cfbda1bedd685 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 8 Jan 2025 16:33:49 +0000 Subject: [PATCH 03/11] Update CosmosClientBuilderTest.java --- .../src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java index 503669653eb7..cedc1a338962 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java @@ -240,7 +240,7 @@ public void validateSessionTokenCapturingForAccountDefaultConsistency(boolean sh } // set env variable to COSMOS.SESSION_CAPTURING_TYPE to REGION_SCOPED to test all possible assertions - @Test(groups = "unit", enabled = false) + @Test(groups = "emulator", enabled = false) public void validateSessionTokenCapturingForAccountDefaultConsistencyWithEnvVariable() { try { From 9f06c6398e05c76a6773aa32aad9fa5802699e72 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 8 Jan 2025 17:24:15 +0000 Subject: [PATCH 04/11] Update CosmosClientBuilderTest.java --- .../test/java/com/azure/cosmos/CosmosClientBuilderTest.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java index cedc1a338962..c078240a4e93 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java @@ -240,7 +240,7 @@ public void validateSessionTokenCapturingForAccountDefaultConsistency(boolean sh } // set env variable to COSMOS.SESSION_CAPTURING_TYPE to REGION_SCOPED to test all possible assertions - @Test(groups = "emulator", enabled = false) + @Test(groups = "unit", enabled = false) public void validateSessionTokenCapturingForAccountDefaultConsistencyWithEnvVariable() { try { @@ -272,7 +272,7 @@ public void validateSessionTokenCapturingForAccountDefaultConsistencyWithEnvVari } } - @Test(groups = "unit") + @Test(groups = "emulator") public void validateContainerCreationInterceptor() { CosmosClient clientWithoutInterceptor = new CosmosClientBuilder() .endpoint(TestConfigurations.HOST) From d8a974117cdfb835bb967fc0987fd0edd64f2255 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 15 Apr 2025 23:30:40 +0000 Subject: [PATCH 05/11] Adding simple extensibility for CosmosPagedFlux --- .../azure/cosmos/CosmosClientBuilderTest.java | 169 +++++++++- .../ImplementationBridgeHelpers.java | 4 + .../com/azure/cosmos/models/FeedResponse.java | 13 + .../azure/cosmos/util/CosmosPagedFlux.java | 296 ++--------------- .../util/CosmosPagedFluxDefaultImpl.java | 306 ++++++++++++++++++ .../util/CosmosPagedFluxStaticListImpl.java | 132 ++++++++ .../azure/cosmos/util/UtilBridgeInternal.java | 4 +- 7 files changed, 636 insertions(+), 288 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFluxDefaultImpl.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFluxStaticListImpl.java diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java index e9ff0d750b36..b2fc6d1ea79b 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java @@ -10,18 +10,27 @@ import com.azure.cosmos.implementation.SessionContainer; import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils; -import com.azure.cosmos.models.CosmosClientTelemetryConfig; import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.SqlParameter; import com.azure.cosmos.models.SqlQuerySpec; import com.azure.cosmos.util.CosmosPagedFlux; +import com.fasterxml.jackson.databind.node.ObjectNode; import org.testng.SkipException; import org.testng.annotations.DataProvider; import org.testng.annotations.Test; import java.net.URISyntaxException; +import java.util.ArrayList; import java.util.Arrays; +import java.util.Comparator; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.stream.Collectors; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.fail; public class CosmosClientBuilderTest { String hostName = "https://sample-account.documents.azure.com:443/"; @@ -162,18 +171,20 @@ public void validateContainerCreationInterceptor() { .userAgentSuffix("noInterceptor") .buildClient(); + ConcurrentMap> queryCache = new ConcurrentHashMap<>(); + CosmosClient clientWithInterceptor = new CosmosClientBuilder() .endpoint(TestConfigurations.HOST) .key(TestConfigurations.MASTER_KEY) .userAgentSuffix("withInterceptor") - .containerCreationInterceptor(originalContainer -> new DisallowQueriesContainer(originalContainer)) + .containerCreationInterceptor(originalContainer -> new CacheAndValidateQueriesContainer(originalContainer, queryCache)) .buildClient(); CosmosAsyncClient asyncClientWithInterceptor = new CosmosClientBuilder() .endpoint(TestConfigurations.HOST) .key(TestConfigurations.MASTER_KEY) .userAgentSuffix("withInterceptor") - .containerCreationInterceptor(originalContainer -> new DisallowQueriesContainer(originalContainer)) + .containerCreationInterceptor(originalContainer -> new CacheAndValidateQueriesContainer(originalContainer, queryCache)) .buildAsyncClient(); CosmosContainer normalContainer = clientWithoutInterceptor @@ -188,34 +199,172 @@ public void validateContainerCreationInterceptor() { .getContainer("TestContainer"); assertThat(customSyncContainer).isNotNull(); assertThat(customSyncContainer.getClass()).isEqualTo(CosmosContainer.class); - assertThat(customSyncContainer.asyncContainer.getClass()).isEqualTo(DisallowQueriesContainer.class); + assertThat(customSyncContainer.asyncContainer.getClass()).isEqualTo(CacheAndValidateQueriesContainer.class); CosmosAsyncContainer customAsyncContainer = asyncClientWithInterceptor .getDatabase("TestDB") .getContainer("TestContainer"); assertThat(customAsyncContainer).isNotNull(); - assertThat(customAsyncContainer.getClass()).isEqualTo(DisallowQueriesContainer.class); + assertThat(customAsyncContainer.getClass()).isEqualTo(CacheAndValidateQueriesContainer.class); + + try { + customSyncContainer.queryItems("SELECT * from c", null, ObjectNode.class); + fail("Unparameterized query should throw"); + } catch (IllegalStateException expectedError) {} + + try { + customAsyncContainer.queryItems("SELECT * from c", null, ObjectNode.class); + fail("Unparameterized query should throw"); + } catch (IllegalStateException expectedError) {} + + try { + customAsyncContainer.queryItems("SELECT * from c", ObjectNode.class); + fail("Unparameterized query should throw"); + } catch (IllegalStateException expectedError) {} + + SqlQuerySpec querySpec = new SqlQuerySpec().setQueryText("SELECT * from c"); + assertThat(queryCache).size().isEqualTo(0); + + try { + List items = customSyncContainer + .queryItems(querySpec, null, ObjectNode.class) + .stream().collect(Collectors.toList()); + } catch (CosmosException cosmosException) { + // Container does not exist - when not cached should fail + assertThat(cosmosException.getStatusCode()).isEqualTo(404); + assertThat(cosmosException.getSubStatusCode()).isEqualTo(1003); + } + + queryCache.putIfAbsent(new CacheKey(ObjectNode.class.getCanonicalName(), querySpec), new ArrayList<>()); + assertThat(queryCache).size().isEqualTo(1); + + // Validate that CacheKey equality check works + queryCache.putIfAbsent(new CacheKey(ObjectNode.class.getCanonicalName(), querySpec), new ArrayList<>()); + assertThat(queryCache).size().isEqualTo(1); + + // Validate that form cache the results can be served + List items = customSyncContainer + .queryItems(querySpec, null, ObjectNode.class) + .stream().collect(Collectors.toList()); + + querySpec = new SqlQuerySpec().setQueryText("SELECT * from c"); + CosmosPagedFlux cachedPagedFlux = customAsyncContainer + .queryItems(querySpec, null, ObjectNode.class); + assertThat(cachedPagedFlux.getClass().getName()).startsWith("com.azure.cosmos.util.CosmosPagedFluxStaticListImpl"); + + // Validate that uncached query form async Container also fails with 404 due to non-existing Container + querySpec = new SqlQuerySpec().setQueryText("SELECT * from r"); + try { + CosmosPagedFlux uncachedPagedFlux = customAsyncContainer + .queryItems(querySpec, null, ObjectNode.class); + } catch (CosmosException cosmosException) { + assertThat(cosmosException.getStatusCode()).isEqualTo(404); + assertThat(cosmosException.getSubStatusCode()).isEqualTo(1003); + } + } + + private static class CacheKey { + private final String className; + private final String queryText; + + private final List parameters; + public CacheKey(String className, SqlQuerySpec querySpec) { + this.className = className; + this.queryText = querySpec.getQueryText(); + List tempParameters = querySpec.getParameters(); + if (tempParameters != null) { + tempParameters.sort(Comparator.comparing(SqlParameter::getName)); + this.parameters = tempParameters; + } else { + this.parameters = new ArrayList<>(); + } + } + + @Override + public int hashCode() { + Object[] temp = new Object[2 + this.parameters.size()]; + temp[0] = this.className; + temp[1] = this.queryText; + for (int i = 0; i < this.parameters.size(); i++) { + temp[2 + i] = this.parameters.get(i).getValue(Object.class); + } + + return Objects.hash(temp); + } + + @Override + public boolean equals(Object obj) { + if (obj == null) { + return false; + } + + if (!(obj instanceof CacheKey)) { + return false; + } + + CacheKey other = (CacheKey)obj; + if (!this.className.equals(other.className)) { + return false; + } + + if (!this.queryText.equals(other.queryText)) { + return false; + } + + if (this.parameters.size() != other.parameters.size()) { + return false; + } + + for (int i = 0; i < this.parameters.size(); i++) { + if (!this.parameters.get(i).getName().equals(other.parameters.get(i).getName())) { + return false; + } + + if (!this.parameters.get(i).getValue(Object.class).equals(other.parameters.get(i).getValue(Object.class))) { + return false; + } + } + + return true; + } } - private static class DisallowQueriesContainer extends CosmosAsyncContainer { + private static class CacheAndValidateQueriesContainer extends CosmosAsyncContainer { + private final ConcurrentMap> queryCache; + + protected CacheAndValidateQueriesContainer( + CosmosAsyncContainer toBeWrappedContainer, + ConcurrentMap> queryCache) { - protected DisallowQueriesContainer(CosmosAsyncContainer toBeWrappedContainer) { super(toBeWrappedContainer); + this.queryCache = queryCache; } @Override public CosmosPagedFlux queryItems(String query, CosmosQueryRequestOptions options, Class classType) { - throw new IllegalStateException("No queries allowed. Use point reads or readMany instead."); + throw new IllegalStateException("No unparameterized queries allowed. Use parameterized query instead."); } @Override public CosmosPagedFlux queryItems(SqlQuerySpec querySpec, Class classType) { - throw new IllegalStateException("No queries allowed. Use point reads or readMany instead."); + return this.queryItems(querySpec, null, classType); } @Override public CosmosPagedFlux queryItems(String query, Class classType) { - throw new IllegalStateException("No queries allowed. Use point reads or readMany instead."); + throw new IllegalStateException("No unparameterized queries allowed. Use parameterized query instead."); + } + + @Override + public CosmosPagedFlux queryItems(SqlQuerySpec querySpec, CosmosQueryRequestOptions options, Class classType) { + CacheKey key = new CacheKey(classType.getCanonicalName(), querySpec); + List cachedResult = this.queryCache.get(key); + if (cachedResult != null) { + return CosmosPagedFlux.createFromList((List)cachedResult, false); + } + + return super + .queryItems(querySpec, options, classType); } } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java index ebcfecd27216..f8e05f8d30a5 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java @@ -1095,6 +1095,10 @@ FeedResponse createFeedResponse(RxDocumentServiceResponse response, CosmosItemSerializer itemSerializer, Class cls); + FeedResponse createNonServiceFeedResponse(List items, + boolean isChangeFeed, + boolean isNoChanges); + FeedResponse createChangeFeedResponse(RxDocumentServiceResponse response, CosmosItemSerializer itemSerializer, Class cls); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java index 0978d5fb2f20..8c9bbef80c8e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/FeedResponse.java @@ -20,10 +20,12 @@ import com.azure.cosmos.implementation.query.QueryInfo; import java.util.ArrayList; +import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; @@ -608,6 +610,17 @@ public FeedResponse createFeedResponse(RxDocumentServiceResponse response return new FeedResponse<>(response.getQueryResponse(itemSerializer, cls), response); } + @Override + public FeedResponse createNonServiceFeedResponse(List items, boolean isChangeFeed, boolean isNoChanges) { + return new FeedResponse<>( + items, + new HashMap<>(), + new ConcurrentHashMap<>(), + isChangeFeed, + isNoChanges + ); + } + @Override public FeedResponse createChangeFeedResponse(RxDocumentServiceResponse response, CosmosItemSerializer itemSerializer, Class cls) { return new FeedResponse<>( 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 43a41bfab8da..152bde84273b 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,27 +3,15 @@ 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.CosmosDiagnosticsContext; import com.azure.cosmos.implementation.CosmosPagedFluxOptions; -import com.azure.cosmos.implementation.DiagnosticsProvider; -import com.azure.cosmos.implementation.FeedOperationState; -import com.azure.cosmos.implementation.ImplementationBridgeHelpers; import com.azure.cosmos.models.FeedResponse; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; import reactor.core.CoreSubscriber; import reactor.core.publisher.Flux; -import java.time.Duration; -import java.time.Instant; -import java.util.concurrent.atomic.AtomicLong; -import java.util.concurrent.atomic.AtomicReference; +import java.util.List; import java.util.function.Consumer; -import java.util.function.Function; /** * Cosmos implementation of {@link ContinuablePagedFlux}. @@ -39,74 +27,13 @@ * @see CosmosPagedFluxOptions * @see FeedResponse */ -public final class CosmosPagedFlux extends ContinuablePagedFlux> { - private static final ImplementationBridgeHelpers.CosmosDiagnosticsContextHelper.CosmosDiagnosticsContextAccessor ctxAccessor = - ImplementationBridgeHelpers.CosmosDiagnosticsContextHelper.getCosmosDiagnosticsContextAccessor(); +public abstract class CosmosPagedFlux extends ContinuablePagedFlux> { + // Ensure there can only be package-internal implementations + CosmosPagedFlux() {} - private final Function>> optionsFluxFunction; - private final Consumer> feedResponseConsumer; - private final int defaultPageSize; + public abstract CosmosPagedFlux handle(Consumer> newFeedResponseConsumer); - CosmosPagedFlux(Function>> optionsFluxFunction) { - this(optionsFluxFunction, null, -1); - } - - CosmosPagedFlux(Function>> optionsFluxFunction, - Consumer> feedResponseConsumer) { - this(optionsFluxFunction, feedResponseConsumer, -1); - } - - CosmosPagedFlux(Function>> optionsFluxFunction, - Consumer> feedResponseConsumer, - int defaultPageSize) { - this.optionsFluxFunction = optionsFluxFunction; - this.feedResponseConsumer = feedResponseConsumer; - this.defaultPageSize = defaultPageSize; - } - - /** - * Handle for invoking "side-effects" on each FeedResponse returned by CosmosPagedFlux - * - * @param newFeedResponseConsumer handler - * @return CosmosPagedFlux instance with attached handler - */ - public CosmosPagedFlux handle(Consumer> newFeedResponseConsumer) { - if (this.feedResponseConsumer != null) { - return new CosmosPagedFlux<>( - this.optionsFluxFunction, - this.feedResponseConsumer.andThen(newFeedResponseConsumer)); - } else { - return new CosmosPagedFlux<>(this.optionsFluxFunction, newFeedResponseConsumer); - } - } - - @Override - public Flux> byPage() { - CosmosPagedFluxOptions cosmosPagedFluxOptions = this.createCosmosPagedFluxOptions(); - return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)); - } - - @Override - public Flux> byPage(String continuationToken) { - CosmosPagedFluxOptions cosmosPagedFluxOptions = this.createCosmosPagedFluxOptions(); - cosmosPagedFluxOptions.setRequestContinuation(continuationToken); - return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)); - } - - @Override - public Flux> byPage(int preferredPageSize) { - CosmosPagedFluxOptions cosmosPagedFluxOptions = this.createCosmosPagedFluxOptions(); - cosmosPagedFluxOptions.setMaxItemCount(preferredPageSize); - return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)); - } - - @Override - public Flux> byPage(String continuationToken, int preferredPageSize) { - CosmosPagedFluxOptions cosmosPagedFluxOptions = this.createCosmosPagedFluxOptions(); - cosmosPagedFluxOptions.setRequestContinuation(continuationToken); - cosmosPagedFluxOptions.setMaxItemCount(preferredPageSize); - return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)); - } + protected abstract CosmosPagedFlux withDefaultPageSize(int pageSize); /** * Subscribe to consume all items of type {@code T} in the sequence respectively. This is recommended for most @@ -115,7 +42,7 @@ public Flux> byPage(String continuationToken, int preferredPageS * @param coreSubscriber The subscriber for this {@link CosmosPagedFlux} */ @Override - public void subscribe(@SuppressWarnings("NullableProblems") CoreSubscriber coreSubscriber) { + public void subscribe(CoreSubscriber coreSubscriber) { Flux> pagedResponse = this.byPage(); pagedResponse.flatMap(tFeedResponse -> { IterableStream elements = tFeedResponse.getElements(); @@ -126,201 +53,18 @@ public void subscribe(@SuppressWarnings("NullableProblems") CoreSubscriber withDefaultPageSize(int pageSize) { - return new CosmosPagedFlux<>(this.optionsFluxFunction, this.feedResponseConsumer, pageSize); - } - - private CosmosPagedFluxOptions createCosmosPagedFluxOptions() { - CosmosPagedFluxOptions cosmosPagedFluxOptions = new CosmosPagedFluxOptions(); - - if (this.defaultPageSize > 0) { - cosmosPagedFluxOptions.setMaxItemCount(this.defaultPageSize); - } - - return cosmosPagedFluxOptions; - } - - private Flux> wrapWithTracingIfEnabled(CosmosPagedFluxOptions pagedFluxOptions, - Flux> publisher, - AtomicLong feedResponseConsumerLatencyInNanos, - Context context) { - FeedOperationState state = pagedFluxOptions.getFeedOperationState(); - DiagnosticsProvider tracerProvider = state != null ? state.getDiagnosticsProvider() : null; - Object lockHolder = new Object(); - if (tracerProvider == null) { - - return publisher - .doOnEach(signal -> { - FeedResponse response = signal.get(); - synchronized (lockHolder) { - switch (signal.getType()) { - case ON_COMPLETE: - case ON_NEXT: - DiagnosticsProvider.recordFeedResponse( - feedResponseConsumer, - pagedFluxOptions.getFeedOperationState(), - () ->pagedFluxOptions.getSamplingRateSnapshot(), - tracerProvider, - response, - feedResponseConsumerLatencyInNanos); - break; - default: - break; - } - } - }); - } - - if (!tracerProvider.isEnabled()) { - pagedFluxOptions.setSamplingRateSnapshot(0, true); - } - - final boolean isSampledOut = tracerProvider.shouldSampleOutOperation(pagedFluxOptions); - final double samplingRateSnapshot = pagedFluxOptions.getSamplingRateSnapshot(); - - Flux> result = tracerProvider - .runUnderSpanInContext(publisher) - .doOnEach(signal -> { - FeedResponse response = signal.get(); - Context traceCtx = DiagnosticsProvider.getContextFromReactorOrNull(signal.getContextView()); - - synchronized (lockHolder) { - switch (signal.getType()) { - case ON_COMPLETE: - if (response != null) { - DiagnosticsProvider.recordFeedResponse( - feedResponseConsumer, - pagedFluxOptions.getFeedOperationState(), - () ->pagedFluxOptions.getSamplingRateSnapshot(), - tracerProvider, - response, - feedResponseConsumerLatencyInNanos); - } - state.mergeDiagnosticsContext(); - - CosmosDiagnosticsContext ctxSnapshot = state.getDiagnosticsContextSnapshot(); - - ctxAccessor - .setSamplingRateSnapshot(ctxSnapshot, samplingRateSnapshot, isSampledOut); - - tracerProvider.recordFeedResponseConsumerLatency( - signal, - ctxSnapshot, - Duration.ofNanos(feedResponseConsumerLatencyInNanos.get())); - - tracerProvider.endSpan(ctxSnapshot, traceCtx, ctxAccessor.isEmptyCompletion(ctxSnapshot), isSampledOut); - - break; - case ON_NEXT: - DiagnosticsProvider.recordFeedResponse( - feedResponseConsumer, - pagedFluxOptions.getFeedOperationState(), - () ->pagedFluxOptions.getSamplingRateSnapshot(), - tracerProvider, - response, - feedResponseConsumerLatencyInNanos); - state.mergeDiagnosticsContext(); - CosmosDiagnosticsContext ctxSnapshotOnNext = state.getDiagnosticsContextSnapshot(); - ctxAccessor - .setSamplingRateSnapshot(ctxSnapshotOnNext, samplingRateSnapshot, isSampledOut); - tracerProvider.endSpan(ctxSnapshotOnNext, traceCtx, false, isSampledOut); - state.resetDiagnosticsContext(); - - DiagnosticsProvider.setContextInReactor(tracerProvider.startSpan( - state.getSpanName(), - state.getDiagnosticsContextSnapshot(), - traceCtx, - isSampledOut)); - - break; - - case ON_ERROR: - state.mergeDiagnosticsContext(); - CosmosDiagnosticsContext ctxSnapshotOnError = state.getDiagnosticsContextSnapshot(); - ctxAccessor - .setSamplingRateSnapshot(ctxSnapshotOnError, samplingRateSnapshot, isSampledOut); - tracerProvider.recordFeedResponseConsumerLatency( - signal, - ctxSnapshotOnError, - Duration.ofNanos(feedResponseConsumerLatencyInNanos.get())); - - // all info is extracted from CosmosException when applicable - tracerProvider.endSpan( - state.getDiagnosticsContextSnapshot(), - traceCtx, - signal.getThrowable(), - isSampledOut - ); - - break; - - default: - break; - } - } - }); - - return Flux - .deferContextual(reactorCtx -> result - .doOnCancel(() -> { - Context traceCtx = DiagnosticsProvider.getContextFromReactorOrNull(reactorCtx); - synchronized (lockHolder) { - state.mergeDiagnosticsContext(); - CosmosDiagnosticsContext ctxSnapshot = state.getDiagnosticsContextSnapshot(); - - ctxAccessor - .setSamplingRateSnapshot(ctxSnapshot, samplingRateSnapshot, isSampledOut); - - tracerProvider.endSpan(ctxSnapshot, traceCtx, false, isSampledOut); - } - }) - .doOnComplete(() -> { - Context traceCtx = DiagnosticsProvider.getContextFromReactorOrNull(reactorCtx); - synchronized(lockHolder) { - state.mergeDiagnosticsContext(); - - CosmosDiagnosticsContext ctxSnapshot = state.getDiagnosticsContextSnapshot(); - ctxAccessor - .setSamplingRateSnapshot(ctxSnapshot, samplingRateSnapshot, isSampledOut); - tracerProvider.endSpan(ctxSnapshot, traceCtx, ctxAccessor.isEmptyCompletion(ctxSnapshot), isSampledOut); - } - })) - .contextWrite(DiagnosticsProvider.setContextInReactor( - tracerProvider.startSpan( - state.getSpanName(), - state.getDiagnosticsContextSnapshot(), - context, - isSampledOut) - )); - } - - - - private Flux> byPage(CosmosPagedFluxOptions pagedFluxOptions, Context context) { - AtomicReference startTime = new AtomicReference<>(); - AtomicLong feedResponseConsumerLatencyInNanos = new AtomicLong(0); - - Flux> result = - wrapWithTracingIfEnabled( - pagedFluxOptions, - this.optionsFluxFunction.apply(pagedFluxOptions), - feedResponseConsumerLatencyInNanos, - context) - .doOnSubscribe(ignoredValue -> { - startTime.set(Instant.now()); - feedResponseConsumerLatencyInNanos.set(0); - }); - - return result; - } - - /////////////////////////////////////////////////////////////////////////////////////////// - // the following helper/accessor only helps to access this class outside of this package.// - /////////////////////////////////////////////////////////////////////////////////////////// - static void initialize() { - ImplementationBridgeHelpers.CosmosPageFluxHelper.setCosmosPageFluxAccessor( - (ImplementationBridgeHelpers.CosmosPageFluxHelper.CosmosPageFluxAccessor) CosmosPagedFlux::new); + /** + * Creates an instance of a CosmosPagedFlux for mocking purposes or when injecting CosmosPagedFlux + * instances from a different data source + * @param items - the list of items to be returned + * @param isChangeFeed - a flag incidcating whether the CosmsoPagedFluy will be returned from a change feed API + * @return an instance of CosmosPagedFlux + * @param The type of the items + */ + public static CosmosPagedFlux createFromList(List items, boolean isChangeFeed) { + return new CosmosPagedFluxStaticListImpl<>( + items, + isChangeFeed + ); } - - static { initialize(); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFluxDefaultImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFluxDefaultImpl.java new file mode 100644 index 000000000000..bac7f707b3c1 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFluxDefaultImpl.java @@ -0,0 +1,306 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +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.CosmosDiagnosticsContext; +import com.azure.cosmos.implementation.CosmosPagedFluxOptions; +import com.azure.cosmos.implementation.DiagnosticsProvider; +import com.azure.cosmos.implementation.FeedOperationState; +import com.azure.cosmos.implementation.ImplementationBridgeHelpers; +import com.azure.cosmos.models.FeedResponse; +import reactor.core.CoreSubscriber; +import reactor.core.publisher.Flux; + +import java.time.Duration; +import java.time.Instant; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import java.util.function.Function; + +/** + * Cosmos implementation of {@link ContinuablePagedFlux}. + *

+ * This type is a Flux that provides the ability to operate on pages of type {@link FeedResponse} and individual items + * in such pages. This type supports {@link String} type continuation tokens, allowing for restarting from a + * previously-retrieved continuation token. + *

+ * For more information on the base type, refer {@link ContinuablePagedFlux} + * + * @param The type of elements in a {@link com.azure.core.util.paging.ContinuablePage} + * @see com.azure.core.util.paging.ContinuablePage + * @see CosmosPagedFluxOptions + * @see FeedResponse + */ +final class CosmosPagedFluxDefaultImpl extends CosmosPagedFlux { + private static final ImplementationBridgeHelpers.CosmosDiagnosticsContextHelper.CosmosDiagnosticsContextAccessor ctxAccessor = + ImplementationBridgeHelpers.CosmosDiagnosticsContextHelper.getCosmosDiagnosticsContextAccessor(); + + private final Function>> optionsFluxFunction; + private final Consumer> feedResponseConsumer; + private final int defaultPageSize; + + CosmosPagedFluxDefaultImpl(Function>> optionsFluxFunction) { + this(optionsFluxFunction, null, -1); + } + + CosmosPagedFluxDefaultImpl(Function>> optionsFluxFunction, + Consumer> feedResponseConsumer) { + this(optionsFluxFunction, feedResponseConsumer, -1); + } + + CosmosPagedFluxDefaultImpl(Function>> optionsFluxFunction, + Consumer> feedResponseConsumer, + int defaultPageSize) { + super(); + this.optionsFluxFunction = optionsFluxFunction; + this.feedResponseConsumer = feedResponseConsumer; + this.defaultPageSize = defaultPageSize; + } + + /** + * Handle for invoking "side-effects" on each FeedResponse returned by CosmosPagedFlux + * + * @param newFeedResponseConsumer handler + * @return CosmosPagedFlux instance with attached handler + */ + public CosmosPagedFlux handle(Consumer> newFeedResponseConsumer) { + if (this.feedResponseConsumer != null) { + return new CosmosPagedFluxDefaultImpl<>( + this.optionsFluxFunction, + this.feedResponseConsumer.andThen(newFeedResponseConsumer)); + } else { + return new CosmosPagedFluxDefaultImpl<>(this.optionsFluxFunction, newFeedResponseConsumer); + } + } + + @Override + protected CosmosPagedFlux withDefaultPageSize(int pageSize) { + return new CosmosPagedFluxDefaultImpl<>(this.optionsFluxFunction, this.feedResponseConsumer, pageSize); + } + + @Override + public Flux> byPage() { + CosmosPagedFluxOptions cosmosPagedFluxOptions = this.createCosmosPagedFluxOptions(); + return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)); + } + + @Override + public Flux> byPage(String continuationToken) { + CosmosPagedFluxOptions cosmosPagedFluxOptions = this.createCosmosPagedFluxOptions(); + cosmosPagedFluxOptions.setRequestContinuation(continuationToken); + return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)); + } + + @Override + public Flux> byPage(int preferredPageSize) { + CosmosPagedFluxOptions cosmosPagedFluxOptions = this.createCosmosPagedFluxOptions(); + cosmosPagedFluxOptions.setMaxItemCount(preferredPageSize); + return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)); + } + + @Override + public Flux> byPage(String continuationToken, int preferredPageSize) { + CosmosPagedFluxOptions cosmosPagedFluxOptions = this.createCosmosPagedFluxOptions(); + cosmosPagedFluxOptions.setRequestContinuation(continuationToken); + cosmosPagedFluxOptions.setMaxItemCount(preferredPageSize); + return FluxUtil.fluxContext(context -> byPage(cosmosPagedFluxOptions, context)); + } + + private CosmosPagedFluxOptions createCosmosPagedFluxOptions() { + CosmosPagedFluxOptions cosmosPagedFluxOptions = new CosmosPagedFluxOptions(); + + if (this.defaultPageSize > 0) { + cosmosPagedFluxOptions.setMaxItemCount(this.defaultPageSize); + } + + return cosmosPagedFluxOptions; + } + + private Flux> wrapWithTracingIfEnabled(CosmosPagedFluxOptions pagedFluxOptions, + Flux> publisher, + AtomicLong feedResponseConsumerLatencyInNanos, + Context context) { + FeedOperationState state = pagedFluxOptions.getFeedOperationState(); + DiagnosticsProvider tracerProvider = state != null ? state.getDiagnosticsProvider() : null; + Object lockHolder = new Object(); + if (tracerProvider == null) { + + return publisher + .doOnEach(signal -> { + FeedResponse response = signal.get(); + synchronized (lockHolder) { + switch (signal.getType()) { + case ON_COMPLETE: + case ON_NEXT: + DiagnosticsProvider.recordFeedResponse( + feedResponseConsumer, + pagedFluxOptions.getFeedOperationState(), + () ->pagedFluxOptions.getSamplingRateSnapshot(), + tracerProvider, + response, + feedResponseConsumerLatencyInNanos); + break; + default: + break; + } + } + }); + } + + if (!tracerProvider.isEnabled()) { + pagedFluxOptions.setSamplingRateSnapshot(0, true); + } + + final boolean isSampledOut = tracerProvider.shouldSampleOutOperation(pagedFluxOptions); + final double samplingRateSnapshot = pagedFluxOptions.getSamplingRateSnapshot(); + + Flux> result = tracerProvider + .runUnderSpanInContext(publisher) + .doOnEach(signal -> { + FeedResponse response = signal.get(); + Context traceCtx = DiagnosticsProvider.getContextFromReactorOrNull(signal.getContextView()); + + synchronized (lockHolder) { + switch (signal.getType()) { + case ON_COMPLETE: + if (response != null) { + DiagnosticsProvider.recordFeedResponse( + feedResponseConsumer, + pagedFluxOptions.getFeedOperationState(), + () ->pagedFluxOptions.getSamplingRateSnapshot(), + tracerProvider, + response, + feedResponseConsumerLatencyInNanos); + } + state.mergeDiagnosticsContext(); + + CosmosDiagnosticsContext ctxSnapshot = state.getDiagnosticsContextSnapshot(); + + ctxAccessor + .setSamplingRateSnapshot(ctxSnapshot, samplingRateSnapshot, isSampledOut); + + tracerProvider.recordFeedResponseConsumerLatency( + signal, + ctxSnapshot, + Duration.ofNanos(feedResponseConsumerLatencyInNanos.get())); + + tracerProvider.endSpan(ctxSnapshot, traceCtx, ctxAccessor.isEmptyCompletion(ctxSnapshot), isSampledOut); + + break; + case ON_NEXT: + DiagnosticsProvider.recordFeedResponse( + feedResponseConsumer, + pagedFluxOptions.getFeedOperationState(), + () ->pagedFluxOptions.getSamplingRateSnapshot(), + tracerProvider, + response, + feedResponseConsumerLatencyInNanos); + state.mergeDiagnosticsContext(); + CosmosDiagnosticsContext ctxSnapshotOnNext = state.getDiagnosticsContextSnapshot(); + ctxAccessor + .setSamplingRateSnapshot(ctxSnapshotOnNext, samplingRateSnapshot, isSampledOut); + tracerProvider.endSpan(ctxSnapshotOnNext, traceCtx, false, isSampledOut); + state.resetDiagnosticsContext(); + + DiagnosticsProvider.setContextInReactor(tracerProvider.startSpan( + state.getSpanName(), + state.getDiagnosticsContextSnapshot(), + traceCtx, + isSampledOut)); + + break; + + case ON_ERROR: + state.mergeDiagnosticsContext(); + CosmosDiagnosticsContext ctxSnapshotOnError = state.getDiagnosticsContextSnapshot(); + ctxAccessor + .setSamplingRateSnapshot(ctxSnapshotOnError, samplingRateSnapshot, isSampledOut); + tracerProvider.recordFeedResponseConsumerLatency( + signal, + ctxSnapshotOnError, + Duration.ofNanos(feedResponseConsumerLatencyInNanos.get())); + + // all info is extracted from CosmosException when applicable + tracerProvider.endSpan( + state.getDiagnosticsContextSnapshot(), + traceCtx, + signal.getThrowable(), + isSampledOut + ); + + break; + + default: + break; + } + } + }); + + return Flux + .deferContextual(reactorCtx -> result + .doOnCancel(() -> { + Context traceCtx = DiagnosticsProvider.getContextFromReactorOrNull(reactorCtx); + synchronized (lockHolder) { + state.mergeDiagnosticsContext(); + CosmosDiagnosticsContext ctxSnapshot = state.getDiagnosticsContextSnapshot(); + + ctxAccessor + .setSamplingRateSnapshot(ctxSnapshot, samplingRateSnapshot, isSampledOut); + + tracerProvider.endSpan(ctxSnapshot, traceCtx, false, isSampledOut); + } + }) + .doOnComplete(() -> { + Context traceCtx = DiagnosticsProvider.getContextFromReactorOrNull(reactorCtx); + synchronized(lockHolder) { + state.mergeDiagnosticsContext(); + + CosmosDiagnosticsContext ctxSnapshot = state.getDiagnosticsContextSnapshot(); + ctxAccessor + .setSamplingRateSnapshot(ctxSnapshot, samplingRateSnapshot, isSampledOut); + tracerProvider.endSpan(ctxSnapshot, traceCtx, ctxAccessor.isEmptyCompletion(ctxSnapshot), isSampledOut); + } + })) + .contextWrite(DiagnosticsProvider.setContextInReactor( + tracerProvider.startSpan( + state.getSpanName(), + state.getDiagnosticsContextSnapshot(), + context, + isSampledOut) + )); + } + + private Flux> byPage(CosmosPagedFluxOptions pagedFluxOptions, Context context) { + AtomicReference startTime = new AtomicReference<>(); + AtomicLong feedResponseConsumerLatencyInNanos = new AtomicLong(0); + + Flux> result = + wrapWithTracingIfEnabled( + pagedFluxOptions, + this.optionsFluxFunction.apply(pagedFluxOptions), + feedResponseConsumerLatencyInNanos, + context) + .doOnSubscribe(ignoredValue -> { + startTime.set(Instant.now()); + feedResponseConsumerLatencyInNanos.set(0); + }); + + return result; + } + + /////////////////////////////////////////////////////////////////////////////////////////// + // the following helper/accessor only helps to access this class outside of this package.// + /////////////////////////////////////////////////////////////////////////////////////////// + static void initialize() { + ImplementationBridgeHelpers.CosmosPageFluxHelper.setCosmosPageFluxAccessor( + (ImplementationBridgeHelpers.CosmosPageFluxHelper.CosmosPageFluxAccessor) CosmosPagedFluxDefaultImpl::new); + } + + static { initialize(); } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFluxStaticListImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFluxStaticListImpl.java new file mode 100644 index 000000000000..faa30a662977 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFluxStaticListImpl.java @@ -0,0 +1,132 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.util; + +import com.azure.cosmos.implementation.ImplementationBridgeHelpers; +import com.azure.cosmos.models.FeedResponse; +import reactor.core.publisher.Flux; + +import java.util.ArrayList; +import java.util.List; +import java.util.function.Consumer; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +final class CosmosPagedFluxStaticListImpl extends CosmosPagedFlux { + private static final ImplementationBridgeHelpers.FeedResponseHelper.FeedResponseAccessor feedResponseHlp = + ImplementationBridgeHelpers.FeedResponseHelper.getFeedResponseAccessor(); + + private static final int DEFAULT_PAGE_SIZE = 100; + + private final Consumer> feedResponseConsumer; + + private final List items; + private final boolean isChangeFeed; + private final int defaultPageSize; + + CosmosPagedFluxStaticListImpl(List items, boolean isChangeFeed) { + this(items, isChangeFeed, null, DEFAULT_PAGE_SIZE); + } + + CosmosPagedFluxStaticListImpl( + List items, + boolean isChangeFeed, + Consumer> feedResponseConsumer, + int defaultPageSize) { + + super(); + + checkNotNull(items, "Argument 'items' must not be null."); + this.items = items; + this.isChangeFeed = isChangeFeed; + this.feedResponseConsumer = feedResponseConsumer; + this.defaultPageSize = defaultPageSize; + } + + @Override + public Flux> byPage() { + return byPage(null, DEFAULT_PAGE_SIZE); + } + + @Override + public Flux> byPage(String continuation) { + return byPage(continuation, DEFAULT_PAGE_SIZE); + } + + @Override + public Flux> byPage(int pageSize) { + return byPage(null, pageSize); + } + + @Override + public Flux> byPage(String continuation, int pageSize) { + int i = -1; + + if (continuation != null && !continuation.isEmpty()) { + i = Integer.parseUnsignedInt(continuation); + } + + List> pages = new ArrayList<>(); + if (i >= this.items.size()) { + if (!this.isChangeFeed) { + pages.add(feedResponseHlp.createNonServiceFeedResponse( + new ArrayList<>(), + false, + false)); + } else { + pages.add(feedResponseHlp.createNonServiceFeedResponse( + new ArrayList<>(), + true, + true)); + + } + + return Flux.fromIterable(pages); + } + + int pageCount = (this.items.size() - i - 1) / pageSize + 1; + for (int p = 0; p < pageCount; p++) { + List itemsForPage = new ArrayList<>(); + while ( ++i < this.items.size() && itemsForPage.size() < pageSize) { + itemsForPage.add(this.items.get(i)); + } + + pages.add(feedResponseHlp.createNonServiceFeedResponse( + new ArrayList<>(), + this.isChangeFeed, + false)); + } + + if (this.isChangeFeed) { + pages.add(feedResponseHlp.createNonServiceFeedResponse( + new ArrayList<>(), + true, + true)); + } + + return Flux.fromIterable(pages); + } + + @Override + public CosmosPagedFlux handle(Consumer> newFeedResponseConsumer) { + if (this.feedResponseConsumer != null) { + return new CosmosPagedFluxStaticListImpl<>( + this.items, + this.isChangeFeed, + this.feedResponseConsumer.andThen(newFeedResponseConsumer), + this.defaultPageSize); + } else { + return new CosmosPagedFluxStaticListImpl<>( + this.items, + this.isChangeFeed, + newFeedResponseConsumer, + this.defaultPageSize); + } + } + + @Override + protected CosmosPagedFlux withDefaultPageSize(int pageSize) { + return new CosmosPagedFluxStaticListImpl<>(this.items, this.isChangeFeed, this.feedResponseConsumer, pageSize); + } +} 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 9fd45e218d85..18e013c29c61 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,7 +24,7 @@ private UtilBridgeInternal() {} @Warning(value = INTERNAL_USE_ONLY_WARNING) public static CosmosPagedFlux createCosmosPagedFlux(Function>> pagedFluxOptionsFluxFunction) { - return new CosmosPagedFlux<>(pagedFluxOptionsFluxFunction); + return new CosmosPagedFluxDefaultImpl<>(pagedFluxOptionsFluxFunction); } @Warning(value = INTERNAL_USE_ONLY_WARNING) @@ -34,6 +34,6 @@ public static CosmosPagedIterable createCosmosPagedIterable(CosmosPagedFl @Warning(value = INTERNAL_USE_ONLY_WARNING) public static void initializeAllAccessors() { - CosmosPagedFlux.initialize(); + CosmosPagedFluxDefaultImpl.initialize(); } } From 5bcc3991ec64d1d8f77c79ec87a3a82e2db5e665 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Wed, 16 Apr 2025 16:09:40 +0000 Subject: [PATCH 06/11] Refactoring the refactoring to avoid LINT errors --- .../azure/cosmos/util/CosmosPagedFlux.java | 30 +++++++++++++++++-- .../util/CosmosPagedFluxDefaultImpl.java | 4 +-- .../util/CosmosPagedFluxStaticListImpl.java | 2 +- 3 files changed, 29 insertions(+), 7 deletions(-) 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 152bde84273b..929dafc02a33 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 @@ -27,13 +27,37 @@ * @see CosmosPagedFluxOptions * @see FeedResponse */ -public abstract class CosmosPagedFlux extends ContinuablePagedFlux> { +public class CosmosPagedFlux extends ContinuablePagedFlux> { // Ensure there can only be package-internal implementations CosmosPagedFlux() {} - public abstract CosmosPagedFlux handle(Consumer> newFeedResponseConsumer); + @Override + public Flux> byPage() { + throw new UnsupportedOperationException("Has to be overridden in child classes."); + } + + @Override + public Flux> byPage(String s) { + throw new UnsupportedOperationException("Has to be overridden in child classes."); + } + + @Override + public Flux> byPage(int i) { + throw new UnsupportedOperationException("Has to be overridden in child classes."); + } + + @Override + public Flux> byPage(String s, int i) { + throw new UnsupportedOperationException("Has to be overridden in child classes."); + } - protected abstract CosmosPagedFlux withDefaultPageSize(int pageSize); + public CosmosPagedFlux handle(Consumer> newFeedResponseConsumer) { + throw new UnsupportedOperationException("Has to be overridden in child classes."); + } + + CosmosPagedFlux withDefaultPageSize(int pageSize) { + throw new UnsupportedOperationException("Has to be overridden in child classes."); + } /** * Subscribe to consume all items of type {@code T} in the sequence respectively. This is recommended for most diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFluxDefaultImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFluxDefaultImpl.java index bac7f707b3c1..d289a474a637 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFluxDefaultImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFluxDefaultImpl.java @@ -5,7 +5,6 @@ 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.CosmosDiagnosticsContext; import com.azure.cosmos.implementation.CosmosPagedFluxOptions; @@ -13,7 +12,6 @@ import com.azure.cosmos.implementation.FeedOperationState; import com.azure.cosmos.implementation.ImplementationBridgeHelpers; import com.azure.cosmos.models.FeedResponse; -import reactor.core.CoreSubscriber; import reactor.core.publisher.Flux; import java.time.Duration; @@ -80,7 +78,7 @@ public CosmosPagedFlux handle(Consumer> newFeedResponseConsum } @Override - protected CosmosPagedFlux withDefaultPageSize(int pageSize) { + CosmosPagedFlux withDefaultPageSize(int pageSize) { return new CosmosPagedFluxDefaultImpl<>(this.optionsFluxFunction, this.feedResponseConsumer, pageSize); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFluxStaticListImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFluxStaticListImpl.java index faa30a662977..0c1c412f988a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFluxStaticListImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFluxStaticListImpl.java @@ -126,7 +126,7 @@ public CosmosPagedFlux handle(Consumer> newFeedResponseConsum } @Override - protected CosmosPagedFlux withDefaultPageSize(int pageSize) { + CosmosPagedFlux withDefaultPageSize(int pageSize) { return new CosmosPagedFluxStaticListImpl<>(this.items, this.isChangeFeed, this.feedResponseConsumer, pageSize); } } From 7c0f9b9b15bee1f0380d685d9510e6fa0e233160 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 22 Apr 2025 19:41:25 +0200 Subject: [PATCH 07/11] Update sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java Co-authored-by: Tomas Varon <70857381+tvaron3@users.noreply.github.com> --- .../src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 1e808cbecd53..ce4de922f157 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 @@ -822,7 +822,7 @@ public CosmosAsyncContainer getContainer(String id) { if (asyncContainer == null) { throw new IllegalStateException( - "The implementation of teh custom container creation interceptor must not return null."); + "The implementation of the custom container creation interceptor must not return null."); } return asyncContainer; } From d54c91bab9fc7e20ceabd644b7d789954569a395 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 22 Apr 2025 18:39:52 +0000 Subject: [PATCH 08/11] Reacting to code review feedback --- .../src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java | 1 + .../src/main/java/com/azure/cosmos/CosmosClient.java | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java index b2fc6d1ea79b..2808f3f8c279 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java @@ -229,6 +229,7 @@ public void validateContainerCreationInterceptor() { List items = customSyncContainer .queryItems(querySpec, null, ObjectNode.class) .stream().collect(Collectors.toList()); + fail("Not yet cached - the query above should always throw"); } catch (CosmosException cosmosException) { // Container does not exist - when not cached should fail assertThat(cosmosException.getStatusCode()).isEqualTo(404); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java index ba5adcb057e7..9a8b2e7e7bb4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java @@ -5,7 +5,6 @@ import com.azure.core.annotation.ServiceClient; import com.azure.cosmos.implementation.ImplementationBridgeHelpers; -import com.azure.cosmos.implementation.guava25.base.Function; import com.azure.cosmos.models.CosmosContainerIdentity; import com.azure.cosmos.models.CosmosDatabaseProperties; import com.azure.cosmos.models.CosmosDatabaseRequestOptions; From 4f280d051ab956ab17e884e2e34777d82546634c Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 22 Apr 2025 18:41:11 +0000 Subject: [PATCH 09/11] Revert "Reacting to code review feedback" This reverts commit d54c91bab9fc7e20ceabd644b7d789954569a395. --- .../src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java | 1 - .../src/main/java/com/azure/cosmos/CosmosClient.java | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java index 2808f3f8c279..b2fc6d1ea79b 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java @@ -229,7 +229,6 @@ public void validateContainerCreationInterceptor() { List items = customSyncContainer .queryItems(querySpec, null, ObjectNode.class) .stream().collect(Collectors.toList()); - fail("Not yet cached - the query above should always throw"); } catch (CosmosException cosmosException) { // Container does not exist - when not cached should fail assertThat(cosmosException.getStatusCode()).isEqualTo(404); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java index 9a8b2e7e7bb4..ba5adcb057e7 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java @@ -5,6 +5,7 @@ import com.azure.core.annotation.ServiceClient; import com.azure.cosmos.implementation.ImplementationBridgeHelpers; +import com.azure.cosmos.implementation.guava25.base.Function; import com.azure.cosmos.models.CosmosContainerIdentity; import com.azure.cosmos.models.CosmosDatabaseProperties; import com.azure.cosmos.models.CosmosDatabaseRequestOptions; From 7729c22659c65cfc4889f997633c7d23fcab90f5 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 22 Apr 2025 18:42:04 +0000 Subject: [PATCH 10/11] Reapply "Reacting to code review feedback" This reverts commit 4f280d051ab956ab17e884e2e34777d82546634c. --- .../src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java | 1 + .../src/main/java/com/azure/cosmos/CosmosClient.java | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java index b2fc6d1ea79b..2808f3f8c279 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java @@ -229,6 +229,7 @@ public void validateContainerCreationInterceptor() { List items = customSyncContainer .queryItems(querySpec, null, ObjectNode.class) .stream().collect(Collectors.toList()); + fail("Not yet cached - the query above should always throw"); } catch (CosmosException cosmosException) { // Container does not exist - when not cached should fail assertThat(cosmosException.getStatusCode()).isEqualTo(404); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java index ba5adcb057e7..9a8b2e7e7bb4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClient.java @@ -5,7 +5,6 @@ import com.azure.core.annotation.ServiceClient; import com.azure.cosmos.implementation.ImplementationBridgeHelpers; -import com.azure.cosmos.implementation.guava25.base.Function; import com.azure.cosmos.models.CosmosContainerIdentity; import com.azure.cosmos.models.CosmosDatabaseProperties; import com.azure.cosmos.models.CosmosDatabaseRequestOptions; From 4a59db3f9b0383f4c0711e655808a5843d834b28 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 22 Apr 2025 18:48:51 +0000 Subject: [PATCH 11/11] Update CosmosAsyncClient.java --- .../src/main/java/com/azure/cosmos/CosmosAsyncClient.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java index 7040a9c1f2e0..789388810fdf 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 @@ -29,7 +29,6 @@ import com.azure.cosmos.implementation.clienttelemetry.TagName; import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdMetrics; import com.azure.cosmos.implementation.faultinjection.IFaultInjectorProvider; -import com.azure.cosmos.implementation.guava25.base.Function; import com.azure.cosmos.implementation.throughputControl.config.ThroughputControlGroupInternal; import com.azure.cosmos.models.CosmosAuthorizationTokenResolver; import com.azure.cosmos.models.CosmosClientTelemetryConfig; @@ -60,6 +59,7 @@ import java.util.EnumSet; import java.util.List; import java.util.Objects; +import java.util.function.Function; import java.util.stream.Collectors; import static com.azure.core.util.FluxUtil.withContext;