Added API to inject custom serializer for point read and query

sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/ReadMyWriteWorkflow.java

@@ -3,9 +3,9 @@
 
 package com.azure.cosmos.benchmark;
 
-import com.azure.cosmos.BridgeInternal;
 import com.azure.cosmos.CosmosBridgeInternal;
 import com.azure.cosmos.CosmosException;
+import com.azure.cosmos.CosmosItemSerializer;
 import com.azure.cosmos.implementation.AsyncDocumentClient;
 import com.azure.cosmos.implementation.CosmosPagedFluxOptions;
 import com.azure.cosmos.implementation.Database;
@@ -194,12 +194,12 @@ private Flux<Document> writeDocument(Integer i) {
         String randomVal = UUID.randomUUID().toString();
         Document document = new Document();
         document.setId(idString);
-        BridgeInternal.setProperty(document, partitionKey, idString);
-        BridgeInternal.setProperty(document, QUERY_FIELD_NAME, randomVal);
-        BridgeInternal.setProperty(document, "dataField1", randomVal);
-        BridgeInternal.setProperty(document, "dataField2", randomVal);
-        BridgeInternal.setProperty(document, "dataField3", randomVal);
-        BridgeInternal.setProperty(document, "dataField4", randomVal);
+        document.set(partitionKey, idString, CosmosItemSerializer.DEFAULT_SERIALIZER);
+        document.set(QUERY_FIELD_NAME, randomVal, CosmosItemSerializer.DEFAULT_SERIALIZER);
+        document.set("dataField1", randomVal, CosmosItemSerializer.DEFAULT_SERIALIZER);
+        document.set("dataField2", randomVal, CosmosItemSerializer.DEFAULT_SERIALIZER);
+        document.set("dataField3", randomVal, CosmosItemSerializer.DEFAULT_SERIALIZER);
+        document.set("dataField4", randomVal, CosmosItemSerializer.DEFAULT_SERIALIZER);
 
         Integer key = i == null ? cacheKey() : i;
         return client.createDocument(getCollectionLink(), document, null, false)

sdk/cosmos/azure-cosmos-encryption/CHANGELOG.md

@@ -3,6 +3,7 @@
 ### 2.11.0-beta.1 (Unreleased)
 
 #### Features Added
+* Added public APIs `getCustomeSerializer` and `setCustomSerializer` to allow customers to specify custom payload transformations or serialization settings. - See [PR 38997](https://github.com/Azure/azure-sdk-for-java/pull/38997)
 
 #### Breaking Changes
 

sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/CosmosEncryptionAsyncClient.java

@@ -11,6 +11,7 @@
 import com.azure.cosmos.CosmosAsyncContainer;
 import com.azure.cosmos.CosmosAsyncDatabase;
 import com.azure.cosmos.CosmosException;
+import com.azure.cosmos.CosmosItemSerializer;
 import com.azure.cosmos.encryption.implementation.Constants;
 import com.azure.cosmos.encryption.implementation.EncryptionImplementationBridgeHelpers;
 import com.azure.cosmos.encryption.implementation.keyprovider.EncryptionKeyStoreProviderImpl;
@@ -46,6 +47,8 @@ public final class CosmosEncryptionAsyncClient implements Closeable {
     private final EncryptionKeyStoreProviderImpl encryptionKeyStoreProviderImpl;
     private final static ImplementationBridgeHelpers.CosmosAsyncClientEncryptionKeyHelper.CosmosAsyncClientEncryptionKeyAccessor cosmosAsyncClientEncryptionKeyAccessor = ImplementationBridgeHelpers.CosmosAsyncClientEncryptionKeyHelper.getCosmosAsyncClientEncryptionKeyAccessor();
 
+    private final static ImplementationBridgeHelpers.CosmosAsyncClientHelper.CosmosAsyncClientAccessor cosmosAsyncClientAccessor = ImplementationBridgeHelpers.CosmosAsyncClientHelper.getCosmosAsyncClientAccessor();
+
     CosmosEncryptionAsyncClient(CosmosAsyncClient cosmosAsyncClient,
                                 KeyEncryptionKeyResolver keyEncryptionKeyResolver,
                                 String keyEncryptionKeyResolverName) {
@@ -230,6 +233,15 @@ private CosmosContainerProperties getContainerPropertiesWithVersionValidation(Co
         return cosmosContainerResponse.getProperties();
     }
 
+    CosmosItemSerializer getEffectiveItemSerializer(
+        CosmosItemSerializer requestOptionsItemSerializer) {
+
+        return cosmosAsyncClientAccessor
+            .getEffectiveItemSerializer(
+                this.getCosmosAsyncClient(),
+                requestOptionsItemSerializer);
+    }
+
     static {
         EncryptionImplementationBridgeHelpers.CosmosEncryptionAsyncClientHelper.seCosmosEncryptionAsyncClientAccessor(new EncryptionImplementationBridgeHelpers.CosmosEncryptionAsyncClientHelper.CosmosEncryptionAsyncClientAccessor() {
             @Override

sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/implementation/CosmosEncryptionQueryTransformer.java

@@ -3,49 +3,58 @@
 package com.azure.cosmos.encryption.implementation;
 
 import com.azure.cosmos.BridgeInternal;
+import com.azure.cosmos.CosmosItemSerializer;
 import com.azure.cosmos.implementation.CosmosPagedFluxOptions;
-import com.azure.cosmos.implementation.ItemDeserializer;
+import com.azure.cosmos.implementation.ObjectNodeMap;
+import com.azure.cosmos.implementation.PrimitiveJsonNodeMap;
+import com.azure.cosmos.implementation.Utils;
 import com.azure.cosmos.implementation.query.Transformer;
 import com.azure.cosmos.models.FeedResponse;
 import com.azure.cosmos.models.ModelBridgeInternal;
 import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.scheduler.Scheduler;
 
 import java.util.List;
+import java.util.Map;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
+import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
+
 public class CosmosEncryptionQueryTransformer<T> implements Transformer<T> {
     private final Scheduler encryptionScheduler;
     private final EncryptionProcessor encryptionProcessor;
-    private final ItemDeserializer itemDeserializer;
     private final Class<T> classType;
     private final boolean isChangeFeed;
 
     public CosmosEncryptionQueryTransformer(
         Scheduler encryptionScheduler,
         EncryptionProcessor encryptionProcessor,
-        ItemDeserializer itemDeserializer,
         Class<T> classType,
         Boolean isChangeFeed) {
+
         this.encryptionScheduler = encryptionScheduler;
         this.encryptionProcessor = encryptionProcessor;
-        this.itemDeserializer = itemDeserializer;
         this.classType = classType;
         this.isChangeFeed = isChangeFeed;
     }
 
     @Override
-    public Function<CosmosPagedFluxOptions, Flux<FeedResponse<T>>> transform(Function<CosmosPagedFluxOptions, Flux<FeedResponse<JsonNode>>> func) {
-        return queryDecryptionTransformer(this.classType, this.isChangeFeed, func);
+    public Function<CosmosPagedFluxOptions, Flux<FeedResponse<T>>> transform(
+        Function<CosmosPagedFluxOptions, Flux<FeedResponse<JsonNode>>> func,
+        CosmosItemSerializer effectiveSerializer) {
+
+        return queryDecryptionTransformer(this.classType, this.isChangeFeed, func, effectiveSerializer);
     }
 
-    private <T> Function<CosmosPagedFluxOptions, Flux<FeedResponse<T>>> queryDecryptionTransformer(
-        Class<T> classType,
+    private <TTransform> Function<CosmosPagedFluxOptions, Flux<FeedResponse<TTransform>>> queryDecryptionTransformer(
+        Class<TTransform> classType,
         boolean isChangeFeed,
-        Function<CosmosPagedFluxOptions, Flux<FeedResponse<JsonNode>>> func) {
+        Function<CosmosPagedFluxOptions, Flux<FeedResponse<JsonNode>>> func,
+        CosmosItemSerializer effectiveSerializer) {
         return func.andThen(flux ->
             flux.publishOn(encryptionScheduler)
                 .flatMap(
@@ -55,9 +64,27 @@ private <T> Function<CosmosPagedFluxOptions, Flux<FeedResponse<T>>> queryDecrypt
                             ModelBridgeInternal.getNoChangesFromFeedResponse(page)
                             : false;
                         List<Mono<JsonNode>> jsonNodeArrayMonoList =
-                            page.getResults().stream().map(jsonNode -> decryptResponseNode(jsonNode)).collect(Collectors.toList());
+                            page.getResults().stream().map(jsonNode -> decryptResponseNode(jsonNode))
+                                .collect(Collectors.toList());
                         return Flux.concat(jsonNodeArrayMonoList).map(
-                            item -> this.itemDeserializer.convert(classType, item)
+                            item -> {
+                                Map<String, Object> decryptedJsonTree;
+
+                                if (item instanceof ObjectNode) {
+                                    decryptedJsonTree = new ObjectNodeMap((ObjectNode)item);
+                                } else if (item.isValueNode()) {
+                                    decryptedJsonTree = new PrimitiveJsonNodeMap(item);
+                                } else {
+                                    return Utils.getSimpleObjectMapper()
+                                                             .convertValue(
+                                                                 item,
+                                                                 classType);
+                                }
+
+                                return effectiveSerializer.deserialize(
+                                    decryptedJsonTree,
+                                    classType);
+                            }
                         ).collectList().map(itemList -> BridgeInternal.createFeedResponseWithQueryMetrics(itemList,
                             page.getResponseHeaders(),
                             BridgeInternal.queryMetricsFromFeedResponse(page),

sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/implementation/CosmosResponseFactory.java

@@ -3,9 +3,9 @@
 
 package com.azure.cosmos.encryption.implementation;
 
+import com.azure.cosmos.CosmosItemSerializer;
 import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
 import com.azure.cosmos.implementation.ImplementationBridgeHelpers.CosmosItemResponseHelper.CosmosItemResponseBuilderAccessor;
-import com.azure.cosmos.implementation.ItemDeserializer;
 import com.azure.cosmos.models.CosmosItemResponse;
 
 public class CosmosResponseFactory {
@@ -17,10 +17,11 @@ public CosmosResponseFactory() {
     }
 
     public <T> CosmosItemResponse<T> createItemResponse(CosmosItemResponse<byte[]> responseMessage,
-                                                        Class<T> classType) {
+                                                        Class<T> classType,
+                                                        CosmosItemSerializer effectiveItemSerializer) {
         return cosmosItemResponseBuilderAccessor.createCosmosItemResponse(
             responseMessage,
             classType,
-            new ItemDeserializer.JsonDeserializer());
+            effectiveItemSerializer);
     }
 }

sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/implementation/EncryptionProcessor.java

@@ -5,6 +5,7 @@
 
 import com.azure.cosmos.BridgeInternal;
 import com.azure.cosmos.CosmosAsyncContainer;
+import com.azure.cosmos.CosmosItemSerializer;
 import com.azure.cosmos.encryption.CosmosEncryptionAsyncClient;
 import com.azure.cosmos.encryption.implementation.keyprovider.EncryptionKeyStoreProviderImpl;
 import com.azure.cosmos.encryption.implementation.mdesrc.cryptography.EncryptionType;
@@ -54,10 +55,10 @@
 
 public class EncryptionProcessor {
     private final static Logger LOGGER = LoggerFactory.getLogger(EncryptionProcessor.class);
-    private CosmosEncryptionAsyncClient encryptionCosmosClient;
-    private CosmosAsyncContainer cosmosAsyncContainer;
-    private EncryptionSettings encryptionSettings;
-    private AtomicBoolean isEncryptionSettingsInitDone;
+    private final CosmosEncryptionAsyncClient encryptionCosmosClient;
+    private final CosmosAsyncContainer cosmosAsyncContainer;
+    private final EncryptionSettings encryptionSettings;
+    private final AtomicBoolean isEncryptionSettingsInitDone;
     private ClientEncryptionPolicy clientEncryptionPolicy;
     private String containerRid;
     private String databaseRid;
@@ -254,12 +255,15 @@ public Mono<byte[]> encrypt(byte[] payload) {
                 payload == null ? null : payload.length,
                 Thread.currentThread().getName());
         }
-        ObjectNode itemJObj = Utils.parse(payload, ObjectNode.class);
+        ObjectNode itemJObj = Utils.parse(payload, ObjectNode.class, CosmosItemSerializer.DEFAULT_SERIALIZER);
         return encrypt(itemJObj);
     }
 
     public Mono<byte[]> encrypt(JsonNode itemJObj) {
-        return encryptObjectNode(itemJObj).map(encryptedObjectNode -> EncryptionUtils.serializeJsonToByteArray(EncryptionUtils.getSimpleObjectMapper(), encryptedObjectNode));
+        return encryptObjectNode(itemJObj).map(
+            encryptedObjectNode -> EncryptionUtils.serializeJsonToByteArray(
+                CosmosItemSerializer.DEFAULT_SERIALIZER,
+                encryptedObjectNode));
     }
 
     public Mono<JsonNode> encryptPatchNode(JsonNode itemObj, String patchPropertyPath) {
@@ -273,7 +277,8 @@ public Mono<JsonNode> encryptPatchNode(JsonNode itemObj, String patchPropertyPat
 
             for (ClientEncryptionIncludedPath includedPath : this.clientEncryptionPolicy.getIncludedPaths()) {
                 String propertyName = includedPath.getPath().substring(1);
-                if (patchPropertyPath.substring(1).equals(propertyName)) {
+                String relativePatchPropertyPath = patchPropertyPath.substring(1);
+                if (relativePatchPropertyPath.equals(propertyName) || relativePatchPropertyPath.startsWith(propertyName + "/")) {
                     if (itemObj.isValueNode()) {
                         return this.encryptionSettings.getEncryptionSettingForPropertyAsync(propertyName,
                             this).flatMap(settings -> {
@@ -497,7 +502,7 @@ public byte[] encryptAndSerializeValue(EncryptionSettings encryptionSettings, Ob
         return cipherTextWithTypeMarker;
     }
 
-    public Mono<Pair<byte[], JsonNode>> decrypt(byte[] input) {
+    public Mono<Pair<byte[], JsonNode>> decrypt(byte[] input, CosmosItemSerializer itemSerializer) {
         if (LOGGER.isDebugEnabled()) {
             LOGGER.debug("Encrypting byte[] of size [{}] on thread [{}]",
                 input == null ? null : input.length,
@@ -508,13 +513,13 @@ public Mono<Pair<byte[], JsonNode>> decrypt(byte[] input) {
             return Mono.empty();
         }
 
-        ObjectNode itemJObj = Utils.parse(input, ObjectNode.class);
-        return decrypt(itemJObj);
+        ObjectNode itemJObj = Utils.parse(input, ObjectNode.class, CosmosItemSerializer.DEFAULT_SERIALIZER);
+        return decrypt(itemJObj, itemSerializer);
     }
 
-    public Mono<Pair<byte[], JsonNode>> decrypt(JsonNode itemJObj) {
+    public Mono<Pair<byte[], JsonNode>> decrypt(JsonNode itemJObj, CosmosItemSerializer itemSerializer) {
         return decryptJsonNode(itemJObj).map(decryptedObjectNode -> Pair.of(
-            EncryptionUtils.serializeJsonToByteArray(EncryptionUtils.getSimpleObjectMapper(), decryptedObjectNode),
+            EncryptionUtils.serializeJsonToByteArray(itemSerializer, decryptedObjectNode),
             decryptedObjectNode));
     }
 

sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/implementation/EncryptionUtils.java

@@ -3,6 +3,7 @@
 
 package com.azure.cosmos.encryption.implementation;
 
+import com.azure.cosmos.CosmosItemSerializer;
 import com.azure.cosmos.implementation.Utils;
 import com.fasterxml.jackson.core.JsonParser;
 import com.fasterxml.jackson.databind.DeserializationFeature;
@@ -21,8 +22,8 @@ public class EncryptionUtils {
         EncryptionUtils.simpleObjectMapper.configure(DeserializationFeature.ACCEPT_FLOAT_AS_INT, false);
     }
 
-    public static byte[] serializeJsonToByteArray(ObjectMapper objectMapper, Object object) {
-        return toByteArray(Utils.serializeJsonToByteBuffer(objectMapper, object));
+    public static byte[] serializeJsonToByteArray(CosmosItemSerializer itemSerializer, Object object) {
+        return toByteArray(Utils.serializeJsonToByteBuffer(itemSerializer, object, null));
     }
 
     public static ObjectMapper getSimpleObjectMapper() {

sdk/cosmos/azure-cosmos-encryption/src/main/java/com/azure/cosmos/encryption/models/SqlQuerySpecWithEncryption.java

@@ -3,6 +3,7 @@
 
 package com.azure.cosmos.encryption.models;
 
+import com.azure.cosmos.CosmosItemSerializer;
 import com.azure.cosmos.encryption.CosmosEncryptionAsyncContainer;
 import com.azure.cosmos.encryption.implementation.Constants;
 import com.azure.cosmos.encryption.implementation.EncryptionImplementationBridgeHelpers;
@@ -27,8 +28,8 @@
  * Represents a SQL query with encryption parameters in the Azure Cosmos DB database service.
  */
 public final class SqlQuerySpecWithEncryption {
-    private SqlQuerySpec sqlQuerySpec;
-    private HashMap<String, SqlParameter> encryptionParamMap = new HashMap<>();
+    private final SqlQuerySpec sqlQuerySpec;
+    private final HashMap<String, SqlParameter> encryptionParamMap = new HashMap<>();
     private final EncryptionImplementationBridgeHelpers.CosmosEncryptionAsyncContainerHelper.CosmosEncryptionAsyncContainerAccessor cosmosEncryptionAsyncContainerAccessor = EncryptionImplementationBridgeHelpers.CosmosEncryptionAsyncContainerHelper.getCosmosEncryptionAsyncContainerAccessor();
 
     /**
@@ -83,9 +84,9 @@ Mono<Void> addEncryptionParameterAsync(String path, SqlParameter sqlParameter,
                                     }
                                 }
                                 byte[] valueByte =
-                                    EncryptionUtils.serializeJsonToByteArray(EncryptionUtils.getSimpleObjectMapper(),
+                                    EncryptionUtils.serializeJsonToByteArray(CosmosItemSerializer.DEFAULT_SERIALIZER,
                                         sqlParameter.getValue(Object.class));
-                                JsonNode itemJObj = Utils.parse(valueByte, JsonNode.class);
+                                JsonNode itemJObj = Utils.parse(valueByte, JsonNode.class, CosmosItemSerializer.DEFAULT_SERIALIZER);
                                 Pair<EncryptionProcessor.TypeMarker, byte[]> typeMarkerPair =
                                     EncryptionProcessor.toByteArray(itemJObj);
                                 byte[] cipherText =

sdk/cosmos/azure-cosmos-encryption/src/main/java/module-info.java

@@ -10,4 +10,5 @@
     // public API surface area
     exports com.azure.cosmos.encryption;
     exports com.azure.cosmos.encryption.models;
+    opens com.azure.cosmos.encryption to com.fasterxml.jackson.databind, com.fasterxml.jackson.module.afterburner;
 }

sdk/cosmos/azure-cosmos-encryption/src/test/java/com/azure/cosmos/encryption/CosmosEncryptionAsyncClientTest.java

@@ -17,7 +17,6 @@
 import java.lang.reflect.Method;
 
 public abstract class CosmosEncryptionAsyncClientTest implements ITest {
-
     private static final ImplementationBridgeHelpers.CosmosClientBuilderHelper.CosmosClientBuilderAccessor cosmosClientBuilderAccessor =
         ImplementationBridgeHelpers.CosmosClientBuilderHelper.getCosmosClientBuilderAccessor();
 
@@ -36,17 +35,13 @@ public final CosmosClientBuilder getClientBuilder() {
         return this.clientBuilder;
     }
 
-    public final ConnectionPolicy getConnectionPolicy() {
-        return cosmosClientBuilderAccessor.getConnectionPolicy(this.clientBuilder);
-    }
-
     @Override
     public final String getTestName() {
         return this.testName;
     }
 
     @BeforeMethod(alwaysRun = true)
-    public final void setTestName(Method method) {
+    public final void setTestName(Method method, Object[] row) {
         String testClassAndMethodName = Strings.lenientFormat("%s::%s",
             method.getDeclaringClass().getSimpleName(),
             method.getName());
@@ -66,10 +61,19 @@ public final void setTestName(Method method) {
         } else {
             this.testName = testClassAndMethodName;
         }
+
+        String suffix = this.resolveTestNameSuffix(row);
+        if (suffix != null && !suffix.isEmpty()) {
+            this.testName += "(" + suffix + ")";
+        }
     }
 
     @AfterMethod(alwaysRun = true)
     public final void unsetTestName() {
         this.testName = null;
     }
+
+    public String resolveTestNameSuffix(Object[] row) {
+        return "";
+    }
 }