Modify PutObject API to add delete functionality

Author: G8XSU

app/src/main/java/org/vss/impl/postgres/PostgresBackendImpl.java

@@ -66,23 +66,31 @@ public PutObjectResponse put(PutObjectRequest request) {
 
     String storeId = request.getStoreId();
 
-    List<VssDbRecord> vssRecords = new ArrayList<>(request.getTransactionItemsList().stream()
+    List<VssDbRecord> vssPutRecords = new ArrayList<>(request.getTransactionItemsList().stream()
+        .map(kv -> buildVssRecord(storeId, kv)).toList());
+
+    List<VssDbRecord> vssDeleteRecords = new ArrayList<>(request.getDeleteItemsList().stream()
         .map(kv -> buildVssRecord(storeId, kv)).toList());
 
     if (request.hasGlobalVersion()) {
       VssDbRecord globalVersionRecord = buildVssRecord(storeId,
           KeyValue.newBuilder()
               .setKey(GLOBAL_VERSION_KEY)
               .setVersion(request.getGlobalVersion())
+              .setValue(ByteString.EMPTY)
               .build());
 
-      vssRecords.add(globalVersionRecord);
+      vssPutRecords.add(globalVersionRecord);
     }
 
     context.transaction((ctx) -> {
       DSLContext dsl = ctx.dsl();
-      List<Query> batchQueries = vssRecords.stream()
-          .map(vssRecord -> buildPutObjectQuery(dsl, vssRecord)).toList();
+      List<Query> batchQueries = new ArrayList<>();
+
+      batchQueries.addAll(vssPutRecords.stream()
+          .map(vssRecord -> buildPutObjectQuery(dsl, vssRecord)).toList());
+      batchQueries.addAll(vssDeleteRecords.stream()
+          .map(vssRecord -> buildDeleteObjectQuery(dsl, vssRecord)).toList());
 
       int[] batchResult = dsl.batch(batchQueries).execute();
 
@@ -97,6 +105,12 @@ public PutObjectResponse put(PutObjectRequest request) {
     return PutObjectResponse.newBuilder().build();
   }
 
+  private Query buildDeleteObjectQuery(DSLContext dsl, VssDbRecord vssRecord) {
+    return dsl.deleteFrom(VSS_DB).where(VSS_DB.STORE_ID.eq(vssRecord.getStoreId())
+        .and(VSS_DB.KEY.eq(vssRecord.getKey()))
+        .and(VSS_DB.VERSION.eq(vssRecord.getVersion())));
+  }
+
   private Query buildPutObjectQuery(DSLContext dsl, VssDbRecord vssRecord) {
     return vssRecord.getVersion() == 0 ? buildInsertRecordQuery(dsl, vssRecord)
         : buildUpdateRecordQuery(dsl, vssRecord);

app/src/test/java/org/vss/AbstractKVStoreIntegrationTest.java

@@ -131,6 +131,40 @@ void putShouldSucceedWhenNoGlobalVersionIsGiven() {
     assertThat(getObject(KVStore.GLOBAL_VERSION_KEY).getVersion(), is(0L));
   }
 
+  @Test
+  void putAndDeleteShouldSucceedAsAtomicTransaction() {
+    assertDoesNotThrow(() -> putObjects(null, List.of(kv("k1", "k1v1", 0))));
+    // Put and Delete succeeds
+    assertDoesNotThrow(() -> putAndDeleteObjects(null, List.of(kv("k2", "k2v1", 0)), List.of(kv("k1", "", 1))));
+
+    KeyValue response = getObject("k2");
+    assertThat(response.getKey(), is("k2"));
+    assertThat(response.getVersion(), is(1L));
+    assertThat(response.getValue().toStringUtf8(), is("k2v1"));
+
+    assertTrue(getObject("k1").getValue().isEmpty());
+
+    // Delete fails (and hence put as well) due to mismatched version for the deleted item.
+    assertThrows(ConflictException.class, () -> putAndDeleteObjects(null, List.of(kv("k3", "k3v1", 0)), List.of(kv("k2", "", 3))));
+
+    assertTrue(getObject("k3").getValue().isEmpty());
+    assertFalse(getObject("k2").getValue().isEmpty());
+
+    // Put fails (and hence delete as well) due to mismatched version for the put item.
+    assertThrows(ConflictException.class, () -> putAndDeleteObjects(null, List.of(kv("k3", "k3v1", 1)), List.of(kv("k2", "", 1))));
+
+    assertTrue(getObject("k3").getValue().isEmpty());
+    assertFalse(getObject("k2").getValue().isEmpty());
+
+    // Put and delete both fail due to mismatched global version.
+    assertThrows(ConflictException.class, () -> putAndDeleteObjects(2L, List.of(kv("k3", "k3v1", 0)), List.of(kv("k2", "", 1))));
+
+    assertTrue(getObject("k3").getValue().isEmpty());
+    assertFalse(getObject("k2").getValue().isEmpty());
+
+    assertThat(getObject(KVStore.GLOBAL_VERSION_KEY).getVersion(), is(0L));
+  }
+
   @Test
   void getShouldReturnEmptyResponseWhenKeyDoesNotExist() {
     KeyValue response = getObject("non_existent_key");
@@ -370,6 +404,19 @@ private void putObjects(@Nullable Long globalVersion, List<KeyValue> keyValues)
     this.kvStore.put(putObjectRequestBuilder.build());
   }
 
+  private void putAndDeleteObjects(@Nullable Long globalVersion, List<KeyValue> putKeyValues, List<KeyValue> deleteKeyValues) {
+    PutObjectRequest.Builder putObjectRequestBuilder = PutObjectRequest.newBuilder()
+        .setStoreId(STORE_ID)
+        .addAllTransactionItems(putKeyValues)
+        .addAllDeleteItems(deleteKeyValues);
+
+    if (Objects.nonNull(globalVersion)) {
+      putObjectRequestBuilder.setGlobalVersion(globalVersion);
+    }
+
+    this.kvStore.put(putObjectRequestBuilder.build());
+  }
+
   private ListKeyVersionsResponse list(@Nullable String nextPageToken, @Nullable Integer pageSize,
       @Nullable String keyPrefix) {
     ListKeyVersionsRequest.Builder listRequestBuilder = ListKeyVersionsRequest.newBuilder()