Adds CustomItemSerializer to CosmsoRequestContext/CosmosRequestOptions (#48963)

* Adds CustomItemSerializer to CosmsoRequestContext/CosmosRequestOptions

* Update CHANGELOG.md

Author: FabianMeiswinkel

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

@@ -7,10 +7,12 @@
 package com.azure.cosmos;
 
 import com.azure.cosmos.SuperFlakyTestRetryAnalyzer;
+import com.azure.cosmos.implementation.DefaultCosmosItemSerializer;
 import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
 import com.azure.cosmos.implementation.InternalObjectNode;
 import com.azure.cosmos.implementation.OverridableRequestOptions;
 import com.azure.cosmos.implementation.TestConfigurations;
+import com.azure.cosmos.implementation.Utils;
 import com.azure.cosmos.models.CosmosBatch;
 import com.azure.cosmos.models.CosmosBatchResponse;
 import com.azure.cosmos.models.CosmosBulkExecutionOptions;
@@ -79,8 +81,11 @@ public class OperationPoliciesTest extends TestSuiteBase {
     private static final String MAX_PREFETCH_PAGE_COUNT = "maxPrefetchPageCount";
     private static final String QUERY_NAME = "queryName";
     private static final String KEYWORD_IDENTIFIERS = "keywordIdentifiers";
-    private static final String[] optionLabels = {E2E_TIMEOUT, CONSISTENCY_LEVEL, CONTENT_RESPONSE_ON_WRITE, NON_IDEMPOTENT_WRITE_RETRIES, BYPASS_CACHE, THROUGHPUT_CONTROL_GROUP_NAME, REQUEST_CHARGE_THRESHOLD, SCAN_IN_QUERY, EXCLUDE_REGIONS, MAX_DEGREE_OF_PARALLELISM, MAX_BUFFERED_ITEM_COUNT, RESPONSE_CONTINUATION_TOKEN_LIMIT_KB, MAX_ITEM_COUNT, QUERY_METRICS, INDEX_METRICS, MAX_PREFETCH_PAGE_COUNT, QUERY_NAME, KEYWORD_IDENTIFIERS, READ_CONSISTENCY_STRATEGY};
-    private static final String[] initialOptions = {"20", "Session", "true", "false", "false", "default", "2000", "false", "East US 2", "2", "100", "200", "30", "false", "false", "10", "QueryName", "59409493805", "Default"};
+    private static final String CUSTOM_ITEM_SERIALIZER = "customItemSerializer";
+    private static final String[] optionLabels = {E2E_TIMEOUT, CONSISTENCY_LEVEL, CONTENT_RESPONSE_ON_WRITE, NON_IDEMPOTENT_WRITE_RETRIES, BYPASS_CACHE, THROUGHPUT_CONTROL_GROUP_NAME, REQUEST_CHARGE_THRESHOLD, SCAN_IN_QUERY, EXCLUDE_REGIONS, MAX_DEGREE_OF_PARALLELISM, MAX_BUFFERED_ITEM_COUNT, RESPONSE_CONTINUATION_TOKEN_LIMIT_KB, MAX_ITEM_COUNT, QUERY_METRICS, INDEX_METRICS, MAX_PREFETCH_PAGE_COUNT, QUERY_NAME, KEYWORD_IDENTIFIERS, READ_CONSISTENCY_STRATEGY, CUSTOM_ITEM_SERIALIZER};
+    private static final String[] initialOptions = {"20", "Session", "true", "false", "false", "default", "2000", "false", "East US 2", "2", "100", "200", "30", "false", "false", "10", "QueryName", "59409493805", "Default", "initial"};
+    private static final CosmosItemSerializer INITIAL_CUSTOM_ITEM_SERIALIZER = CosmosItemSerializer.DEFAULT_SERIALIZER;
+    private static final CosmosItemSerializer CHANGED_CUSTOM_ITEM_SERIALIZER = new DefaultCosmosItemSerializer(Utils.getSimpleObjectMapper());
 
     @Factory(dataProvider = "clientBuildersWithApplyPolicies")
     public OperationPoliciesTest(CosmosClientBuilder clientBuilder) {
@@ -107,7 +112,8 @@ private static void createReadDeleteBatchEtcOptions(String operationType, String
                     .setIntegratedCacheBypassed(Boolean.parseBoolean(prop.getProperty(BYPASS_CACHE))))
                 .setThroughputControlGroupName(prop.getProperty(THROUGHPUT_CONTROL_GROUP_NAME))
                 .setExcludeRegions(new ArrayList<>(Arrays.asList(prop.getProperty(EXCLUDE_REGIONS).split(","))))
-                .setKeywordIdentifiers(new HashSet<>(Arrays.asList(prop.getProperty(KEYWORD_IDENTIFIERS).split(","))));
+                .setKeywordIdentifiers(new HashSet<>(Arrays.asList(prop.getProperty(KEYWORD_IDENTIFIERS).split(","))))
+                .setCustomItemSerializer(resolveCustomItemSerializer());
         }
     }
 
@@ -133,7 +139,8 @@ private static void createQueryReadAllItemsOptions(String operationType, String
                 .setQueryName(prop.getProperty(QUERY_NAME))
                 .setConsistencyLevel(ConsistencyLevel.fromServiceSerializedFormat(prop.getProperty(CONSISTENCY_LEVEL)))
                 .setReadConsistencyStrategy(ReadConsistencyStrategy.fromServiceSerializedFormat(prop.getProperty(READ_CONSISTENCY_STRATEGY)))
-                .setKeywordIdentifiers(new HashSet<>(Arrays.asList(prop.getProperty(KEYWORD_IDENTIFIERS).split(","))));
+                .setKeywordIdentifiers(new HashSet<>(Arrays.asList(prop.getProperty(KEYWORD_IDENTIFIERS).split(","))))
+                .setCustomItemSerializer(resolveCustomItemSerializer());
         }
     }
 
@@ -152,15 +159,17 @@ private static void createReadManyOptions(String spanName, CosmosRequestOptions
                     .setIndexMetricsEnabled(Boolean.parseBoolean(prop.getProperty(INDEX_METRICS)))
                     .setKeywordIdentifiers(new HashSet<>(Arrays.asList(prop.getProperty(KEYWORD_IDENTIFIERS).split(","))))
                     .setConsistencyLevel(ConsistencyLevel.fromServiceSerializedFormat(prop.getProperty(CONSISTENCY_LEVEL)))
-                    .setReadConsistencyStrategy(ReadConsistencyStrategy.fromServiceSerializedFormat(prop.getProperty(READ_CONSISTENCY_STRATEGY)));
+                    .setReadConsistencyStrategy(ReadConsistencyStrategy.fromServiceSerializedFormat(prop.getProperty(READ_CONSISTENCY_STRATEGY)))
+                    .setCustomItemSerializer(resolveCustomItemSerializer());
         }
     }
 
     private static void createBulkOptions(String operationType, String spanName, CosmosRequestOptions cosmosRequestOptions) {
         if (operationType.equals("Batch") && spanName.contains("nonTransactionalBatch")) {
                 cosmosRequestOptions.setExcludeRegions((new ArrayList<>(Arrays.asList(prop.getProperty(EXCLUDE_REGIONS).split(",")))))
                     .setThroughputControlGroupName(prop.getProperty(THROUGHPUT_CONTROL_GROUP_NAME))
-                    .setKeywordIdentifiers(new HashSet<>(Arrays.asList(prop.getProperty(KEYWORD_IDENTIFIERS).split(","))));
+                    .setKeywordIdentifiers(new HashSet<>(Arrays.asList(prop.getProperty(KEYWORD_IDENTIFIERS).split(","))))
+                    .setCustomItemSerializer(resolveCustomItemSerializer());
         }
     }
 
@@ -171,8 +180,25 @@ private static void createChangeFeedOptions(String spanName, CosmosRequestOption
                     .setDiagnosticsThresholds(new CosmosDiagnosticsThresholds().setRequestChargeThreshold(Float.parseFloat(prop.getProperty(REQUEST_CHARGE_THRESHOLD))))
                     .setMaxPrefetchPageCount(Integer.parseInt(prop.getProperty(MAX_PREFETCH_PAGE_COUNT)))
                     .setMaxItemCount(Integer.parseInt(prop.getProperty(MAX_ITEM_COUNT)))
-                    .setKeywordIdentifiers(new HashSet<>(Arrays.asList(prop.getProperty(KEYWORD_IDENTIFIERS).split(","))));
+                    .setKeywordIdentifiers(new HashSet<>(Arrays.asList(prop.getProperty(KEYWORD_IDENTIFIERS).split(","))))
+                    .setCustomItemSerializer(resolveCustomItemSerializer());
+        }
+    }
+
+    private static CosmosItemSerializer resolveCustomItemSerializer() {
+        String value = prop.getProperty(CUSTOM_ITEM_SERIALIZER);
+        if ("changed".equals(value)) {
+            return CHANGED_CUSTOM_ITEM_SERIALIZER;
+        }
+        return INITIAL_CUSTOM_ITEM_SERIALIZER;
+    }
+
+    private static CosmosItemSerializer resolveExpectedCustomItemSerializer(String[] options) {
+        String value = options[19];
+        if ("changed".equals(value)) {
+            return CHANGED_CUSTOM_ITEM_SERIALIZER;
         }
+        return INITIAL_CUSTOM_ITEM_SERIALIZER;
     }
 
     @DataProvider
@@ -242,8 +268,8 @@ public void afterMethod() {
     @DataProvider(name = "changedOptions")
     private String[][] createChangedOptions() {
         return new String[][] {
-            { "8", "ConsistentPrefix", "true", "false", "true", "defaultChanged", "1000", "true", "West US 2", "4", "200", "400", "100", "false", "true", "20", "QueryNameChanged", "112", "LatestCommitted" },
-            { "4", "Eventual", "false", "true", "true", "defaultChanged", "1000", "true", "West US 2", "4", "200", "400", "100", "true", "false", "20", "QueryNameChanged", "221", "LatestCommitted" },
+            { "8", "ConsistentPrefix", "true", "false", "true", "defaultChanged", "1000", "true", "West US 2", "4", "200", "400", "100", "false", "true", "20", "QueryNameChanged", "112", "LatestCommitted", "changed" },
+            { "4", "Eventual", "false", "true", "true", "defaultChanged", "1000", "true", "West US 2", "4", "200", "400", "100", "true", "false", "20", "QueryNameChanged", "221", "LatestCommitted", "changed" },
             initialOptions
         };
     }
@@ -699,6 +725,7 @@ private void validateOptions(String[] options, CosmosItemResponse<?> response, b
        assertThat(requestOptions.getDiagnosticsThresholds().getRequestChargeThreshold()).isEqualTo(Float.parseFloat(options[6]));
        assertThat(requestOptions.getExcludedRegions()).isEqualTo(new ArrayList<>(Arrays.asList(options[8].split(","))));
        assertThat(requestOptions.getKeywordIdentifiers()).isEqualTo(new HashSet<>(Arrays.asList(options[17].split(","))));
+       assertThat(requestOptions.getCustomItemSerializer()).isSameAs(resolveExpectedCustomItemSerializer(options));
     }
 
     private void validateOptions(String[] options, CosmosBatchResponse response) {
@@ -708,6 +735,7 @@ private void validateOptions(String[] options, CosmosBatchResponse response) {
         assertThat(requestOptions.getDiagnosticsThresholds().getRequestChargeThreshold()).isEqualTo(Float.parseFloat(options[6]));
         assertThat(requestOptions.getExcludedRegions()).isEqualTo(new ArrayList<>(Arrays.asList(options[8].split(","))));
         assertThat(requestOptions.getKeywordIdentifiers()).isEqualTo(new HashSet<>(Arrays.asList(options[17].split(","))));
+        assertThat(requestOptions.getCustomItemSerializer()).isSameAs(resolveExpectedCustomItemSerializer(options));
     }
 
     private void validateOptions(String[] options, CosmosBulkItemResponse response) {
@@ -716,6 +744,7 @@ private void validateOptions(String[] options, CosmosBulkItemResponse response)
         assertThat(requestOptions.getExcludedRegions()).isEqualTo(new ArrayList<>(Arrays.asList(options[8].split(","))));
         assertThat(requestOptions.getThroughputControlGroupName()).isEqualTo(options[5]);
         assertThat(requestOptions.getKeywordIdentifiers()).isEqualTo(new HashSet<>(Arrays.asList(options[17].split(","))));
+        assertThat(requestOptions.getCustomItemSerializer()).isSameAs(resolveExpectedCustomItemSerializer(options));
     }
 
     private void validateOptions(String[] changedOptions, FeedResponse<InternalObjectNode> response, boolean isChangeFeed, boolean isReadMany) {
@@ -757,6 +786,7 @@ private void validateOptions(String[] changedOptions, FeedResponse<InternalObjec
             assertThat(requestOptions.getQueryNameOrDefault("")).isEqualTo(changedOptions[16]);
         }
         assertThat(requestOptions.getKeywordIdentifiers()).isEqualTo(new HashSet<>(Arrays.asList(changedOptions[17].split(","))));
+        assertThat(requestOptions.getCustomItemSerializer()).isSameAs(resolveExpectedCustomItemSerializer(changedOptions));
     }
 
     private void changeProperties(String[] values) {

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -7,6 +7,7 @@
 * Added new `readManyByPartitionKeys` API on `CosmosAsyncContainer` / `CosmosContainer` to bulk-query all documents matching a list of partition key values with better efficiency than issuing individual queries. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801)
 * Added `CosmosReadManyByPartitionKeysRequestOptions` - a dedicated request-options type for `readManyByPartitionKeys` that exposes `setContinuationToken(String)` for resuming previous invocations and `setMaxConcurrentBatchPrefetch(int)` to bound per-call prefetch parallelism. See [PR 48801](https://github.com/Azure/azure-sdk-for-java/pull/48801)
 * Added `CosmosReadManyByPartitionKeysRequestOptions.setMaxBatchSize(Integer)` to set the max. number of partition keys used for a single batch. See [PR 48930](https://github.com/Azure/azure-sdk-for-java/pull/48930)
+* Added `getCustomItemSerializer()` to `CosmosRequestContext` and `setCustomItemSerializer(CosmosItemSerializer)` to `CosmosRequestOptions` to allow overriding the custom item serializer via operation policies. - See [PR 48963](https://github.com/Azure/azure-sdk-for-java/pull/48963)
 
 #### Breaking Changes
 

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

@@ -203,6 +203,15 @@ public Set<String> getKeywordIdentifiers() {
         return requestOptions.getKeywordIdentifiers();
     }
 
+    /**
+     * Gets the custom item serializer.
+     *
+     * @return the custom item serializer. It could be null if not defined or called on an irrelevant operation.
+     */
+    public CosmosItemSerializer getCustomItemSerializer() {
+        return requestOptions.getCustomItemSerializer();
+    }
+
     ///////////////////////////////////////////////////////////////////////////////////////////
     // the following helper/accessor only helps to access this class outside of this package.//
     ///////////////////////////////////////////////////////////////////////////////////////////

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

@@ -139,6 +139,7 @@ public void setMinTargetMicroBatchSize(int minTargetMicroBatchSize) {
         this.minTargetMicroBatchSize = minTargetMicroBatchSize;
     }
 
+    @Override
     public CosmosItemSerializer getCustomItemSerializer() {
         return this.customSerializer;
     }
@@ -357,6 +358,7 @@ public void override(CosmosRequestOptions cosmosRequestOptions) {
         this.excludeRegions = overrideOption(cosmosRequestOptions.getExcludedRegions(), this.excludeRegions);
         this.throughputControlGroupName = overrideOption(cosmosRequestOptions.getThroughputControlGroupName(), this.throughputControlGroupName);
         this.keywordIdentifiers = overrideOption(cosmosRequestOptions.getKeywordIdentifiers(), this.keywordIdentifiers);
+        this.customSerializer = overrideOption(cosmosRequestOptions.getCustomItemSerializer(), this.customSerializer);
     }
 
 }

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

@@ -216,6 +216,7 @@ public Integer getResponseContinuationTokenLimitInKb() {
         return null;
     }
 
+    @Override
     public CosmosItemSerializer getCustomItemSerializer() {
         return this.customSerializer;
     }
@@ -421,6 +422,7 @@ public void override(CosmosRequestOptions cosmosRequestOptions) {
         this.throughputControlGroupName = overrideOption(cosmosRequestOptions.getThroughputControlGroupName(), this.throughputControlGroupName);
         this.thresholds = overrideOption(cosmosRequestOptions.getDiagnosticsThresholds(), this.thresholds);
         this.keywordIdentifiers = overrideOption(cosmosRequestOptions.getKeywordIdentifiers(), this.keywordIdentifiers);
+        this.customSerializer = overrideOption(cosmosRequestOptions.getCustomItemSerializer(), this.customSerializer);
     }
 
 }

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

@@ -511,6 +511,7 @@ public CosmosEndToEndOperationLatencyPolicyConfig getCosmosEndToEndLatencyPolicy
      * Gets the custom item serializer defined for this instance of request options
      * @return the custom item serializer
      */
+    @Override
     public CosmosItemSerializer getCustomItemSerializer() {
         return this.customSerializer;
     }
@@ -556,6 +557,7 @@ public void override(CosmosRequestOptions cosmosRequestOptions) {
         this.queryMetricsEnabled = overrideOption(cosmosRequestOptions.isQueryMetricsEnabled(), this.queryMetricsEnabled);
         this.responseContinuationTokenLimitInKb = overrideOption(cosmosRequestOptions.getResponseContinuationTokenLimitInKb(), this.responseContinuationTokenLimitInKb);
         this.keywordIdentifiers = overrideOption(cosmosRequestOptions.getKeywordIdentifiers(), this.keywordIdentifiers);
+        this.customSerializer = overrideOption(cosmosRequestOptions.getCustomItemSerializer(), this.customSerializer);
     }
 
     public RequestOptions applyToRequestOptions(RequestOptions requestOptions) {

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

@@ -72,6 +72,7 @@ public CosmosTransactionalBulkExecutionOptionsImpl() {
         this.customOptions = new HashMap<>();
     }
 
+    @Override
     public CosmosItemSerializer getCustomItemSerializer() {
         return this.customSerializer;
     }
@@ -286,6 +287,7 @@ public void override(CosmosRequestOptions cosmosRequestOptions) {
         this.excludeRegions = overrideOption(cosmosRequestOptions.getExcludedRegions(), this.excludeRegions);
         this.throughputControlGroupName = overrideOption(cosmosRequestOptions.getThroughputControlGroupName(), this.throughputControlGroupName);
         this.keywordIdentifiers = overrideOption(cosmosRequestOptions.getKeywordIdentifiers(), this.keywordIdentifiers);
+        this.customSerializer = overrideOption(cosmosRequestOptions.getCustomItemSerializer(), this.customSerializer);
     }
 
 }

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

@@ -5,6 +5,7 @@
 import com.azure.cosmos.ConsistencyLevel;
 import com.azure.cosmos.CosmosDiagnosticsThresholds;
 import com.azure.cosmos.CosmosEndToEndOperationLatencyPolicyConfig;
+import com.azure.cosmos.CosmosItemSerializer;
 import com.azure.cosmos.ReadConsistencyStrategy;
 import com.azure.cosmos.models.DedicatedGatewayRequestOptions;
 
@@ -156,4 +157,11 @@ public interface ReadOnlyRequestOptions {
      * @return the custom correlated ids.
      */
     Set<String> getKeywordIdentifiers();
+
+    /**
+     * Gets the custom item serializer.
+     *
+     * @return the custom item serializer. It could be null if not defined or called on an irrelevant operation.
+     */
+    CosmosItemSerializer getCustomItemSerializer();
 }

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

@@ -654,12 +654,18 @@ public void override(CosmosRequestOptions cosmosCommonRequestOptions) {
         this.thresholds = overrideOption(cosmosCommonRequestOptions.getDiagnosticsThresholds(), this.thresholds);
         this.endToEndOperationLatencyConfig = overrideOption(cosmosCommonRequestOptions.getCosmosEndToEndLatencyPolicyConfig(), this.endToEndOperationLatencyConfig);
         this.keywordIdentifiers = overrideOption(cosmosCommonRequestOptions.getKeywordIdentifiers(), this.keywordIdentifiers);
+        this.effectiveItemSerializer = overrideOption(cosmosCommonRequestOptions.getCustomItemSerializer(), this.effectiveItemSerializer);
     }
 
     public CosmosItemSerializer getEffectiveItemSerializer() {
         return this.effectiveItemSerializer;
     }
 
+    @Override
+    public CosmosItemSerializer getCustomItemSerializer() {
+        return this.effectiveItemSerializer;
+    }
+
     public void setEffectiveItemSerializer(CosmosItemSerializer serializer) {
         this.effectiveItemSerializer = serializer;
     }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/CosmosRequestOptions.java

@@ -5,6 +5,7 @@
 import com.azure.cosmos.ConsistencyLevel;
 import com.azure.cosmos.CosmosDiagnosticsThresholds;
 import com.azure.cosmos.CosmosEndToEndOperationLatencyPolicyConfig;
+import com.azure.cosmos.CosmosItemSerializer;
 import com.azure.cosmos.ReadConsistencyStrategy;
 import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList;
 import com.azure.cosmos.util.Beta;
@@ -40,6 +41,7 @@ public final class CosmosRequestOptions {
     private String queryName;
     private Set<String> keywordIdentifiers;
     private ReadConsistencyStrategy readConsistencyStrategy;
+    private CosmosItemSerializer customItemSerializer;
     private static final Set<String> EMPTY_KEYWORD_IDENTIFIERS = Collections.unmodifiableSet(new HashSet<>());
 
     /**
@@ -453,4 +455,24 @@ public CosmosEndToEndOperationLatencyPolicyConfig getCosmosEndToEndLatencyPolicy
     public Set<String> getKeywordIdentifiers() {
         return this.keywordIdentifiers;
     }
+
+    /**
+     * Gets the custom item serializer.
+     *
+     * @return the custom item serializer.
+     */
+    public CosmosItemSerializer getCustomItemSerializer() {
+        return this.customItemSerializer;
+    }
+
+    /**
+     * Sets the custom item serializer.
+     *
+     * @param customItemSerializer the custom item serializer.
+     * @return current CosmosRequestOptions.
+     */
+    public CosmosRequestOptions setCustomItemSerializer(CosmosItemSerializer customItemSerializer) {
+        this.customItemSerializer = customItemSerializer;
+        return this;
+    }
 }