feat: bulk delete documents (#77)

Author: skjindal93

document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreTest.java

@@ -19,6 +19,7 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
 import org.apache.commons.lang3.RandomUtils;
@@ -981,6 +982,23 @@ public void testDelete(String dataStoreName) throws IOException {
     Assertions.assertEquals(collection.count(), 0);
   }
 
+  @ParameterizedTest
+  @MethodSource("databaseContextProvider")
+  public void testBulkDelete(String dataStoreName) throws IOException {
+    Datastore datastore = datastoreMap.get(dataStoreName);
+    Collection collection = datastore.getCollection(COLLECTION_NAME);
+    SingleValueKey docKey1 = new SingleValueKey("default", "testKey1");
+    collection.upsert(docKey1, Utils.createDocument("foo1", "bar1"));
+    SingleValueKey docKey2 = new SingleValueKey("default", "testKey2");
+    collection.upsert(docKey2, Utils.createDocument("foo2", "bar2"));
+    SingleValueKey docKey3 = new SingleValueKey("default", "testKey3");
+    collection.upsert(docKey3, Utils.createDocument("foo3", "bar3"));
+
+    Assertions.assertEquals(collection.count(), 3);
+    collection.delete(Set.of(docKey1, docKey2));
+    Assertions.assertEquals(collection.count(), 1);
+  }
+
   @ParameterizedTest
   @MethodSource("databaseContextProvider")
   public void testDeleteAll(String dataStoreName) throws IOException {

document-store/src/main/java/org/hypertrace/core/documentstore/BulkDeleteResult.java

@@ -0,0 +1,8 @@
+package org.hypertrace.core.documentstore;
+
+public class BulkDeleteResult extends DeleteResult {
+
+  public BulkDeleteResult(long deletedCount) {
+    super(deletedCount);
+  }
+}

document-store/src/main/java/org/hypertrace/core/documentstore/Collection.java

@@ -3,6 +3,7 @@
 import java.io.IOException;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 /** Interface spec for common operations on a collection of documents */
 public interface Collection {
@@ -87,6 +88,13 @@ public interface Collection {
    */
   boolean delete(Key key);
 
+  /**
+   * Delete the documents for the given keys
+   *
+   * @param keys {@link Key}s of the document to be deleted
+   */
+  BulkDeleteResult delete(Set<Key> keys);
+
   /**
    * Deletes a sub document
    *

document-store/src/main/java/org/hypertrace/core/documentstore/DeleteResult.java

@@ -0,0 +1,13 @@
+package org.hypertrace.core.documentstore;
+
+public class DeleteResult {
+  private long deletedCount;
+
+  public DeleteResult(long deletedCount) {
+    this.deletedCount = deletedCount;
+  }
+
+  public long getDeletedCount() {
+    return deletedCount;
+  }
+}

document-store/src/main/java/org/hypertrace/core/documentstore/mongo/MongoCollection.java

@@ -40,6 +40,7 @@
 import org.bson.json.JsonMode;
 import org.bson.json.JsonWriterSettings;
 import org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest;
+import org.hypertrace.core.documentstore.BulkDeleteResult;
 import org.hypertrace.core.documentstore.BulkUpdateRequest;
 import org.hypertrace.core.documentstore.BulkUpdateResult;
 import org.hypertrace.core.documentstore.CloseableIterator;
@@ -479,6 +480,12 @@ public boolean delete(Key key) {
     return deleteResult.getDeletedCount() > 0;
   }
 
+  @Override
+  public BulkDeleteResult delete(Set<Key> keys) {
+    DeleteResult deleteResult = collection.deleteMany(this.selectionCriteriaForKeys(keys));
+    return new BulkDeleteResult(deleteResult.getDeletedCount());
+  }
+
   @Override
   public boolean deleteSubDoc(Key key, String subDocPath) {
     BasicDBObject unsetObject = new BasicDBObject("$unset", new BasicDBObject(subDocPath, ""));

document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java

@@ -21,6 +21,7 @@
 import java.util.stream.Collectors;
 import lombok.SneakyThrows;
 import org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest;
+import org.hypertrace.core.documentstore.BulkDeleteResult;
 import org.hypertrace.core.documentstore.BulkUpdateRequest;
 import org.hypertrace.core.documentstore.BulkUpdateResult;
 import org.hypertrace.core.documentstore.CloseableIterator;
@@ -344,6 +345,31 @@ public boolean delete(Key key) {
     return false;
   }
 
+  @Override
+  public BulkDeleteResult delete(Set<Key> keys) {
+    String ids =
+        keys.stream().map(key -> "'" + key.toString() + "'").collect(Collectors.joining(", "));
+    String space = " ";
+    String deleteSQL =
+        new StringBuilder("DELETE FROM")
+            .append(space)
+            .append(collectionName)
+            .append(" WHERE ")
+            .append(ID)
+            .append(" IN ")
+            .append("(")
+            .append(ids)
+            .append(")")
+            .toString();
+    try (PreparedStatement preparedStatement = client.prepareStatement(deleteSQL)) {
+      int deletedCount = preparedStatement.executeUpdate();
+      return new BulkDeleteResult(deletedCount);
+    } catch (SQLException e) {
+      LOGGER.error("SQLException deleting documents. keys: {}", keys, e);
+    }
+    return new BulkDeleteResult(0);
+  }
+
   @Override
   public boolean deleteSubDoc(Key key, String subDocPath) {
     String deleteSubDocSQL =