From 2f886c6923b8d201bd207c3fa242a987bf5553c0 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Mon, 20 Jul 2026 14:18:42 +0530 Subject: [PATCH 1/4] Added DeepBulkUpdateResult to give fine-grained information to the callers regarding sucessful/failed keys --- .../core/documentstore/Collection.java | 5 +- .../documentstore/DeepBulkUpdateResult.java | 37 ++++++ .../postgres/FlatPostgresCollection.java | 38 ++++-- .../postgres/FlatPostgresCollectionTest.java | 122 ++++++++++++++++++ 4 files changed, 187 insertions(+), 15 deletions(-) create mode 100644 document-store/src/main/java/org/hypertrace/core/documentstore/DeepBulkUpdateResult.java diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/Collection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/Collection.java index a79e3c99..5f5ea678 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/Collection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/Collection.java @@ -438,10 +438,11 @@ CloseableIterator 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> updates, UpdateOptions updateOptions) throws IOException { throw new UnsupportedOperationException("bulkUpdate is not supported!"); diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/DeepBulkUpdateResult.java b/document-store/src/main/java/org/hypertrace/core/documentstore/DeepBulkUpdateResult.java new file mode 100644 index 00000000..8263d96d --- /dev/null +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/DeepBulkUpdateResult.java @@ -0,0 +1,37 @@ +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 failedKeys; + private final Set successfulKeys; + + public DeepBulkUpdateResult( + long successfullyUpdated, Set failedKeys, Set 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(); + } +} diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java index cbbcbfa3..75d6b9e7 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java @@ -45,6 +45,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; @@ -869,47 +870,54 @@ public CloseableIterator bulkUpdate( } @Override - public BulkUpdateResult bulkUpdate( + public DeepBulkUpdateResult bulkUpdate( Map> 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 failedKeys = new HashSet<>(); + Set successfulKeys = new HashSet<>(); try (Connection connection = client.getPooledConnection()) { // Group keys by their "SQL shape" (same SET-clause fragments AND param count) Map 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 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( @@ -963,7 +971,10 @@ private boolean executeKeyUpdate( * nested-JSONB REMOVE_ALL_FROM_LIST emitting 1+N placeholders) will land in distinct groups. */ private Map groupKeysByUpdateShape( - Connection connection, Map> updates, String tableName) { + Connection connection, + Map> updates, + String tableName, + Set failedKeys) { Map groups = new LinkedHashMap<>(); @@ -995,7 +1006,8 @@ private Map 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); } } @@ -1057,7 +1069,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: {}, Error: {}", sql, e.getMessage()); throw e; } } diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollectionTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollectionTest.java index af9e724a..62d4ef04 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollectionTest.java +++ b/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,123 @@ 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.getFailedKeys().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()); + } + + @Test + @DisplayName("Map bulkUpdate: null UpdateOptions throws IllegalArgumentException") + void testMapBulkUpdateThrowsOnNullOptions() { + Map> 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 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> updates = new LinkedHashMap<>(); + updates.put(Key.from("k1"), List.of(SubDocumentUpdate.of("price", 100))); + updates.put(Key.from("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.getFailedKeys().isEmpty()); + 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 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> 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()); + assertEquals( + Set.of(k3, k4), + result.getFailedKeys(), + "failedKeys must contain exactly the keys from the aborted batch group"); + 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> 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 From 4a23c52f0e5d35952d3a7da57424403c60d5007e Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Mon, 20 Jul 2026 14:21:39 +0530 Subject: [PATCH 2/4] WIP --- .../java/org/hypertrace/core/documentstore/Collection.java | 4 ++-- .../hypertrace/core/documentstore/DeepBulkUpdateResult.java | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/Collection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/Collection.java index 5f5ea678..a3cb7d8c 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/Collection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/Collection.java @@ -438,8 +438,8 @@ CloseableIterator 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 {@link DeepBulkUpdateResult} containing the count of successfully updated documents - * and the set of keys whose updates could not be applied. + * @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 DeepBulkUpdateResult bulkUpdate( diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/DeepBulkUpdateResult.java b/document-store/src/main/java/org/hypertrace/core/documentstore/DeepBulkUpdateResult.java index 8263d96d..64f37007 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/DeepBulkUpdateResult.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/DeepBulkUpdateResult.java @@ -14,6 +14,8 @@ public class DeepBulkUpdateResult extends BulkUpdateResult { private final Set failedKeys; private final Set successfulKeys; + // all keys attempted can be derived from failedKeys + successfulKeys + public DeepBulkUpdateResult( long successfullyUpdated, Set failedKeys, Set successfulKeys) { super(successfullyUpdated); From 705c00ecf70a61436908daf534f810e5412dcd60 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Mon, 20 Jul 2026 14:26:23 +0530 Subject: [PATCH 3/4] WIP --- .../core/documentstore/postgres/FlatPostgresCollection.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java index 75d6b9e7..cc90117b 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java @@ -1069,7 +1069,7 @@ private int executeBatchUpdate( LOGGER.debug("Batch update affected {} rows out of {} keys", totalUpdated, keys.size()); return totalUpdated; } catch (SQLException e) { - LOGGER.error("Failed to execute batch update. SQL: {}, Error: {}", sql, e.getMessage()); + LOGGER.error("Failed to execute batch update. SQL: {}", sql, e); throw e; } } From cb58b8a6311e586029483257f337d1a0ed1857bb Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Mon, 20 Jul 2026 14:33:02 +0530 Subject: [PATCH 4/4] Added UTs for successful keys --- .../postgres/FlatPostgresCollectionTest.java | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollectionTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollectionTest.java index 62d4ef04..b3cfa4de 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollectionTest.java +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollectionTest.java @@ -481,7 +481,9 @@ void testMapBulkUpdateEmptyMap() throws Exception { assertEquals(0, result.getUpdatedCount()); assertFalse(result.hasFailures()); + assertTrue(result.wasSuccessful()); assertTrue(result.getFailedKeys().isEmpty()); + assertTrue(result.getSuccessfulKeys().isEmpty()); } @Test @@ -492,6 +494,8 @@ void testMapBulkUpdateNullMap() throws Exception { assertEquals(0, result.getUpdatedCount()); assertFalse(result.hasFailures()); + assertTrue(result.wasSuccessful()); + assertTrue(result.getSuccessfulKeys().isEmpty()); } @Test @@ -514,15 +518,22 @@ void testMapBulkUpdateAllSuccess() throws Exception { when(mockPreparedStatement.executeBatch()).thenReturn(new int[] {1, 1}); Map> updates = new LinkedHashMap<>(); - updates.put(Key.from("k1"), List.of(SubDocumentUpdate.of("price", 100))); - updates.put(Key.from("k2"), List.of(SubDocumentUpdate.of("price", 200))); + 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(); } @@ -563,10 +574,15 @@ void testMapBulkUpdatePartialGroupFailure() throws Exception { 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(); }