Encryption filling gaps (#13634)

* minor cleanup * minor cleanup * encryptor * static field name correction, fixed pre-condition check * move encryption hooks out of sdk * encryption * undid public api change to CosmosItemResponse * cleanup * cleanup * code review comments addressed * fixed compilation error * minor code style cleanup * check style fix * javadoc fix * javadoc fix * check-style * query basic support * encryption for query * minor package move, cleanup * code style fix * fixed check style * upsert, delete, replace * more tests * addressed code review comments
AzureSDKAutomation · Jul 30, 2020 · 8ea6bb9 · 8ea6bb9
1 parent acbcf2c
commit 8ea6bb9
Show file tree

Hide file tree

Showing 4 changed files with 497 additions and 38 deletions.
diff --git a/...-encryption/src/main/java/com/azure/cosmos/encryption/EncryptionCosmosAsyncContainer.java b/...-encryption/src/main/java/com/azure/cosmos/encryption/EncryptionCosmosAsyncContainer.java
@@ -56,7 +56,61 @@ private Mono<CosmosItemResponse<byte[]>> createItemStream(byte[] payload,
             encryptionItemRequestOptions);
 
         return response
-            .subscribeOn(Schedulers.elastic())
+            .publishOn(Schedulers.elastic())
+            .map(rsp -> {
+                    EncryptionModelBridgeInternal.setByteArrayContent(
+                        rsp,
+                        decryptResponseAsync(
+                            EncryptionModelBridgeInternal.getByteArrayContent(rsp),
+                            encryptionItemRequestOptions.getDecryptionResultHandler()));
+                    return rsp;
+                }
+            );
+    }
+
+    private Mono<CosmosItemResponse<byte[]>> replaceItemStream(byte[] payload,
+                                                               String itemId,
+                                                               PartitionKey partitionKey,
+                                                               EncryptionItemRequestOptions encryptionItemRequestOptions) {
+        Preconditions.checkNotNull(payload, "payload can't be null");
+
+        // TODO: add diagnostics
+        assert encryptionItemRequestOptions != null && encryptionItemRequestOptions.getEncryptionOptions() != null;
+        payload = EncryptionProcessor.encryptAsync(payload, encryptor,
+            encryptionItemRequestOptions.getEncryptionOptions());
+
+        Mono<CosmosItemResponse<byte[]>> response = container.replaceItem(payload,
+            itemId,
+            partitionKey,
+            encryptionItemRequestOptions);
+
+        return response
+            .publishOn(Schedulers.elastic())
+            .map(rsp -> {
+                    EncryptionModelBridgeInternal.setByteArrayContent(
+                        rsp,
+                        decryptResponseAsync(
+                            EncryptionModelBridgeInternal.getByteArrayContent(rsp),
+                            encryptionItemRequestOptions.getDecryptionResultHandler()));
+                    return rsp;
+                }
+            );
+    }
+
+    private Mono<CosmosItemResponse<byte[]>> upsertItemStream(byte[] payload,
+                                                              PartitionKey partitionKey,
+                                                              EncryptionItemRequestOptions encryptionItemRequestOptions) {
+        Preconditions.checkNotNull(payload, "payload can't be null");
+
+        // TODO: add diagnostics
+        assert encryptionItemRequestOptions != null && encryptionItemRequestOptions.getEncryptionOptions() != null;
+        payload = EncryptionProcessor.encryptAsync(payload, encryptor,
+            encryptionItemRequestOptions.getEncryptionOptions());
+
+        Mono<CosmosItemResponse<byte[]>> response = container.upsertItem(payload, partitionKey,
+            encryptionItemRequestOptions);
+
+        return response
             .publishOn(Schedulers.elastic())
             .map(rsp -> {
                     EncryptionModelBridgeInternal.setByteArrayContent(
@@ -107,6 +161,103 @@ public <T> Mono<CosmosItemResponse<T>> createItem(T item,
         }
     }
 
+    /**
+     * Deletes the item.
+     * <p>
+     * After subscription the operation will be performed.
+     * The {@link Mono} upon successful completion will contain a single Cosmos item response with the deleted item.
+     *
+     * @param itemId id of the item.
+     * @param partitionKey partitionKey of the item.
+     * @param options the request options.
+     * @return an {@link Mono} containing the Cosmos item resource response.
+     */
+    public Mono<CosmosItemResponse<Object>> deleteItem(String itemId,
+                                                       PartitionKey partitionKey,
+                                                       CosmosItemRequestOptions options) {
+
+        return container.deleteItem(itemId, partitionKey, options);
+    }
+
+
+    /**
+     * upserts item and encrypts the requested fields
+     *
+     * @param item the Cosmos item represented as a POJO or Cosmos item object.
+     * @param partitionKey the partition key.
+     * @param requestOptions request option
+     * @param <T> serialization class type
+     * @return result
+     */
+    public <T> Mono<CosmosItemResponse<T>> upsertItem(T item,
+                                                      PartitionKey partitionKey,
+                                                      CosmosItemRequestOptions requestOptions) {
+        Preconditions.checkNotNull(item, "item");
+        if (requestOptions == null) {
+            requestOptions = new CosmosItemRequestOptions();
+        }
+
+        EncryptionItemRequestOptions encryptionItemRequestOptions = Utils.as(requestOptions,
+            EncryptionItemRequestOptions.class);
+
+        if (encryptionItemRequestOptions != null && encryptionItemRequestOptions.getEncryptionOptions() != null) {
+            Preconditions.checkArgument(partitionKey != null, "partitionKey cannot be null for operations using "
+                + "EncryptionContainer.");
+
+            byte[] payload = cosmosSerializerToStream(item);
+            Mono<CosmosItemResponse<byte[]>> result = this.upsertItemStream(payload, partitionKey,
+                encryptionItemRequestOptions);
+
+            return result.map(rsp -> (CosmosItemResponse<T>) this.responseFactory.createItemResponse(rsp,
+                item.getClass()));
+
+        } else {
+            return container.upsertItem(item, partitionKey, requestOptions);
+        }
+    }
+
+    /**
+     * replaces item and encrypts the requested fields
+     *
+     * @param item the Cosmos item represented as a POJO or Cosmos item object.
+     * @param itemId the item id.
+     * @param partitionKey the partition key.
+     * @param requestOptions request option
+     * @param <T> serialization class type
+     * @return result
+     */
+    public <T> Mono<CosmosItemResponse<T>> replaceItem(T item,
+                                                       String itemId,
+                                                       PartitionKey partitionKey,
+                                                       CosmosItemRequestOptions requestOptions) {
+        Preconditions.checkNotNull(item, "item");
+        if (requestOptions == null) {
+            requestOptions = new CosmosItemRequestOptions();
+        }
+
+        EncryptionItemRequestOptions encryptionItemRequestOptions = Utils.as(requestOptions,
+            EncryptionItemRequestOptions.class);
+
+        if (encryptionItemRequestOptions != null && encryptionItemRequestOptions.getEncryptionOptions() != null) {
+            Preconditions.checkArgument(partitionKey != null, "partitionKey cannot be null for operations using "
+                + "EncryptionContainer.");
+
+            byte[] payload = cosmosSerializerToStream(item);
+            Mono<CosmosItemResponse<byte[]>> result = this.replaceItemStream(payload,
+                itemId,
+                partitionKey,
+                encryptionItemRequestOptions);
+
+            return result.map(rsp -> (CosmosItemResponse<T>) this.responseFactory.createItemResponse(rsp,
+                item.getClass()));
+
+        } else {
+            return container.replaceItem(item, itemId, partitionKey, requestOptions);
+        }
+    }
+
+
+
     private Mono<CosmosItemResponse<byte[]>> readItemStream(String id,
                                                             PartitionKey partitionKey,
                                                             CosmosItemRequestOptions requestOptions) {
@@ -117,7 +268,6 @@ private Mono<CosmosItemResponse<byte[]>> readItemStream(String id,
 
         return responseMessageAsync
             .publishOn(Schedulers.elastic())
-            .subscribeOn(Schedulers.elastic())
             .map(
                 responseMessage -> {
                     Consumer<DecryptionResult> decryptionErroHandler = null;
@@ -247,3 +397,4 @@ private byte[] decryptResponseAsync(
         }
     }
 }
+