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..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,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..64f37007 --- /dev/null +++ b/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 failedKeys; + private final Set successfulKeys; + + // all keys attempted can be derived from failedKeys + 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 2c8f263c..19927abb 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 @@ -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 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( @@ -969,7 +977,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<>(); @@ -1001,7 +1012,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); } } @@ -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; } } 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..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 @@ -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> 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<>(); + 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 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()); + 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> 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