[Bug fix] Check if document exists before put operation

Signed-off-by: Arjun kumar Giri <[email protected]>
arjunkumargiri · Jul 31, 2024 · ec237fa · ec237fa
1 parent b23935b
commit ec237fa
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 8 deletions.
diff --git a/plugin/src/main/java/org/opensearch/ml/sdkclient/DDBOpenSearchClient.java b/plugin/src/main/java/org/opensearch/ml/sdkclient/DDBOpenSearchClient.java
@@ -128,8 +128,8 @@ public CompletionStage<PutDataObjectResponse> putDataObjectAsync(PutDataObjectRe
                 item.put(SOURCE, AttributeValue.builder().m(sourceMap).build());
                 item.put(SEQ_NO_KEY, AttributeValue.builder().n(sequenceNumber.toString()).build());
                 Builder builder = PutItemRequest.builder().tableName(tableName).item(item);
-                if (!request.overwriteIfExists()) {
-                    builder.conditionExpression("attribute_not_exists(" + HASH_KEY + ") AND attribute_not_exists(" + RANGE_KEY + ")");
+                if (!request.overwriteIfExists() && getItemResponse != null && getItemResponse.item() != null) {
+                    throw new OpenSearchStatusException("Existing data object for ID: " + request.id(), RestStatus.CONFLICT);
                 }
                 final PutItemRequest putItemRequest = builder.build();
 

diff --git a/plugin/src/test/java/org/opensearch/ml/sdkclient/DDBOpenSearchClientTests.java b/plugin/src/test/java/org/opensearch/ml/sdkclient/DDBOpenSearchClientTests.java
@@ -160,11 +160,6 @@ public void testPutDataObject_HappyCase() throws IOException {
         Assert.assertEquals(TEST_ID, putItemRequest.item().get(RANGE_KEY).s());
         Assert.assertEquals(TENANT_ID, putItemRequest.item().get(HASH_KEY).s());
         Assert.assertEquals("0", putItemRequest.item().get(SEQ_NUM).n());
-        Assert
-            .assertEquals(
-                "attribute_not_exists(" + HASH_KEY + ") AND attribute_not_exists(" + RANGE_KEY + ")",
-                putItemRequest.conditionExpression()
-            );
         Assert.assertEquals("foo", putItemRequest.item().get("_source").m().get("data").s());
     }
 
@@ -174,7 +169,6 @@ public void testPutDataObject_ExistingDocument_UpdatesSequenceNumber() throws IO
             .builder()
             .index(TEST_INDEX)
             .id(TEST_ID)
-            .overwriteIfExists(false)
             .tenantId(TENANT_ID)
             .dataObject(testDataObject)
             .build();
@@ -193,6 +187,27 @@ public void testPutDataObject_ExistingDocument_UpdatesSequenceNumber() throws IO
         Assert.assertEquals("6", putItemRequest.item().get(SEQ_NUM).n());
     }
 
+    @Test
+    public void testPutDataObject_ExistingDocument_DisableOverwrite() throws IOException {
+        PutDataObjectRequest putRequest = PutDataObjectRequest
+            .builder()
+            .index(TEST_INDEX)
+            .id(TEST_ID)
+            .tenantId(TENANT_ID)
+            .overwriteIfExists(false)
+            .dataObject(testDataObject)
+            .build();
+        Mockito
+            .when(dynamoDbClient.getItem(Mockito.any(GetItemRequest.class)))
+            .thenReturn(GetItemResponse.builder().item(ImmutableMap.of(SEQ_NUM, AttributeValue.builder().n("5").build())).build());
+        Mockito.when(dynamoDbClient.putItem(Mockito.any(PutItemRequest.class))).thenReturn(PutItemResponse.builder().build());
+        CompletableFuture<PutDataObjectResponse> response = sdkClient
+            .putDataObjectAsync(putRequest, testThreadPool.executor(GENERAL_THREAD_POOL))
+            .toCompletableFuture();
+        CompletionException ce = assertThrows(CompletionException.class, () -> response.join());
+        assertEquals(OpenSearchStatusException.class, ce.getCause().getClass());
+    }
+
     @Test
     public void testPutDataObject_WithComplexData() throws IOException {
         ComplexDataObject complexDataObject = ComplexDataObject