Allowing customers to wrap CosmosAsyncContainer (Azure#43724)

* Initial draft

* Added unit test

* Update CosmosClientBuilderTest.java

* Update CosmosClientBuilderTest.java

* Adding simple extensibility for CosmosPagedFlux

* Refactoring the refactoring to avoid LINT errors

* Update sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java

Co-authored-by: Tomas Varon <[email protected]>

* Reacting to code review feedback

* Revert "Reacting to code review feedback"

This reverts commit d54c91b.

* Reapply "Reacting to code review feedback"

This reverts commit 4f280d0.

* Update CosmosAsyncClient.java

---------

Co-authored-by: Tomas Varon <[email protected]>

Author: FabianMeiswinkel

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosClientBuilderTest.java

@@ -10,14 +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.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/";
@@ -149,4 +162,210 @@ public void validateSessionTokenCapturingForAccountDefaultConsistencyWithEnvVari
             System.clearProperty("COSMOS.SESSION_CAPTURING_TYPE");
         }
     }
+
+    @Test(groups = "emulator")
+    public void validateContainerCreationInterceptor() {
+        CosmosClient clientWithoutInterceptor = new CosmosClientBuilder()
+            .endpoint(TestConfigurations.HOST)
+            .key(TestConfigurations.MASTER_KEY)
+            .userAgentSuffix("noInterceptor")
+            .buildClient();
+
+        ConcurrentMap<CacheKey, List<?>> queryCache = new ConcurrentHashMap<>();
+
+        CosmosClient clientWithInterceptor = new CosmosClientBuilder()
+            .endpoint(TestConfigurations.HOST)
+            .key(TestConfigurations.MASTER_KEY)
+            .userAgentSuffix("withInterceptor")
+            .containerCreationInterceptor(originalContainer -> new CacheAndValidateQueriesContainer(originalContainer, queryCache))
+            .buildClient();
+
+        CosmosAsyncClient asyncClientWithInterceptor = new CosmosClientBuilder()
+            .endpoint(TestConfigurations.HOST)
+            .key(TestConfigurations.MASTER_KEY)
+            .userAgentSuffix("withInterceptor")
+            .containerCreationInterceptor(originalContainer -> new CacheAndValidateQueriesContainer(originalContainer, queryCache))
+            .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(CacheAndValidateQueriesContainer.class);
+
+        CosmosAsyncContainer customAsyncContainer = asyncClientWithInterceptor
+            .getDatabase("TestDB")
+            .getContainer("TestContainer");
+        assertThat(customAsyncContainer).isNotNull();
+        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<ObjectNode> 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);
+            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<ObjectNode> items = customSyncContainer
+            .queryItems(querySpec, null, ObjectNode.class)
+            .stream().collect(Collectors.toList());
+
+        querySpec = new SqlQuerySpec().setQueryText("SELECT * from c");
+        CosmosPagedFlux<ObjectNode> 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<ObjectNode> 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<SqlParameter> parameters;
+        public CacheKey(String className, SqlQuerySpec querySpec) {
+            this.className = className;
+            this.queryText = querySpec.getQueryText();
+            List<SqlParameter> 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 CacheAndValidateQueriesContainer extends CosmosAsyncContainer {
+        private final ConcurrentMap<CacheKey, List<?>> queryCache;
+
+        protected CacheAndValidateQueriesContainer(
+            CosmosAsyncContainer toBeWrappedContainer,
+            ConcurrentMap<CacheKey, List<?>> queryCache) {
+
+            super(toBeWrappedContainer);
+            this.queryCache = queryCache;
+        }
+
+        @Override
+        public <T> CosmosPagedFlux<T> queryItems(String query, CosmosQueryRequestOptions options, Class<T> classType) {
+            throw new IllegalStateException("No unparameterized queries allowed. Use parameterized query instead.");
+        }
+
+        @Override
+        public <T> CosmosPagedFlux<T> queryItems(SqlQuerySpec querySpec, Class<T> classType) {
+            return this.queryItems(querySpec, null, classType);
+        }
+
+        @Override
+        public <T> CosmosPagedFlux<T> queryItems(String query, Class<T> classType) {
+            throw new IllegalStateException("No unparameterized queries allowed. Use parameterized query instead.");
+        }
+
+        @Override
+        public <T> CosmosPagedFlux<T> queryItems(SqlQuerySpec querySpec, CosmosQueryRequestOptions options, Class<T> classType) {
+            CacheKey key = new CacheKey(classType.getCanonicalName(), querySpec);
+            List<?> cachedResult = this.queryCache.get(key);
+            if (cachedResult != null) {
+                return CosmosPagedFlux.createFromList((List<T>)cachedResult, false);
+            }
+
+            return super
+                .queryItems(querySpec, options, classType);
+        }
+    }
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java

@@ -59,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;
@@ -88,6 +89,9 @@ public final class CosmosAsyncClient implements Closeable {
         .CosmosClientTelemetryConfigHelper
         .getCosmosClientTelemetryConfigAccessor();
 
+    private final static Function<CosmosAsyncContainer, CosmosAsyncContainer> DEFAULT_CONTAINER_FACTORY =
+        (originalContainer) -> originalContainer;
+
     private final AsyncDocumentClient asyncDocumentClient;
     private final String serviceEndpoint;
     private final ConnectionPolicy connectionPolicy;
@@ -103,6 +107,7 @@ public final class CosmosAsyncClient implements Closeable {
     private final WriteRetryPolicy nonIdempotentWriteRetryPolicy;
     private final List<CosmosOperationPolicy> requestPolicies;
     private final CosmosItemSerializer defaultCustomSerializer;
+    private final java.util.function.Function<CosmosAsyncContainer, CosmosAsyncContainer> containerFactory;
 
     CosmosAsyncClient(CosmosClientBuilder builder) {
         // Async Cosmos client wrapper
@@ -121,6 +126,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();
 
@@ -788,6 +798,10 @@ String getUserAgent() {
         return this.asyncDocumentClient.getUserAgent();
     }
 
+    java.util.function.Function<CosmosAsyncContainer, CosmosAsyncContainer> getContainerCreationInterceptor() {
+        return this.containerFactory;
+    }
+
     ///////////////////////////////////////////////////////////////////////////////////////////
     // the following helper/accessor only helps to access this class outside of this package.//
     ///////////////////////////////////////////////////////////////////////////////////////////

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;

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncDatabase.java

@@ -815,7 +815,16 @@ public CosmosPagedFlux<CosmosContainerProperties> 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 the custom container creation interceptor must not return null.");
+        }
+        return asyncContainer;
     }
 
     /**

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosClientBuilder.java

@@ -37,6 +37,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;
@@ -149,6 +150,8 @@ public class CosmosClientBuilder implements
     private boolean isRegionScopedSessionCapturingEnabled = false;
     private boolean serverCertValidationDisabled = false;
 
+    private Function<CosmosAsyncContainer, CosmosAsyncContainer> containerFactory = null;
+
     /**
      * Instantiates a new Cosmos client builder.
      */
@@ -172,6 +175,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<CosmosAsyncContainer, CosmosAsyncContainer> containerCreationInterceptor() {
+        return this.containerFactory;
+    }
+
+    /**
+     * Sets a function that allows intercepting container creation - for example to wrap the original
+     * 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}
+     */
+    public CosmosClientBuilder containerCreationInterceptor(
+        Function<CosmosAsyncContainer, CosmosAsyncContainer> 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

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ImplementationBridgeHelpers.java

@@ -1095,6 +1095,10 @@ <T> FeedResponse<T> createFeedResponse(RxDocumentServiceResponse response,
                                                    CosmosItemSerializer itemSerializer,
                                                    Class<T> cls);
 
+            <T> FeedResponse<T> createNonServiceFeedResponse(List<T> items,
+                                                   boolean isChangeFeed,
+                                                   boolean isNoChanges);
+
             <T> FeedResponse<T> createChangeFeedResponse(RxDocumentServiceResponse response,
                                                    CosmosItemSerializer itemSerializer,
                                                    Class<T> cls);