adding missing API upsert with PartitionKey (Azure#13497)

* partitionKey support for upsert * partitionKey support for upsert * Fixed potential null pointer issue Co-authored-by: Kushagra Thapar <[email protected]>
jmonty42 · Jul 27, 2020 · 089078c · 089078c
1 parent 7d12ed3
commit 089078c
Show file tree

Hide file tree

Showing 4 changed files with 137 additions and 60 deletions.
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
@@ -301,6 +301,25 @@ public <T> Mono<CosmosItemResponse<T>> upsertItem(T item, CosmosItemRequestOptio
         return withContext(context -> upsertItemInternal(item, requestOptions, context));
     }
 
+    /**
+     * Upserts an item.
+     * <p>
+     * After subscription the operation will be performed. The {@link Mono} upon
+     * successful completion will contain a single resource response with the
+     * upserted item. In case of failure the {@link Mono} will error.
+     *
+     * @param <T> the type parameter.
+     * @param item the item represented as a POJO or Item object to upsert.
+     * @param partitionKey the partition key.
+     * @param options the request options.
+     * @return an {@link Mono} containing the single resource response with the upserted item or an error.
+     */
+    public <T> Mono<CosmosItemResponse<T>> upsertItem(T item, PartitionKey partitionKey, CosmosItemRequestOptions options) {
+        final CosmosItemRequestOptions requestOptions = options == null ? new CosmosItemRequestOptions() : options;
+        ModelBridgeInternal.setPartitionKey(requestOptions, partitionKey);
+        return withContext(context -> upsertItemInternal(item, requestOptions, context));
+    }
+
     /**
      * Reads all the items in the current container.
      * <p>

diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java
@@ -199,6 +199,19 @@ public <T> CosmosItemResponse<T> upsertItem(T item, CosmosItemRequestOptions opt
         return this.blockItemResponse(this.asyncContainer.upsertItem(item, options));
     }
 
+    /**
+     * Upserts a item Cosmos sync item while specifying additional options.
+     *
+     * @param <T> the type parameter.
+     * @param item the item.
+     * @param partitionKey the partitionKey.
+     * @param options the options.
+     * @return the Cosmos item response.
+     */
+    public <T> CosmosItemResponse<T> upsertItem(T item, PartitionKey partitionKey, CosmosItemRequestOptions options) {
+        return this.blockItemResponse(this.asyncContainer.upsertItem(item, partitionKey, options));
+    }
+
     /**
      * Block cosmos item response.
      *

diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TestObject.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TestObject.java
@@ -0,0 +1,90 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos;
+
+import com.azure.cosmos.implementation.guava25.collect.ImmutableList;
+import com.azure.cosmos.rx.DocumentCrudTest;
+
+import java.util.List;
+import java.util.Objects;
+import java.util.UUID;
+
+public class TestObject {
+    private String id;
+    private String mypk;
+    private List<List<Integer>> sgmts;
+    private String stringProp;
+
+    public TestObject() {
+    }
+
+    public TestObject(String id, String mypk, List<List<Integer>> sgmts, String stringProp) {
+        this.id = id;
+        this.mypk = mypk;
+        this.sgmts = sgmts;
+        this.stringProp = stringProp;
+    }
+
+    public static TestObject create() {
+        return new TestObject(UUID.randomUUID().toString(), UUID.randomUUID().toString(), ImmutableList.of(ImmutableList.of(5)), UUID.randomUUID().toString());
+    }
+
+    public String getId() {
+        return id;
+    }
+
+    public void setId(String id) {
+        this.id = id;
+    }
+
+    public String getMypk() {
+        return mypk;
+    }
+
+    public void setMypk(String mypk) {
+        this.mypk = mypk;
+    }
+
+    public List<List<Integer>> getSgmts() {
+        return sgmts;
+    }
+
+    public void setSgmts(List<List<Integer>> sgmts) {
+        this.sgmts = sgmts;
+    }
+
+    /**
+     * Getter for property 'stringProp'.
+     *
+     * @return Value for property 'stringProp'.
+     */
+    public String getStringProp() {
+        return stringProp;
+    }
+
+    /**
+     * Setter for property 'stringProp'.
+     *
+     * @param stringProp Value to set for property 'stringProp'.
+     */
+    public void setStringProp(String stringProp) {
+        this.stringProp = stringProp;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        TestObject that = (TestObject) o;
+        return Objects.equals(id, that.id) &&
+            Objects.equals(mypk, that.mypk) &&
+            Objects.equals(sgmts, that.sgmts) &&
+            Objects.equals(stringProp, that.stringProp);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(id, mypk, sgmts, stringProp);
+    }
+}
diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/DocumentCrudTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/rx/DocumentCrudTest.java
@@ -5,6 +5,7 @@
 import com.azure.cosmos.BridgeInternal;
 import com.azure.cosmos.CosmosAsyncClient;
 import com.azure.cosmos.CosmosAsyncContainer;
+import com.azure.cosmos.TestObject;
 import com.azure.cosmos.implementation.HttpConstants;
 import com.azure.cosmos.models.CosmosItemResponse;
 import com.azure.cosmos.CosmosClientBuilder;
@@ -363,6 +364,19 @@ public void upsertDocument_ReplaceDocument(String documentId) throws Throwable {
         this.validateItemSuccess(readObservable, validator);
     }
 
+    @Test(groups = { "simple" }, timeOut = TIMEOUT)
+    public void upsertDocument_ReplaceDocumentWithPartitionKey() throws Throwable {
+        TestObject item = TestObject.create();
+        CosmosItemResponse<TestObject> response = container.createItem(item,  new PartitionKey(item.getMypk()), new CosmosItemRequestOptions()).block();
+
+        item.setStringProp( UUID.randomUUID().toString());
+
+        CosmosItemResponse<TestObject> replaceResponse = container.upsertItem(item,  new CosmosItemRequestOptions()).block();
+
+        // Validate result
+        assertThat(replaceResponse.getItem()).isEqualTo(item);
+    }
+
     @Test(groups = {"simple"}, timeOut = TIMEOUT)
     public void typedItems() throws Throwable {
         String docId = "1234";
@@ -384,7 +398,7 @@ public void typedItems() throws Throwable {
         TestObject resultObject = itemResponseMono.block().getItem();
         compareTestObjs(newTestObject, resultObject);
 
-        Mono<CosmosItemResponse<TestObject>> readResponseMono = container.readItem(newTestObject.id,
+        Mono<CosmosItemResponse<TestObject>> readResponseMono = container.readItem(newTestObject.getId(),
                                                                                         new PartitionKey(newTestObject
                                                                                                              .getMypk()),
                                                                                         TestObject.class);
@@ -407,65 +421,6 @@ private void compareTestObjs(TestObject newTestObject, TestObject resultObject)
         assertThat(newTestObject.getStringProp()).isEqualTo(resultObject.getStringProp());
     }
 
-    static class TestObject {
-        private String id;
-        private String mypk;
-        private List<List<Integer>> sgmts;
-        private String stringProp;
-
-        public TestObject() {
-        }
-
-        public TestObject(String id, String mypk, List<List<Integer>> sgmts, String stringProp) {
-            this.id = id;
-            this.mypk = mypk;
-            this.sgmts = sgmts;
-            this.stringProp = stringProp;
-        }
-
-        public String getId() {
-            return id;
-        }
-
-        public void setId(String id) {
-            this.id = id;
-        }
-
-        public String getMypk() {
-            return mypk;
-        }
-
-        public void setMypk(String mypk) {
-            this.mypk = mypk;
-        }
-
-        public List<List<Integer>> getSgmts() {
-            return sgmts;
-        }
-
-        public void setSgmts(List<List<Integer>> sgmts) {
-            this.sgmts = sgmts;
-        }
-
-        /**
-         * Getter for property 'stringProp'.
-         *
-         * @return Value for property 'stringProp'.
-         */
-        public String getStringProp() {
-            return stringProp;
-        }
-
-        /**
-         * Setter for property 'stringProp'.
-         *
-         * @param stringProp Value to set for property 'stringProp'.
-         */
-        public void setStringProp(String stringProp) {
-            this.stringProp = stringProp;
-        }
-    }
-
 
     @BeforeClass(groups = { "simple" }, timeOut = SETUP_TIMEOUT)
     public void before_DocumentCrudTest() {