Add DeepBulkUpdateResult to give fine-grained information to the caller about failed/successful keys in batch updates

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

@@ -438,10 +438,11 @@ CloseableIterator<Document> bulkUpdate(
    * @param updates Map of Key to Collection of SubDocumentUpdate operations. Each key's updates are
    *     applied atomically, but no cross-key atomicity is guaranteed.
    * @param updateOptions Options for the update operation
-   * @return BulkUpdateResult containing the count of successfully updated documents
+   * @return {@link DeepBulkUpdateResult} containing the count of successfully updated documents and
+   *     the set of keys whose updates could not be applied.
    * @throws IOException if the update operation fails
    */
-  default BulkUpdateResult bulkUpdate(
+  default DeepBulkUpdateResult bulkUpdate(
       Map<Key, java.util.Collection<SubDocumentUpdate>> updates, UpdateOptions updateOptions)
       throws IOException {
     throw new UnsupportedOperationException("bulkUpdate is not supported!");

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

@@ -0,0 +1,39 @@
+package org.hypertrace.core.documentstore;
+
+import java.util.Collections;
+import java.util.Set;
+import lombok.Getter;
+
+/**
+ * Carries information about failed and successful keys. Can be enhanced in future to carry
+ * information like failure reasons, etc.
+ */
+@Getter
+public class DeepBulkUpdateResult extends BulkUpdateResult {
+
+  private final Set<Key> failedKeys;
+  private final Set<Key> successfulKeys;
+
+  // all keys attempted can be derived from failedKeys + successfulKeys
+
+  public DeepBulkUpdateResult(
+      long successfullyUpdated, Set<Key> failedKeys, Set<Key> successfulKeys) {
+    super(successfullyUpdated);
+    this.failedKeys =
+        failedKeys == null || failedKeys.isEmpty()
+            ? Collections.emptySet()
+            : Collections.unmodifiableSet(failedKeys);
+    this.successfulKeys =
+        successfulKeys == null || successfulKeys.isEmpty()
+            ? Collections.emptySet()
+            : Collections.unmodifiableSet(successfulKeys);
+  }
+
+  public boolean hasFailures() {
+    return !failedKeys.isEmpty();
+  }
+
+  public boolean wasSuccessful() {
+    return failedKeys.isEmpty();
+  }
+}

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

@@ -47,6 +47,7 @@
 import org.hypertrace.core.documentstore.CloseableIterator;
 import org.hypertrace.core.documentstore.CreateResult;
 import org.hypertrace.core.documentstore.CreateStatus;
+import org.hypertrace.core.documentstore.DeepBulkUpdateResult;
 import org.hypertrace.core.documentstore.Document;
 import org.hypertrace.core.documentstore.DocumentType;
 import org.hypertrace.core.documentstore.Filter;
@@ -875,47 +876,54 @@ public CloseableIterator<Document> bulkUpdate(
   }
 
   @Override
-  public BulkUpdateResult bulkUpdate(
+  public DeepBulkUpdateResult bulkUpdate(
       Map<Key, Collection<SubDocumentUpdate>> updates, UpdateOptions updateOptions)
       throws IOException {
 
     if (updates == null || updates.isEmpty()) {
-      return new BulkUpdateResult(0);
+      return new DeepBulkUpdateResult(0, Set.of(), Set.of());
     }
 
     Preconditions.checkArgument(updateOptions != null, "UpdateOptions cannot be NULL");
 
+    // todo: Honour update options like in bulkUpdate
     String tableName = tableIdentifier.getTableName();
     String quotedPkColumn = PostgresUtils.wrapFieldNamesWithDoubleQuotes(getPKForTable(tableName));
     long batchUpdateTimestamp = System.currentTimeMillis();
 
     int totalUpdated = 0;
+    Set<Key> failedKeys = new HashSet<>();
+    Set<Key> successfulKeys = new HashSet<>();
 
     try (Connection connection = client.getPooledConnection()) {
       // Group keys by their "SQL shape" (same SET-clause fragments AND param count)
       Map<String, KeyUpdateGroup> keyGroups =
-          groupKeysByUpdateShape(connection, updates, tableName);
+          groupKeysByUpdateShape(connection, updates, tableName, failedKeys);
 
       // Execute one multi-row UPDATE per group (or fallback to single-key if group size = 1)
       for (Map.Entry<String, KeyUpdateGroup> entry : keyGroups.entrySet()) {
+        KeyUpdateGroup group = entry.getValue();
         try {
           int updated =
               executeBatchUpdate(
-                  connection, entry.getValue(), tableName, quotedPkColumn, batchUpdateTimestamp);
+                  connection, group, tableName, quotedPkColumn, batchUpdateTimestamp);
           totalUpdated += updated;
+          successfulKeys.addAll(group.getKeys());
         } catch (Exception e) {
-          LOGGER.warn(
-              "Failed to update key group (size: {}): {}",
-              entry.getValue().getKeys().size(),
-              e.getMessage());
-          // Continue with other groups - no cross-group atomicity
+          LOGGER.error(
+              "Failed to update key group (size: {}). Will continue with remaining batches",
+              group.getKeys().size(),
+              e);
+          // Batch was aborted as a whole (e.g. Postgres deadlock); surface every key in the
+          // group to the caller so it can decide what to do. No cross-group atomicity either.
+          failedKeys.addAll(group.getKeys());
         }
       }
     } catch (SQLException e) {
       throw new IOException("Failed to get connection for bulk update", e);
     }
 
-    return new BulkUpdateResult(totalUpdated);
+    return new DeepBulkUpdateResult(totalUpdated, failedKeys, successfulKeys);
   }
 
   private boolean executeKeyUpdate(
@@ -969,7 +977,10 @@ private boolean executeKeyUpdate(
    * nested-JSONB REMOVE_ALL_FROM_LIST emitting 1+N placeholders) will land in distinct groups.
    */
   private Map<String, KeyUpdateGroup> groupKeysByUpdateShape(
-      Connection connection, Map<Key, Collection<SubDocumentUpdate>> updates, String tableName) {
+      Connection connection,
+      Map<Key, Collection<SubDocumentUpdate>> updates,
+      String tableName,
+      Set<Key> failedKeys) {
 
     Map<String, KeyUpdateGroup> groups = new LinkedHashMap<>();
 
@@ -1001,7 +1012,8 @@ private Map<String, KeyUpdateGroup> groupKeysByUpdateShape(
             .addKeyWithParams(key, params);
 
       } catch (Exception e) {
-        LOGGER.warn("Failed to group key {}: {}", key, e.getMessage());
+        LOGGER.warn("Failed to group key {}", key, e);
+        failedKeys.add(key);
       }
     }
 
@@ -1070,7 +1082,7 @@ private int executeBatchUpdate(
       LOGGER.debug("Batch update affected {} rows out of {} keys", totalUpdated, keys.size());
       return totalUpdated;
     } catch (SQLException e) {
-      LOGGER.warn("Failed to execute batch update. SQL: {}, Error: {}", sql, e.getMessage());
+      LOGGER.error("Failed to execute batch update. SQL: {}", sql, e);
       throw e;
     }
   }

document-store/src/test/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollectionTest.java

@@ -17,19 +17,24 @@
 import static org.mockito.Mockito.when;
 
 import java.io.IOException;
+import java.sql.BatchUpdateException;
 import java.sql.Connection;
 import java.sql.Date;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.Statement;
 import java.sql.Timestamp;
 import java.time.Instant;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import org.hypertrace.core.documentstore.CloseableIterator;
+import org.hypertrace.core.documentstore.DeepBulkUpdateResult;
 import org.hypertrace.core.documentstore.Document;
 import org.hypertrace.core.documentstore.JSONDocument;
 import org.hypertrace.core.documentstore.Key;
@@ -466,6 +471,139 @@ void testBulkUpdateClosesIteratorOnException() throws Exception {
 
       verify(mockIterator).close();
     }
+
+    @Test
+    @DisplayName("Map bulkUpdate: empty updates map -> zero updated, no failures")
+    void testMapBulkUpdateEmptyMap() throws Exception {
+      DeepBulkUpdateResult result =
+          flatPostgresCollection.bulkUpdate(
+              Collections.emptyMap(), UpdateOptions.DEFAULT_UPDATE_OPTIONS);
+
+      assertEquals(0, result.getUpdatedCount());
+      assertFalse(result.hasFailures());
+      assertTrue(result.wasSuccessful());
+      assertTrue(result.getFailedKeys().isEmpty());
+      assertTrue(result.getSuccessfulKeys().isEmpty());
+    }
+
+    @Test
+    @DisplayName("Map bulkUpdate: null updates map -> zero updated, no failures")
+    void testMapBulkUpdateNullMap() throws Exception {
+      DeepBulkUpdateResult result =
+          flatPostgresCollection.bulkUpdate(null, UpdateOptions.DEFAULT_UPDATE_OPTIONS);
+
+      assertEquals(0, result.getUpdatedCount());
+      assertFalse(result.hasFailures());
+      assertTrue(result.wasSuccessful());
+      assertTrue(result.getSuccessfulKeys().isEmpty());
+    }
+
+    @Test
+    @DisplayName("Map bulkUpdate: null UpdateOptions throws IllegalArgumentException")
+    void testMapBulkUpdateThrowsOnNullOptions() {
+      Map<Key, Collection<SubDocumentUpdate>> updates = new LinkedHashMap<>();
+      updates.put(Key.from("k1"), List.of(SubDocumentUpdate.of("price", 100)));
+
+      assertThrows(
+          IllegalArgumentException.class, () -> flatPostgresCollection.bulkUpdate(updates, null));
+    }
+
+    @Test
+    @DisplayName("Map bulkUpdate: happy path - all keys updated, no failures")
+    void testMapBulkUpdateAllSuccess() throws Exception {
+      Map<String, PostgresColumnMetadata> schema = createBasicSchema();
+      setupCommonMocks(schema);
+
+      // Every SET on "price" lands in one shape group; two keys => two batch entries.
+      when(mockPreparedStatement.executeBatch()).thenReturn(new int[] {1, 1});
+
+      Map<Key, Collection<SubDocumentUpdate>> updates = new LinkedHashMap<>();
+      Key k1 = Key.from("k1");
+      Key k2 = Key.from("k2");
+      updates.put(k1, List.of(SubDocumentUpdate.of("price", 100)));
+      updates.put(k2, List.of(SubDocumentUpdate.of("price", 200)));
+
+      DeepBulkUpdateResult result =
+          flatPostgresCollection.bulkUpdate(updates, UpdateOptions.DEFAULT_UPDATE_OPTIONS);
+
+      assertEquals(2, result.getUpdatedCount());
+      assertFalse(result.hasFailures());
+      assertTrue(result.wasSuccessful());
+      assertTrue(result.getFailedKeys().isEmpty());
+      assertEquals(
+          Set.of(k1, k2),
+          result.getSuccessfulKeys(),
+          "successfulKeys must contain every input key when no batch aborts");
+      verify(mockPreparedStatement, times(1)).executeBatch();
+    }
+
+    @Test
+    @DisplayName(
+        "Map bulkUpdate: one group's executeBatch throws deadlock - only its keys marked failed")
+    void testMapBulkUpdatePartialGroupFailure() throws Exception {
+      Map<String, PostgresColumnMetadata> schema = createBasicSchema();
+      setupCommonMocks(schema);
+
+      // Two distinct SET columns => two shape groups. Group ordering follows input insertion
+      // order via LinkedHashMap; both groupings happen before either executeBatch call.
+      // First executeBatch (group for "price") succeeds; second (group for "item") aborts with
+      // Postgres deadlock_detected (40P01) via BatchUpdateException.
+      BatchUpdateException deadlock =
+          new BatchUpdateException(
+              "Batch entry ... was aborted: ERROR: deadlock detected",
+              "40P01",
+              0,
+              new int[] {Statement.EXECUTE_FAILED},
+              new SQLException("deadlock detected", "40P01"));
+      when(mockPreparedStatement.executeBatch())
+          .thenReturn(new int[] {1, 1}) // group A: SET "price" for k1, k2
+          .thenThrow(deadlock); // group B: SET "item" for k3, k4
+
+      Map<Key, Collection<SubDocumentUpdate>> updates = new LinkedHashMap<>();
+      Key k1 = Key.from("k1");
+      Key k2 = Key.from("k2");
+      Key k3 = Key.from("k3");
+      Key k4 = Key.from("k4");
+      updates.put(k1, List.of(SubDocumentUpdate.of("price", 100)));
+      updates.put(k2, List.of(SubDocumentUpdate.of("price", 200)));
+      updates.put(k3, List.of(SubDocumentUpdate.of("item", "Alpha")));
+      updates.put(k4, List.of(SubDocumentUpdate.of("item", "Bravo")));
+
+      DeepBulkUpdateResult result =
+          flatPostgresCollection.bulkUpdate(updates, UpdateOptions.DEFAULT_UPDATE_OPTIONS);
+
+      assertEquals(2, result.getUpdatedCount(), "only price group's rows should count as updated");
+      assertTrue(result.hasFailures());
+      assertFalse(result.wasSuccessful());
+      assertEquals(
+          Set.of(k3, k4),
+          result.getFailedKeys(),
+          "failedKeys must contain exactly the keys from the aborted batch group");
+      assertEquals(
+          Set.of(k1, k2),
+          result.getSuccessfulKeys(),
+          "successfulKeys must contain exactly the keys from the group whose batch completed");
+      verify(mockPreparedStatement, times(2)).executeBatch();
+    }
+
+    @Test
+    @DisplayName("Map bulkUpdate: connection acquisition failure -> IOException")
+    void testMapBulkUpdateConnectionAcquisitionFails() throws Exception {
+      // Only stubs actually consumed before the connection acquisition are required.
+      when(mockSchemaRegistry.getPrimaryKeyColumn(COLLECTION_NAME)).thenReturn(Optional.of("id"));
+      SQLException sqlException = new SQLException("pool exhausted", "08001");
+      when(mockClient.getPooledConnection()).thenThrow(sqlException);
+
+      Map<Key, Collection<SubDocumentUpdate>> updates = new LinkedHashMap<>();
+      updates.put(Key.from("k1"), List.of(SubDocumentUpdate.of("price", 100)));
+
+      IOException thrown =
+          assertThrows(
+              IOException.class,
+              () ->
+                  flatPostgresCollection.bulkUpdate(updates, UpdateOptions.DEFAULT_UPDATE_OPTIONS));
+      assertEquals(sqlException, thrown.getCause());
+    }
   }
 
   @Nested