From 2b8507a47b20c66055aa4e874a2a72b0b981ea8b Mon Sep 17 00:00:00 2001 From: Bhuvan506 Date: Mon, 24 Aug 2026 16:56:43 +0530 Subject: [PATCH 1/3] Return false from bulkUpsert when write result is incomplete Mongo and Postgres bulkUpsert previously returned true whenever no exception was thrown, ignoring BulkWriteResult / JDBC batch counts. Callers (e.g. attribute-service) then treated partial writes as success. Validate that every requested document is accounted for before reporting success; keep throwing/returning false on hard failures. Co-authored-by: Cursor --- .../core/documentstore/Collection.java | 3 +- .../commons/BatchWriteUtils.java | 27 ++++++ .../documentstore/mongo/MongoCollection.java | 38 +++++++++ .../postgres/FlatPostgresCollection.java | 9 ++ .../postgres/PostgresCollection.java | 16 ++++ .../commons/BatchWriteUtilsTest.java | 37 ++++++++ .../mongo/MongoCollectionTest.java | 84 +++++++++++++++++++ 7 files changed, 213 insertions(+), 1 deletion(-) create mode 100644 document-store/src/main/java/org/hypertrace/core/documentstore/commons/BatchWriteUtils.java create mode 100644 document-store/src/test/java/org/hypertrace/core/documentstore/commons/BatchWriteUtilsTest.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 a3cb7d8ca..1af44fb70 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 @@ -217,7 +217,8 @@ long count( /** * @param documents to be upserted in bulk - * @return true if the operation succeeded + * @return true if every requested document was upserted successfully; false if the write failed + * or only a subset of documents were written */ boolean bulkUpsert(Map documents); diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/commons/BatchWriteUtils.java b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/BatchWriteUtils.java new file mode 100644 index 000000000..9f9e3546e --- /dev/null +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/BatchWriteUtils.java @@ -0,0 +1,27 @@ +package org.hypertrace.core.documentstore.commons; + +import java.sql.Statement; + +/** Shared helpers for validating JDBC batch write outcomes. */ +public final class BatchWriteUtils { + + private BatchWriteUtils() {} + + /** + * Returns true when every batch entry completed without {@link Statement#EXECUTE_FAILED} and the + * result length matches the number of operations submitted. + * + *

{@link Statement#SUCCESS_NO_INFO} (-2) and positive update counts are treated as success. + */ + public static boolean isBatchFullySuccessful(final int[] updateCounts, final int expectedSize) { + if (updateCounts == null || updateCounts.length != expectedSize) { + return false; + } + for (final int count : updateCounts) { + if (count == Statement.EXECUTE_FAILED) { + return false; + } + } + return true; + } +} diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/MongoCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/MongoCollection.java index f80ef9e0a..29a91e685 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/MongoCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/MongoCollection.java @@ -12,6 +12,7 @@ import static org.hypertrace.core.documentstore.mongo.update.parser.MongoSetOperationParser.SET_CLAUSE; import com.fasterxml.jackson.core.JsonProcessingException; +import com.google.common.annotations.VisibleForTesting; import com.mongodb.BasicDBObject; import com.mongodb.MongoBulkWriteException; import com.mongodb.MongoCommandException; @@ -678,6 +679,17 @@ public boolean bulkUpsert(Map documents) { try { BulkWriteResult result = bulkUpsertImpl(documents); LOGGER.debug(result.toString()); + if (!isBulkUpsertComplete(result, documents.size())) { + LOGGER.error( + "Incomplete bulk upsert for documents. requested={}, matched={}, upserted={}," + + " acknowledged={}, result={}", + documents.size(), + result.wasAcknowledged() ? result.getMatchedCount() : -1, + result.wasAcknowledged() ? result.getUpserts().size() : -1, + result.wasAcknowledged(), + result); + return false; + } return true; } catch (IOException | MongoServerException e) { LOGGER.error("Error during bulk upsert for documents:{}", documents, e); @@ -702,6 +714,21 @@ private BulkWriteResult bulkUpsertImpl(Map documents) .get(() -> collection.bulkWrite(bulkCollection, new BulkWriteOptions().ordered(false))); } + /** + * Each UpdateOne upsert accounts for exactly one matched existing document or one upserted + * document. Incomplete results (or unacknowledged writes) must not be reported as success. + */ + @VisibleForTesting + static boolean isBulkUpsertComplete(final BulkWriteResult result, final int requestedCount) { + if (requestedCount == 0) { + return true; + } + if (!result.wasAcknowledged()) { + return false; + } + return result.getMatchedCount() + result.getUpserts().size() == requestedCount; + } + @Override public CloseableIterator bulkUpsertAndReturnOlderDocuments(Map documents) throws IOException { @@ -714,6 +741,17 @@ public CloseableIterator bulkUpsertAndReturnOlderDocuments(Map documents) { if (LOGGER.isDebugEnabled()) { LOGGER.debug("Bulk upsert results: {}", Arrays.toString(results)); } + if (!BatchWriteUtils.isBatchFullySuccessful(results, parsedDocuments.size())) { + LOGGER.error( + "Incomplete bulkUpsert. requested={}, submitted={}, updateCounts={}", + documents.size(), + parsedDocuments.size(), + Arrays.toString(results)); + return false; + } return true; } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java index 686228aab..6bfc5b2ce 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java @@ -67,6 +67,7 @@ import org.hypertrace.core.documentstore.Key; import org.hypertrace.core.documentstore.Query; import org.hypertrace.core.documentstore.UpdateResult; +import org.hypertrace.core.documentstore.commons.BatchWriteUtils; import org.hypertrace.core.documentstore.commons.CommonUpdateValidator; import org.hypertrace.core.documentstore.commons.DocStoreConstants; import org.hypertrace.core.documentstore.commons.UpdateValidator; @@ -754,6 +755,14 @@ public boolean bulkUpsert(Map documents) { LOGGER.debug("Write result: {}", Arrays.toString(updateCounts)); } + if (!BatchWriteUtils.isBatchFullySuccessful(updateCounts, documents.size())) { + LOGGER.error( + "Incomplete bulk upsert for documents. requested={}, updateCounts={}", + documents.size(), + Arrays.toString(updateCounts)); + return false; + } + return true; } catch (BatchUpdateException e) { LOGGER.error("BatchUpdateException bulk inserting documents.", e); @@ -801,6 +810,13 @@ public CloseableIterator bulkUpsertAndReturnOlderDocuments(Map documents = + Map.of( + new SingleValueKey("default", "k1"), document, + new SingleValueKey("default", "k2"), document); + + when(collection.bulkWrite(anyList(), any(BulkWriteOptions.class))) + .thenReturn(BulkWriteResult.acknowledged(0, 2, 0, 2, emptyList(), emptyList())); + + assertTrue(mongoCollection.bulkUpsert(documents)); + } + + @Test + void returnsTrueWhenDocumentsAreUpserted() throws Exception { + Document document = new JSONDocument("{\"planet\": \"Mars\"}"); + Map documents = + Map.of( + new SingleValueKey("default", "k1"), document, + new SingleValueKey("default", "k2"), document); + + List upserts = + List.of( + new BulkWriteUpsert(0, new BsonString("default:k1")), + new BulkWriteUpsert(1, new BsonString("default:k2"))); + when(collection.bulkWrite(anyList(), any(BulkWriteOptions.class))) + .thenReturn(BulkWriteResult.acknowledged(0, 0, 0, 0, upserts, emptyList())); + + assertTrue(mongoCollection.bulkUpsert(documents)); + } + + @Test + void returnsFalseWhenResultAccountsForFewerDocumentsThanRequested() throws Exception { + Document document = new JSONDocument("{\"planet\": \"Mars\"}"); + Map documents = + Map.of( + new SingleValueKey("default", "k1"), document, + new SingleValueKey("default", "k2"), document); + + // Only one of two requested docs accounted for (matched=1, upserts=0) + when(collection.bulkWrite(anyList(), any(BulkWriteOptions.class))) + .thenReturn(BulkWriteResult.acknowledged(0, 1, 0, 1, emptyList(), emptyList())); + + assertFalse(mongoCollection.bulkUpsert(documents)); + } + + @Test + void returnsFalseForUnacknowledgedResult() throws Exception { + Document document = new JSONDocument("{\"planet\": \"Mars\"}"); + Map documents = Map.of(new SingleValueKey("default", "k1"), document); + + when(collection.bulkWrite(anyList(), any(BulkWriteOptions.class))) + .thenReturn(BulkWriteResult.unacknowledged()); + + assertFalse(mongoCollection.bulkUpsert(documents)); + } + + @Test + void isBulkUpsertComplete_matchedPlusUpsertsEqualsRequested() { + assertTrue( + MongoCollection.isBulkUpsertComplete( + BulkWriteResult.acknowledged(0, 1, 0, 1, emptyList(), emptyList()), 1)); + assertTrue( + MongoCollection.isBulkUpsertComplete( + BulkWriteResult.acknowledged( + 0, 1, 0, 1, List.of(new BulkWriteUpsert(1, new BsonString("id"))), emptyList()), + 2)); + assertFalse( + MongoCollection.isBulkUpsertComplete( + BulkWriteResult.acknowledged(0, 1, 0, 1, emptyList(), emptyList()), 2)); + assertFalse(MongoCollection.isBulkUpsertComplete(BulkWriteResult.unacknowledged(), 1)); + assertTrue(MongoCollection.isBulkUpsertComplete(BulkWriteResult.unacknowledged(), 0)); + } + } } From a1eae2fc08fcba4b1defd23dd57ce234e2b2d309 Mon Sep 17 00:00:00 2001 From: Bhuvan506 Date: Thu, 27 Aug 2026 07:21:46 +0530 Subject: [PATCH 2/3] Move JDBC batch success check onto PostgresCollection Address review feedback: keep the helper with the Postgres hierarchy (FlatPostgresCollection already extends it) instead of a commons util. Co-authored-by: Cursor --- .../commons/BatchWriteUtils.java | 27 -------------- .../postgres/FlatPostgresCollection.java | 3 +- .../postgres/PostgresCollection.java | 24 ++++++++++-- .../commons/BatchWriteUtilsTest.java | 37 ------------------- .../postgres/PostgresCollectionTest.java | 34 +++++++++++++++++ 5 files changed, 56 insertions(+), 69 deletions(-) delete mode 100644 document-store/src/main/java/org/hypertrace/core/documentstore/commons/BatchWriteUtils.java delete mode 100644 document-store/src/test/java/org/hypertrace/core/documentstore/commons/BatchWriteUtilsTest.java diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/commons/BatchWriteUtils.java b/document-store/src/main/java/org/hypertrace/core/documentstore/commons/BatchWriteUtils.java deleted file mode 100644 index 9f9e3546e..000000000 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/commons/BatchWriteUtils.java +++ /dev/null @@ -1,27 +0,0 @@ -package org.hypertrace.core.documentstore.commons; - -import java.sql.Statement; - -/** Shared helpers for validating JDBC batch write outcomes. */ -public final class BatchWriteUtils { - - private BatchWriteUtils() {} - - /** - * Returns true when every batch entry completed without {@link Statement#EXECUTE_FAILED} and the - * result length matches the number of operations submitted. - * - *

{@link Statement#SUCCESS_NO_INFO} (-2) and positive update counts are treated as success. - */ - public static boolean isBatchFullySuccessful(final int[] updateCounts, final int expectedSize) { - if (updateCounts == null || updateCounts.length != expectedSize) { - return false; - } - for (final int count : updateCounts) { - if (count == Statement.EXECUTE_FAILED) { - return false; - } - } - return true; - } -} 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 30ba93a8a..1aab5c0ba 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 @@ -53,7 +53,6 @@ import org.hypertrace.core.documentstore.Filter; import org.hypertrace.core.documentstore.Key; import org.hypertrace.core.documentstore.UpdateResult; -import org.hypertrace.core.documentstore.commons.BatchWriteUtils; import org.hypertrace.core.documentstore.commons.CommonUpdateValidator; import org.hypertrace.core.documentstore.commons.UpdateValidator; import org.hypertrace.core.documentstore.model.config.postgres.CollectionConfig; @@ -431,7 +430,7 @@ public boolean bulkUpsert(Map documents) { if (LOGGER.isDebugEnabled()) { LOGGER.debug("Bulk upsert results: {}", Arrays.toString(results)); } - if (!BatchWriteUtils.isBatchFullySuccessful(results, parsedDocuments.size())) { + if (!isBatchFullySuccessful(results, parsedDocuments.size())) { LOGGER.error( "Incomplete bulkUpsert. requested={}, submitted={}, updateCounts={}", documents.size(), diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java index 6bfc5b2ce..2719b6a40 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java @@ -67,7 +67,6 @@ import org.hypertrace.core.documentstore.Key; import org.hypertrace.core.documentstore.Query; import org.hypertrace.core.documentstore.UpdateResult; -import org.hypertrace.core.documentstore.commons.BatchWriteUtils; import org.hypertrace.core.documentstore.commons.CommonUpdateValidator; import org.hypertrace.core.documentstore.commons.DocStoreConstants; import org.hypertrace.core.documentstore.commons.UpdateValidator; @@ -755,7 +754,7 @@ public boolean bulkUpsert(Map documents) { LOGGER.debug("Write result: {}", Arrays.toString(updateCounts)); } - if (!BatchWriteUtils.isBatchFullySuccessful(updateCounts, documents.size())) { + if (!isBatchFullySuccessful(updateCounts, documents.size())) { LOGGER.error( "Incomplete bulk upsert for documents. requested={}, updateCounts={}", documents.size(), @@ -810,7 +809,7 @@ public CloseableIterator bulkUpsertAndReturnOlderDocuments(Map documents) throws SQLException, } } + /** + * Returns true when every batch entry completed without {@link Statement#EXECUTE_FAILED} and the + * result length matches the number of operations submitted. + * + *

{@link Statement#SUCCESS_NO_INFO} (-2) and positive update counts are treated as success. + */ + @VisibleForTesting + static boolean isBatchFullySuccessful(final int[] updateCounts, final int expectedSize) { + if (updateCounts == null || updateCounts.length != expectedSize) { + return false; + } + for (final int count : updateCounts) { + if (count == Statement.EXECUTE_FAILED) { + return false; + } + } + return true; + } + @VisibleForTesting JsonNode getJsonNodeAtPath(String path, JsonNode rootNode, boolean createPathIfMissing) { if (StringUtils.isEmpty(path)) { diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/commons/BatchWriteUtilsTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/commons/BatchWriteUtilsTest.java deleted file mode 100644 index 527ae13c9..000000000 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/commons/BatchWriteUtilsTest.java +++ /dev/null @@ -1,37 +0,0 @@ -package org.hypertrace.core.documentstore.commons; - -import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import java.sql.Statement; -import org.junit.jupiter.api.Test; - -class BatchWriteUtilsTest { - - @Test - void isBatchFullySuccessful_allPositive_returnsTrue() { - assertTrue(BatchWriteUtils.isBatchFullySuccessful(new int[] {1, 1, 2}, 3)); - } - - @Test - void isBatchFullySuccessful_successNoInfo_returnsTrue() { - assertTrue( - BatchWriteUtils.isBatchFullySuccessful( - new int[] {Statement.SUCCESS_NO_INFO, Statement.SUCCESS_NO_INFO}, 2)); - } - - @Test - void isBatchFullySuccessful_executeFailed_returnsFalse() { - assertFalse(BatchWriteUtils.isBatchFullySuccessful(new int[] {1, Statement.EXECUTE_FAILED}, 2)); - } - - @Test - void isBatchFullySuccessful_lengthMismatch_returnsFalse() { - assertFalse(BatchWriteUtils.isBatchFullySuccessful(new int[] {1}, 2)); - } - - @Test - void isBatchFullySuccessful_null_returnsFalse() { - assertFalse(BatchWriteUtils.isBatchFullySuccessful(null, 0)); - } -} diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresCollectionTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresCollectionTest.java index f43885383..fa438f560 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresCollectionTest.java +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresCollectionTest.java @@ -34,6 +34,7 @@ import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; +import java.sql.Statement; import java.time.Clock; import java.util.List; import java.util.Optional; @@ -54,6 +55,7 @@ import org.hypertrace.core.documentstore.query.Query; import org.hypertrace.core.documentstore.query.SortingSpec; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.ArgumentCaptor; @@ -1237,6 +1239,38 @@ void testUpsertSQLException() throws SQLException, IOException { verify(mockUpsertPreparedStatement, times(1)).setString(eq(3), any()); } + @Nested + class BatchFullySuccessful { + + @Test + void allPositive_returnsTrue() { + assertTrue(PostgresCollection.isBatchFullySuccessful(new int[] {1, 1, 2}, 3)); + } + + @Test + void successNoInfo_returnsTrue() { + assertTrue( + PostgresCollection.isBatchFullySuccessful( + new int[] {Statement.SUCCESS_NO_INFO, Statement.SUCCESS_NO_INFO}, 2)); + } + + @Test + void executeFailed_returnsFalse() { + assertFalse( + PostgresCollection.isBatchFullySuccessful(new int[] {1, Statement.EXECUTE_FAILED}, 2)); + } + + @Test + void lengthMismatch_returnsFalse() { + assertFalse(PostgresCollection.isBatchFullySuccessful(new int[] {1}, 2)); + } + + @Test + void nullCounts_returnsFalse() { + assertFalse(PostgresCollection.isBatchFullySuccessful(null, 0)); + } + } + private void mockResultSetMetadata() throws SQLException { when(mockResultSetMetaData.getColumnName(1)).thenReturn("quantity"); when(mockResultSetMetaData.getColumnType(1)).thenReturn(INTEGER); From d81d04b1b5402a8f4b7429349eb82906378c3f4b Mon Sep 17 00:00:00 2001 From: Bhuvan506 Date: Fri, 28 Aug 2026 12:17:16 +0530 Subject: [PATCH 3/3] Address review: resource cleanup and Postgres batch failure logging Close Mongo cursors / Postgres ResultSets on bulkUpsertAndReturnOlderDocuments failure paths. For Postgres, rely on BatchUpdateException (log partial updateCounts) instead of a success-path batch length check that always passed. Add Mongo/PG bulkUpsert consistency integration tests. Co-authored-by: Cursor --- .../MongoPostgresWriteConsistencyTest.java | 58 ++++++++++++++++ .../documentstore/mongo/MongoCollection.java | 23 ++++++- .../postgres/FlatPostgresCollection.java | 14 ++-- .../postgres/PostgresCollection.java | 69 +++++++++---------- .../postgres/PostgresCollectionTest.java | 34 --------- 5 files changed, 115 insertions(+), 83 deletions(-) diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/MongoPostgresWriteConsistencyTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/MongoPostgresWriteConsistencyTest.java index 3fc8463c8..4819679d5 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/MongoPostgresWriteConsistencyTest.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/MongoPostgresWriteConsistencyTest.java @@ -132,6 +132,64 @@ void testBulkUpsert(String storeName) throws Exception { } } + @ParameterizedTest(name = "{0}: bulkUpsert returns true and all docs are readable") + @ArgumentsSource(AllStoresProvider.class) + void testBulkUpsertSuccessIsConsistentAcrossStores(String storeName) throws Exception { + Collection collection = getCollection(storeName); + Map documents = new LinkedHashMap<>(); + List docIds = new java.util.ArrayList<>(); + for (int i = 1; i <= 5; i++) { + String docId = generateDocId("bulk-ok-" + i); + docIds.add(docId); + documents.put(createKey(docId), createTestDocument(docId)); + } + + assertTrue( + collection.bulkUpsert(documents), + storeName + " bulkUpsert should return true for a complete batch"); + + for (String docId : docIds) { + Query query = buildQueryById(docId); + try (CloseableIterator iterator = collection.find(query)) { + assertTrue(iterator.hasNext(), storeName + " missing doc " + docId); + Document doc = iterator.next(); + JsonNode json = OBJECT_MAPPER.readTree(doc.toJson()); + assertEquals("TestItem", json.get("item").asText()); + assertFalse(iterator.hasNext()); + } + } + } + + @ParameterizedTest(name = "{0}: bulkUpsertAndReturnOlderDocuments is consistent") + @ArgumentsSource(AllStoresProvider.class) + void testBulkUpsertAndReturnOlderDocumentsIsConsistent(String storeName) throws Exception { + Collection collection = getCollection(storeName); + String docId = generateDocId("bulk-older"); + insertTestDocument(docId, collection); + + Map updates = new HashMap<>(); + ObjectNode updated = OBJECT_MAPPER.createObjectNode(); + updated.put("id", getKeyString(docId)); + updated.put("item", "AfterUpsert"); + updated.put("price", 42); + updates.put(createKey(docId), new JSONDocument(updated)); + + try (CloseableIterator older = + collection.bulkUpsertAndReturnOlderDocuments(updates)) { + assertTrue(older.hasNext()); + JsonNode before = OBJECT_MAPPER.readTree(older.next().toJson()); + assertEquals("TestItem", before.get("item").asText()); + assertFalse(older.hasNext()); + } + + try (CloseableIterator after = collection.find(buildQueryById(docId))) { + assertTrue(after.hasNext()); + JsonNode json = OBJECT_MAPPER.readTree(after.next().toJson()); + assertEquals("AfterUpsert", json.get("item").asText()); + assertEquals(42, json.get("price").asInt()); + } + } + @ParameterizedTest(name = "{0}: bulkUpsert merges fields (does not replace entire document)") @ArgumentsSource(AllStoresProvider.class) void testBulkUpsertMergesFields(String storeName) throws Exception { diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/MongoCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/MongoCollection.java index 29a91e685..60aa7cd8d 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/MongoCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/mongo/MongoCollection.java @@ -732,11 +732,12 @@ static boolean isBulkUpsertComplete(final BulkWriteResult result, final int requ @Override public CloseableIterator bulkUpsertAndReturnOlderDocuments(Map documents) throws IOException { + MongoCursor mongoCursor = null; try { // First get all the documents for the given keys. FindIterable cursor = collection.find(selectionCriteriaForKeys(documents.keySet())); - final MongoCursor mongoCursor = cursor.cursor(); + mongoCursor = cursor.cursor(); // Now go ahead and do the bulk upsert. BulkWriteResult result = bulkUpsertImpl(documents); @@ -750,13 +751,31 @@ public CloseableIterator bulkUpsertAndReturnOlderDocuments(Map iterator = convertToDocumentIterator(mongoCursor); + mongoCursor = null; // ownership transferred to the iterator + return iterator; } catch (JsonProcessingException e) { + closeQuietly(mongoCursor); LOGGER.error("Error during bulk upsert for documents:{}", documents, e); throw new IOException("Error during bulk upsert."); + } catch (RuntimeException e) { + closeQuietly(mongoCursor); + throw e; + } + } + + private static void closeQuietly(final MongoCursor cursor) { + if (cursor != null) { + try { + cursor.close(); + } catch (Exception e) { + LOGGER.warn("Failed to close Mongo cursor after bulk upsert failure", e); + } } } 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 1aab5c0ba..e4c25a1aa 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 @@ -430,19 +430,15 @@ public boolean bulkUpsert(Map documents) { if (LOGGER.isDebugEnabled()) { LOGGER.debug("Bulk upsert results: {}", Arrays.toString(results)); } - if (!isBatchFullySuccessful(results, parsedDocuments.size())) { - LOGGER.error( - "Incomplete bulkUpsert. requested={}, submitted={}, updateCounts={}", - documents.size(), - parsedDocuments.size(), - Arrays.toString(results)); - return false; - } return true; } } catch (BatchUpdateException e) { - LOGGER.error("BatchUpdateException in bulkUpsert", e); + LOGGER.error( + "BatchUpdateException in bulkUpsert. requested={}, updateCounts={}", + documents.size(), + Arrays.toString(e.getUpdateCounts()), + e); } catch (SQLException e) { LOGGER.error( "SQLException in bulkUpsert. SQLState: {} Error Code: {}", diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java index 2719b6a40..bdc102d65 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java @@ -754,17 +754,14 @@ public boolean bulkUpsert(Map documents) { LOGGER.debug("Write result: {}", Arrays.toString(updateCounts)); } - if (!isBatchFullySuccessful(updateCounts, documents.size())) { - LOGGER.error( - "Incomplete bulk upsert for documents. requested={}, updateCounts={}", - documents.size(), - Arrays.toString(updateCounts)); - return false; - } - return true; } catch (BatchUpdateException e) { - LOGGER.error("BatchUpdateException bulk inserting documents.", e); + // Partial application: some entries may have succeeded before EXECUTE_FAILED. + LOGGER.error( + "BatchUpdateException bulk inserting documents. requested={}, updateCounts={}", + documents.size(), + Arrays.toString(e.getUpdateCounts()), + e); } catch (SQLException e) { LOGGER.error( "SQLException bulk inserting documents. SQLState: {} Error Code:{}", @@ -782,6 +779,8 @@ public boolean bulkUpsert(Map documents) { public CloseableIterator bulkUpsertAndReturnOlderDocuments(Map documents) throws IOException { String query = null; + PreparedStatement preparedStatement = null; + ResultSet resultSet = null; try { String collect = documents.keySet().stream() @@ -801,32 +800,45 @@ public CloseableIterator bulkUpsertAndReturnOlderDocuments(Map iterator = new PostgresResultIterator(resultSet); + resultSet = null; // ownership transferred to the iterator + return iterator; + } catch (BatchUpdateException e) { + LOGGER.error( + "BatchUpdateException bulk inserting documents. requested={}, updateCounts={}", + documents.size(), + Arrays.toString(e.getUpdateCounts()), + e); } catch (IOException e) { LOGGER.error("SQLException bulk inserting documents. documents: {}", documents, e); } catch (SQLException e) { LOGGER.error("SQLException querying documents. query: {}", query, e); + } finally { + closeQuietly(resultSet); } throw new IOException("Could not bulk upsert the documents."); } + private static void closeQuietly(final ResultSet resultSet) { + if (resultSet != null) { + try { + resultSet.close(); + } catch (SQLException e) { + LOGGER.warn("Failed to close ResultSet after bulk upsert failure", e); + } + } + } + @Override public void drop() { String dropTableSQL = String.format("DROP TABLE IF EXISTS %s", tableIdentifier); @@ -1052,25 +1064,6 @@ private int[] bulkUpsertImpl(Map documents) throws SQLException, } } - /** - * Returns true when every batch entry completed without {@link Statement#EXECUTE_FAILED} and the - * result length matches the number of operations submitted. - * - *

{@link Statement#SUCCESS_NO_INFO} (-2) and positive update counts are treated as success. - */ - @VisibleForTesting - static boolean isBatchFullySuccessful(final int[] updateCounts, final int expectedSize) { - if (updateCounts == null || updateCounts.length != expectedSize) { - return false; - } - for (final int count : updateCounts) { - if (count == Statement.EXECUTE_FAILED) { - return false; - } - } - return true; - } - @VisibleForTesting JsonNode getJsonNodeAtPath(String path, JsonNode rootNode, boolean createPathIfMissing) { if (StringUtils.isEmpty(path)) { diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresCollectionTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresCollectionTest.java index fa438f560..f43885383 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresCollectionTest.java +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/PostgresCollectionTest.java @@ -34,7 +34,6 @@ import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; -import java.sql.Statement; import java.time.Clock; import java.util.List; import java.util.Optional; @@ -55,7 +54,6 @@ import org.hypertrace.core.documentstore.query.Query; import org.hypertrace.core.documentstore.query.SortingSpec; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.ArgumentCaptor; @@ -1239,38 +1237,6 @@ void testUpsertSQLException() throws SQLException, IOException { verify(mockUpsertPreparedStatement, times(1)).setString(eq(3), any()); } - @Nested - class BatchFullySuccessful { - - @Test - void allPositive_returnsTrue() { - assertTrue(PostgresCollection.isBatchFullySuccessful(new int[] {1, 1, 2}, 3)); - } - - @Test - void successNoInfo_returnsTrue() { - assertTrue( - PostgresCollection.isBatchFullySuccessful( - new int[] {Statement.SUCCESS_NO_INFO, Statement.SUCCESS_NO_INFO}, 2)); - } - - @Test - void executeFailed_returnsFalse() { - assertFalse( - PostgresCollection.isBatchFullySuccessful(new int[] {1, Statement.EXECUTE_FAILED}, 2)); - } - - @Test - void lengthMismatch_returnsFalse() { - assertFalse(PostgresCollection.isBatchFullySuccessful(new int[] {1}, 2)); - } - - @Test - void nullCounts_returnsFalse() { - assertFalse(PostgresCollection.isBatchFullySuccessful(null, 0)); - } - } - private void mockResultSetMetadata() throws SQLException { when(mockResultSetMetaData.getColumnName(1)).thenReturn("quantity"); when(mockResultSetMetaData.getColumnType(1)).thenReturn(INTEGER);