Return false from bulkUpsert when write result is incomplete

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<Key, Document> documents);
 

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<Key, Document> 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<Key, Document> 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<Document> bulkUpsertAndReturnOlderDocuments(Map<Key, Document> documents)
       throws IOException {
@@ -714,6 +741,17 @@ public CloseableIterator<Document> bulkUpsertAndReturnOlderDocuments(Map<Key, Do
       // Now go ahead and do the bulk upsert.
       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);
+        throw new IOException("Incomplete bulk upsert.");
+      }
 
       return convertToDocumentIterator(mongoCursor);
     } catch (JsonProcessingException e) {

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

@@ -430,6 +430,14 @@ public boolean bulkUpsert(Map<Key, Document> 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;
       }
 

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

@@ -754,6 +754,14 @@ public boolean bulkUpsert(Map<Key, Document> 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);
@@ -801,6 +809,13 @@ public CloseableIterator<Document> bulkUpsertAndReturnOlderDocuments(Map<Key, Do
       if (LOGGER.isDebugEnabled()) {
         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));
+        throw new IOException("Incomplete bulk upsert.");
+      }
 
       return new PostgresResultIterator(resultSet);
     } catch (IOException e) {
@@ -1037,6 +1052,25 @@ private int[] bulkUpsertImpl(Map<Key, Document> 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.
+   *
+   * <p>{@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)) {

document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoCollectionTest.java

@@ -31,14 +31,18 @@
 import com.mongodb.BasicDBObject;
 import com.mongodb.MongoNamespace;
 import com.mongodb.ReadPreference;
+import com.mongodb.bulk.BulkWriteResult;
+import com.mongodb.bulk.BulkWriteUpsert;
 import com.mongodb.client.AggregateIterable;
 import com.mongodb.client.FindIterable;
 import com.mongodb.client.MongoCursor;
+import com.mongodb.client.model.BulkWriteOptions;
 import com.mongodb.client.model.FindOneAndUpdateOptions;
 import com.mongodb.client.result.UpdateResult;
 import java.io.IOException;
 import java.time.Clock;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 import org.bson.BsonString;
@@ -49,6 +53,7 @@
 import org.hypertrace.core.documentstore.JSONDocument;
 import org.hypertrace.core.documentstore.Key;
 import org.hypertrace.core.documentstore.Query;
+import org.hypertrace.core.documentstore.SingleValueKey;
 import org.hypertrace.core.documentstore.expression.impl.ConstantExpression;
 import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression;
 import org.hypertrace.core.documentstore.expression.impl.LogicalExpression;
@@ -452,4 +457,83 @@ void testBulkUpdateWithoutUpdates() {
                   UpdateOptions.DEFAULT_UPDATE_OPTIONS));
     }
   }
+
+  @Nested
+  class BulkUpsert {
+
+    @Test
+    void returnsTrueWhenAllDocumentsMatched() throws Exception {
+      Document document = new JSONDocument("{\"planet\": \"Mars\"}");
+      Map<Key, Document> 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<Key, Document> documents =
+          Map.of(
+              new SingleValueKey("default", "k1"), document,
+              new SingleValueKey("default", "k2"), document);
+
+      List<BulkWriteUpsert> 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<Key, Document> 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<Key, Document> 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));
+    }
+  }
 }

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);