From 70abb050213ff90e8e565e349781b590e6518830 Mon Sep 17 00:00:00 2001 From: bhubam Date: Tue, 6 Oct 2026 13:39:19 +0530 Subject: [PATCH 1/6] Add Collection#upsertAndReturnOlderDocument for atomic pre-write snapshots Callers that need the previous version of a document in order to emit a change event currently have to issue a separate read before the write. That read is not atomic with respect to the write, so a concurrent upsert on the same key can land in between and the "previous" snapshot ends up already holding the new value. Add a merge-upsert that hands back the pre-write document as part of the write itself. Mongo uses findOneAndUpdate with ReturnDocument.BEFORE; both Postgres collections do the read and the upsert inside a single transaction, with the read taking a row lock. An empty Optional means the document did not exist and this call created it. Co-authored-by: Cursor --- .../MongoPostgresWriteConsistencyTest.java | 67 ++++++++++++++++ .../postgres/PostgresDocStoreTest.java | 40 ++++++++++ .../core/documentstore/Collection.java | 9 +++ .../documentstore/mongo/MongoCollection.java | 20 +++++ .../postgres/FlatPostgresCollection.java | 79 +++++++++++++++++++ .../postgres/PostgresCollection.java | 39 +++++++++ .../mongo/MongoCollectionTest.java | 47 +++++++++++ 7 files changed, 301 insertions(+) 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..51e6650cb 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 @@ -495,6 +495,73 @@ void testUpsertAndReturnExistingDoc(String storeName) throws Exception { } } + @Nested + @DisplayName("UpsertAndReturnOlderDocument Consistency Tests") + class UpsertAndReturnOlderDocumentConsistencyTests { + + @ParameterizedTest(name = "{0}: upsertAndReturnOlderDocument new document") + @ArgumentsSource(AllStoresProvider.class) + void testUpsertAndReturnOlderDocumentNewDoc(String storeName) throws Exception { + String docId = generateDocId("upsert-return-older-new"); + Key key = createKey(docId); + + Collection collection = getCollection(storeName); + Document document = createTestDocument(docId); + + Optional returned = collection.upsertAndReturnOlderDocument(key, document); + + assertTrue(returned.isEmpty()); + + Query query = buildQueryById(docId); + try (CloseableIterator iterator = collection.find(query)) { + assertTrue(iterator.hasNext()); + Document retrievedDoc = iterator.next(); + JsonNode json = OBJECT_MAPPER.readTree(retrievedDoc.toJson()); + + assertEquals("TestItem", json.get("item").asText()); + assertEquals(100, json.get("price").asInt()); + assertEquals(50, json.get("quantity").asInt()); + } + } + + @ParameterizedTest(name = "{0}: upsertAndReturnOlderDocument existing document") + @ArgumentsSource(AllStoresProvider.class) + void testUpsertAndReturnOlderDocumentExistingDoc(String storeName) throws Exception { + String docId = generateDocId("upsert-return-older-existing"); + Key key = createKey(docId); + + Collection collection = getCollection(storeName); + Document initialDoc = createTestDocument(docId); + collection.upsert(key, initialDoc); + + ObjectNode partialNode = OBJECT_MAPPER.createObjectNode(); + partialNode.put("id", getKeyString(docId)); + partialNode.put("item", "UpdatedItem"); + partialNode.put("price", 999); + Document partialDoc = new JSONDocument(partialNode); + + Optional returned = collection.upsertAndReturnOlderDocument(key, partialDoc); + + assertTrue(returned.isPresent()); + JsonNode returnedJson = OBJECT_MAPPER.readTree(returned.get().toJson()); + assertEquals("TestItem", returnedJson.get("item").asText()); + assertEquals(100, returnedJson.get("price").asInt()); + assertEquals(50, returnedJson.get("quantity").asInt()); + + Query query = buildQueryById(docId); + try (CloseableIterator iterator = collection.find(query)) { + assertTrue(iterator.hasNext()); + Document retrievedDoc = iterator.next(); + JsonNode json = OBJECT_MAPPER.readTree(retrievedDoc.toJson()); + + assertEquals("UpdatedItem", json.get("item").asText()); + assertEquals(999, json.get("price").asInt()); + assertEquals(50, json.get("quantity").asInt()); + assertTrue(json.get("in_stock").asBoolean()); + } + } + } + @Nested @DisplayName("CreateOrReplace Consistency Tests") class CreateOrReplaceConsistencyTests { diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresDocStoreTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresDocStoreTest.java index 2d32e5b52..02aff3da2 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresDocStoreTest.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresDocStoreTest.java @@ -18,12 +18,15 @@ import java.util.Iterator; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Properties; +import org.hypertrace.core.documentstore.CloseableIterator; import org.hypertrace.core.documentstore.Collection; import org.hypertrace.core.documentstore.Datastore; import org.hypertrace.core.documentstore.DatastoreProvider; import org.hypertrace.core.documentstore.Document; import org.hypertrace.core.documentstore.Filter; +import org.hypertrace.core.documentstore.JSONDocument; import org.hypertrace.core.documentstore.Key; import org.hypertrace.core.documentstore.Query; import org.hypertrace.core.documentstore.SingleValueKey; @@ -116,6 +119,43 @@ public void testUpsertAndReturn() throws IOException { assertEquals(document.toJson(), resultDocument.toJson()); } + @Test + public void testUpsertAndReturnOlderDocument() throws IOException { + Collection collection = datastore.getCollection(COLLECTION_NAME); + Key key = new SingleValueKey("default", "olderDocKey"); + + ObjectMapper objectMapper = new ObjectMapper(); + ObjectNode initialNode = objectMapper.createObjectNode(); + initialNode.put("foo1", "bar1"); + initialNode.put("foo2", "bar2"); + Document initialDoc = new JSONDocument(initialNode); + collection.upsert(key, initialDoc); + + ObjectNode partialNode = objectMapper.createObjectNode(); + partialNode.put("foo1", "updated"); + Document partialDoc = new JSONDocument(partialNode); + + Optional olderDocument = collection.upsertAndReturnOlderDocument(key, partialDoc); + + assertTrue(olderDocument.isPresent()); + JsonNode olderJson = objectMapper.readTree(olderDocument.get().toJson()); + assertEquals("bar1", olderJson.get("foo1").asText()); + assertEquals("bar2", olderJson.get("foo2").asText()); + + Query query = new Query(); + query.setFilter(Filter.eq("id", key.toString())); + try (CloseableIterator iterator = collection.search(query)) { + assertTrue(iterator.hasNext()); + JsonNode storedJson = objectMapper.readTree(iterator.next().toJson()); + assertEquals("updated", storedJson.get("foo1").asText()); + } + + Optional createdDocument = + collection.upsertAndReturnOlderDocument( + new SingleValueKey("default", "newOlderDocKey"), Utils.createDocument("foo1", "bar1")); + assertFalse(createdDocument.isPresent()); + } + @Test public void test_getJsonNodeAtPath() throws Exception { ObjectMapper objectMapper = new ObjectMapper(); 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..7848cf54e 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 @@ -57,6 +57,15 @@ public interface Collection { @Deprecated(forRemoval = true) Document upsertAndReturn(Key key, Document document) throws IOException; + /** + * Atomically perform the same merge-upsert as {@link #upsertAndReturn(Key, Document)} but return + * the document as it existed BEFORE the write. + * + * @return the pre-write document, or {@link Optional#empty()} if the document did not exist (i.e. + * this call created it) + */ + Optional upsertAndReturnOlderDocument(Key key, Document document) throws IOException; + /** * Update a sub document * 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..b9bbce3a9 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 @@ -300,6 +300,26 @@ public Document upsertAndReturn(Key key, Document document) throws IOException { return dbObjectToDocument(upsertResult); } + @Override + public Optional upsertAndReturnOlderDocument(Key key, Document document) + throws IOException { + BasicDBObject upsertResult = + Failsafe.with(upsertRetryPolicy) + .get( + () -> + collection.findOneAndUpdate( + this.selectionCriteriaForKey(key), + this.prepareUpsert(key, document), + new FindOneAndUpdateOptions() + .upsert(true) + .returnDocument(ReturnDocument.BEFORE))); + if (upsertResult == null) { + return Optional.empty(); + } + + return Optional.of(dbObjectToDocument(upsertResult)); + } + @SuppressWarnings("MismatchedQueryAndUpdateOfCollection") private List prepareForCreateOrReplace(final Key key, final Document document) throws JsonProcessingException { 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 100449189..1e7100157 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 @@ -55,6 +55,7 @@ import org.hypertrace.core.documentstore.UpdateResult; import org.hypertrace.core.documentstore.commons.CommonUpdateValidator; import org.hypertrace.core.documentstore.commons.UpdateValidator; +import org.hypertrace.core.documentstore.expression.impl.KeyExpression; import org.hypertrace.core.documentstore.model.config.postgres.CollectionConfig; import org.hypertrace.core.documentstore.model.exception.DuplicateDocumentException; import org.hypertrace.core.documentstore.model.exception.SchemaMismatchException; @@ -231,6 +232,39 @@ public Document upsertAndReturn(Key key, Document document) throws IOException { } } + @Override + public Optional upsertAndReturnOlderDocument(Key key, Document document) + throws IOException { + final String tableName = tableIdentifier.getTableName(); + + try (final Connection connection = client.getTransactionalConnection()) { + try { + final Query query = + Query.builder() + .setFilter( + org.hypertrace.core.documentstore.query.Filter.builder() + .expression(KeyExpression.from(key.toString())) + .build()) + .build(); + + final Optional beforeDocument = selectFirstDocument(connection, query); + upsertOnConnection(connection, key, document, tableName); + connection.commit(); + return beforeDocument; + } catch (final Exception e) { + connection.rollback(); + throw e; + } + } catch (final Exception e) { + LOGGER.error( + "Exception upserting document and returning older version. key: {} content:{}", + key, + document, + e); + throw new IOException(e); + } + } + @Override public boolean updateSubDoc(Key key, String subDocPath, Document subDocument) { throw new UnsupportedOperationException(WRITE_NOT_SUPPORTED); @@ -1746,6 +1780,51 @@ private String buildCreateOrReplaceSql(List allTableColumns, String pkCo tableIdentifier, columnList, placeholders, pkColumn, setClause); } + private void upsertOnConnection( + Connection connection, Key key, Document document, String tableName) throws IOException { + final List skippedFields = new ArrayList<>(); + + try { + final TypedDocument parsed = parseDocument(document, tableName, skippedFields); + final String pkColumn = getPKForTable(tableName); + final String quotedPkColumn = PostgresUtils.wrapFieldNamesWithDoubleQuotes(pkColumn); + final PostgresDataType pkType = getPrimaryKeyType(tableName, pkColumn); + parsed.add(quotedPkColumn, key.toString(), pkType, false); + + final List allColumns = getAllQuotedColumns(tableName); + final String sql = buildUpsertSql(allColumns, quotedPkColumn); + executeUpsertOnConnection(connection, sql, allColumns, parsed); + } catch (SQLException e) { + LOGGER.error("SQLException in upsert. key: {} content: {}", key, document, e); + throw new IOException(e); + } + } + + private void executeUpsertOnConnection( + Connection connection, String sql, List allColumns, TypedDocument parsed) + throws SQLException { + try (PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(connection, sql)) { + int index = 1; + final Set parsedColumns = new HashSet<>(parsed.getColumns()); + for (String column : allColumns) { + if (parsedColumns.contains(column)) { + setParameter( + connection, + ps, + index++, + parsed.getValue(column), + parsed.getType(column), + parsed.isArray(column)); + } else { + ps.setObject(index++, null); + } + } + try (ResultSet rs = ps.executeQuery()) { + rs.next(); + } + } + } + private boolean executeUpsert(String sql, List allColumns, TypedDocument parsed) throws SQLException { long ta = System.nanoTime(); 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..dc8d9389b 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 @@ -299,6 +299,45 @@ public Document upsertAndReturn(Key key, Document document) throws IOException { } } + @Override + public Optional upsertAndReturnOlderDocument(Key key, Document document) + throws IOException { + final String selectQuery = + String.format("SELECT %s FROM %s WHERE %s = ? FOR UPDATE", DOCUMENT, tableIdentifier, ID); + + try (final Connection connection = client.getTransactionalConnection()) { + try { + Optional beforeDocument = empty(); + try (final PreparedStatement selectStatement = connection.prepareStatement(selectQuery)) { + selectStatement.setString(1, key.toString()); + beforeDocument = getFirstDocument(selectStatement.executeQuery()); + } + + try (final PreparedStatement upsertStatement = + connection.prepareStatement(getUpsertSQL(), Statement.RETURN_GENERATED_KEYS)) { + final String jsonString = prepareDocument(key, document); + upsertStatement.setString(1, key.toString()); + upsertStatement.setString(2, jsonString); + upsertStatement.setString(3, jsonString); + upsertStatement.executeUpdate(); + } + + connection.commit(); + return beforeDocument; + } catch (final Exception e) { + connection.rollback(); + throw e; + } + } catch (final Exception e) { + LOGGER.error( + "Exception upserting document and returning older version. key: {} content:{}", + key, + document, + e); + throw new IOException(e); + } + } + /** * Update the sub document based on subDocPath based on longest key match. * diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoCollectionTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoCollectionTest.java index 172d71b13..afc5e9554 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoCollectionTest.java +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoCollectionTest.java @@ -1,6 +1,7 @@ package org.hypertrace.core.documentstore.mongo; import static com.mongodb.client.model.ReturnDocument.AFTER; +import static com.mongodb.client.model.ReturnDocument.BEFORE; import static java.util.Collections.emptyList; import static org.hypertrace.core.documentstore.expression.operators.LogicalOperator.AND; import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.EQ; @@ -49,6 +50,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 +454,49 @@ void testBulkUpdateWithoutUpdates() { UpdateOptions.DEFAULT_UPDATE_OPTIONS)); } } + + @Nested + class UpsertAndReturnOlderDocumentTest { + + private final Key key = new SingleValueKey("default", "testKey"); + private final Document document; + + UpsertAndReturnOlderDocumentTest() throws IOException { + document = new JSONDocument("{\"item\":\"Soap\",\"price\":100}"); + } + + @Test + void testUpsertAndReturnOlderDocumentReturnsEmptyWhenCreated() throws IOException { + final ArgumentCaptor options = + ArgumentCaptor.forClass(FindOneAndUpdateOptions.class); + + when(collection.findOneAndUpdate( + any(BasicDBObject.class), any(BasicDBObject.class), options.capture())) + .thenReturn(null); + + final Optional result = mongoCollection.upsertAndReturnOlderDocument(key, document); + + assertFalse(result.isPresent()); + assertTrue(options.getValue().isUpsert()); + assertEquals(BEFORE, options.getValue().getReturnDocument()); + } + + @Test + void testUpsertAndReturnOlderDocumentReturnsBeforeImage() throws IOException { + final ArgumentCaptor options = + ArgumentCaptor.forClass(FindOneAndUpdateOptions.class); + final BasicDBObject beforeDocument = new BasicDBObject("item", "OldSoap"); + + when(collection.findOneAndUpdate( + any(BasicDBObject.class), any(BasicDBObject.class), options.capture())) + .thenReturn(beforeDocument); + + final Optional result = mongoCollection.upsertAndReturnOlderDocument(key, document); + + assertTrue(result.isPresent()); + assertTrue(result.get().toJson().contains("OldSoap")); + assertTrue(options.getValue().isUpsert()); + assertEquals(BEFORE, options.getValue().getReturnDocument()); + } + } } From 81a2aef3a8db1f42490ddd31dd7981c189252b83 Mon Sep 17 00:00:00 2001 From: bhubam Date: Tue, 13 Oct 2026 00:30:11 +0530 Subject: [PATCH 2/6] Keep Collection#upsertAndReturnOlderDocument source-compatible and avoid a Postgres deadlock. Add a default UnsupportedOperationException implementation so existing Collection clients still compile, and reuse the locked before-image when preparing the JSON upsert so the follow-up createdTime read cannot block on FOR UPDATE. Co-authored-by: Cursor --- .../core/documentstore/Collection.java | 5 +++- .../postgres/PostgresCollection.java | 26 ++++++++++++++++--- 2 files changed, 26 insertions(+), 5 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 7848cf54e..ced562e3f 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 @@ -64,7 +64,10 @@ public interface Collection { * @return the pre-write document, or {@link Optional#empty()} if the document did not exist (i.e. * this call created it) */ - Optional upsertAndReturnOlderDocument(Key key, Document document) throws IOException; + default Optional upsertAndReturnOlderDocument(Key key, Document document) + throws IOException { + throw new UnsupportedOperationException("upsertAndReturnOlderDocument is not supported"); + } /** * Update a sub document 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 dc8d9389b..7bc9bed21 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 @@ -303,19 +303,22 @@ public Document upsertAndReturn(Key key, Document document) throws IOException { public Optional upsertAndReturnOlderDocument(Key key, Document document) throws IOException { final String selectQuery = - String.format("SELECT %s FROM %s WHERE %s = ? FOR UPDATE", DOCUMENT, tableIdentifier, ID); + String.format("SELECT * FROM %s WHERE %s = ? FOR UPDATE", tableIdentifier, ID); try (final Connection connection = client.getTransactionalConnection()) { try { Optional beforeDocument = empty(); try (final PreparedStatement selectStatement = connection.prepareStatement(selectQuery)) { selectStatement.setString(1, key.toString()); - beforeDocument = getFirstDocument(selectStatement.executeQuery()); + beforeDocument = + getFirstDocument(new PostgresResultIterator(selectStatement.executeQuery())); } try (final PreparedStatement upsertStatement = connection.prepareStatement(getUpsertSQL(), Statement.RETURN_GENERATED_KEYS)) { - final String jsonString = prepareDocument(key, document); + // createdTime must come from this locked row; a second connection would deadlock. + final String jsonString = + prepareDocument(key, document, createdTimeFrom(beforeDocument)); upsertStatement.setString(1, key.toString()); upsertStatement.setString(2, jsonString); upsertStatement.setString(3, jsonString); @@ -1188,6 +1191,11 @@ private long bulkUpdateRequestsWithoutFilter(List requestsWit } private String prepareDocument(Key key, Document document) throws IOException { + return prepareDocument(key, document, getCreatedTime(key)); + } + + private String prepareDocument(Key key, Document document, Optional existingCreatedTime) + throws IOException { String jsonString = document.toJson(); ObjectNode jsonNode = (ObjectNode) MAPPER.readTree(jsonString); @@ -1195,13 +1203,23 @@ private String prepareDocument(Key key, Document document) throws IOException { // update time fields long now = System.currentTimeMillis(); - Optional existingCreatedTime = getCreatedTime(key); jsonNode.put(DocStoreConstants.CREATED_TIME, existingCreatedTime.orElse(now)); jsonNode.put(DocStoreConstants.LAST_UPDATED_TIME, now); return MAPPER.writeValueAsString(jsonNode); } + private Optional createdTimeFrom(Optional document) throws IOException { + if (document.isEmpty()) { + return empty(); + } + JsonNode existingDocument = getDocAsJSON(document.get()); + if (existingDocument.has(DocStoreConstants.CREATED_TIME)) { + return Optional.of(existingDocument.get(DocStoreConstants.CREATED_TIME).asLong()); + } + return empty(); + } + private String prepareDocumentForCreation(final Key key, final Document document) throws IOException { final String jsonString = document.toJson(); From a8de2b7acc4de68401c93eacb8e80137b7b4b08a Mon Sep 17 00:00:00 2001 From: bhubam Date: Tue, 13 Oct 2026 00:37:25 +0530 Subject: [PATCH 3/6] Fix Spotless formatting in PostgresCollection. Co-authored-by: Cursor --- .../core/documentstore/postgres/PostgresCollection.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) 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 7bc9bed21..8c43890d2 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 @@ -317,8 +317,7 @@ public Optional upsertAndReturnOlderDocument(Key key, Document documen try (final PreparedStatement upsertStatement = connection.prepareStatement(getUpsertSQL(), Statement.RETURN_GENERATED_KEYS)) { // createdTime must come from this locked row; a second connection would deadlock. - final String jsonString = - prepareDocument(key, document, createdTimeFrom(beforeDocument)); + final String jsonString = prepareDocument(key, document, createdTimeFrom(beforeDocument)); upsertStatement.setString(1, key.toString()); upsertStatement.setString(2, jsonString); upsertStatement.setString(3, jsonString); From 4537f35c2fc7fd31513cb806e7059d9f8244b92a Mon Sep 17 00:00:00 2001 From: bhubam Date: Tue, 13 Oct 2026 14:39:55 +0530 Subject: [PATCH 4/6] AAP-13921: Return before and after documents from a single atomic write Co-authored-by: Cursor --- .../MongoPostgresWriteConsistencyTest.java | 219 +++++++++++++++++- .../core/documentstore/Collection.java | 25 ++ .../core/documentstore/DocumentChange.java | 57 +++++ .../documentstore/mongo/MongoCollection.java | 50 +++- .../postgres/FlatPostgresCollection.java | 146 +++++++++++- .../documentstore/DocumentChangeTest.java | 30 +++ .../mongo/MongoCollectionTest.java | 52 +++++ 7 files changed, 569 insertions(+), 10 deletions(-) create mode 100644 document-store/src/main/java/org/hypertrace/core/documentstore/DocumentChange.java create mode 100644 document-store/src/test/java/org/hypertrace/core/documentstore/DocumentChangeTest.java 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 51e6650cb..aacc42e95 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 @@ -6,10 +6,13 @@ import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.fail; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.ObjectNode; +import com.typesafe.config.ConfigFactory; import java.io.IOException; +import java.util.ArrayList; import java.util.HashMap; import java.util.HashSet; import java.util.LinkedHashMap; @@ -17,6 +20,10 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.stream.Stream; import org.hypertrace.core.documentstore.model.options.ReturnDocumentType; import org.hypertrace.core.documentstore.model.options.UpdateOptions; @@ -36,7 +43,10 @@ import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.ArgumentsProvider; import org.junit.jupiter.params.provider.ArgumentsSource; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.containers.wait.strategy.Wait; import org.testcontainers.junit.jupiter.Testcontainers; +import org.testcontainers.utility.DockerImageName; /*Validates write consistency b/w Mongo and Postgres*/ @Testcontainers @@ -50,7 +60,7 @@ public class MongoPostgresWriteConsistencyTest extends BaseWriteTest { public static void init() throws IOException { // Start MongoDB and PostgreSQL using BaseWriteTest setup initMongo(); - initPostgres(); + initPostgresForFlatCollection(COLLECTION_NAME); datastoreMap.put(MONGO_STORE, mongoDatastore); datastoreMap.put(POSTGRES_FLAT_STORE, postgresDatastore); @@ -69,6 +79,37 @@ public static void init() throws IOException { LOGGER.info("Test setup complete. Collections ready for both Mongo and PostgresFlat."); } + /** + * Same as {@link BaseWriteTest#initPostgres()} but configures {@code timestampFields} so flat + * collections populate and preserve {@code created_at} / {@code last_updated} (see {@link + * FlatCollectionWriteTest}). + */ + private static void initPostgresForFlatCollection(String tableName) { + postgresContainer = + new GenericContainer<>(DockerImageName.parse("postgres:13.1")) + .withEnv("POSTGRES_PASSWORD", "postgres") + .withEnv("POSTGRES_USER", "postgres") + .withExposedPorts(5432) + .waitingFor(Wait.forListeningPort()); + postgresContainer.start(); + + String postgresConnectionUrl = + String.format("jdbc:postgresql://localhost:%s/", postgresContainer.getMappedPort(5432)); + + Map postgresConfig = new HashMap<>(); + postgresConfig.put("url", postgresConnectionUrl); + postgresConfig.put("user", "postgres"); + postgresConfig.put("password", "postgres"); + postgresConfig.put( + "postgres.collectionConfigs." + tableName + ".timestampFields.created", "created_at"); + postgresConfig.put( + "postgres.collectionConfigs." + tableName + ".timestampFields.lastUpdated", "last_updated"); + + postgresDatastore = + DatastoreProvider.getDatastore("Postgres", ConfigFactory.parseMap(postgresConfig)); + LOGGER.info("Postgres datastore initialized with flat timestampFields for {}", tableName); + } + @BeforeEach public void clearCollections() { collectionMap.get(MONGO_STORE).deleteAll(); @@ -562,6 +603,182 @@ void testUpsertAndReturnOlderDocumentExistingDoc(String storeName) throws Except } } + @Nested + @DisplayName("DocumentChange atomic write tests") + class DocumentChangeConsistencyTests { + + private DocumentChange atomicWrite( + String storeName, Collection collection, Key key, Document document) throws IOException { + if (MONGO_STORE.equals(storeName)) { + return collection.upsertAndReturnChange(key, document); + } + return collection.createOrReplaceAndReturnChange(key, document); + } + + private String createdTimestampField(String storeName) { + return MONGO_STORE.equals(storeName) ? "createdTime" : "created_at"; + } + + @ParameterizedTest(name = "{0}: atomic write on insert returns empty before") + @ArgumentsSource(AllStoresProvider.class) + void testAtomicWriteInsert(String storeName) throws Exception { + String docId = generateDocId("atomic-insert"); + Key key = createKey(docId); + Collection collection = getCollection(storeName); + + DocumentChange change = atomicWrite(storeName, collection, key, createTestDocument(docId)); + + assertTrue(change.getBefore().isEmpty()); + JsonNode afterJson = OBJECT_MAPPER.readTree(change.getAfter().toJson()); + assertEquals("TestItem", afterJson.get("item").asText()); + assertEquals(100, afterJson.get("price").asInt()); + assertNotNull(afterJson.get(createdTimestampField(storeName))); + } + + @ParameterizedTest(name = "{0}: atomic write on update preserves created timestamp") + @ArgumentsSource(AllStoresProvider.class) + void testAtomicWriteUpdate(String storeName) throws Exception { + String docId = generateDocId("atomic-update"); + Key key = createKey(docId); + Collection collection = getCollection(storeName); + String createdField = createdTimestampField(storeName); + + collection.upsert(key, createTestDocument(docId)); + + JsonNode initialJson; + Query query = buildQueryById(docId); + try (CloseableIterator iterator = collection.find(query)) { + initialJson = OBJECT_MAPPER.readTree(iterator.next().toJson()); + } + assertNotNull(initialJson.get(createdField)); + + ObjectNode updateNode = OBJECT_MAPPER.createObjectNode(); + updateNode.put("id", getKeyString(docId)); + updateNode.put("item", "UpdatedItem"); + updateNode.put("price", 999); + Document updateDoc = new JSONDocument(updateNode); + + DocumentChange change = atomicWrite(storeName, collection, key, updateDoc); + + assertTrue(change.getBefore().isPresent()); + JsonNode beforeJson = OBJECT_MAPPER.readTree(change.getBefore().get().toJson()); + assertEquals("TestItem", beforeJson.get("item").asText()); + assertEquals(100, beforeJson.get("price").asInt()); + + JsonNode afterJson = OBJECT_MAPPER.readTree(change.getAfter().toJson()); + assertEquals("UpdatedItem", afterJson.get("item").asText()); + assertEquals(999, afterJson.get("price").asInt()); + assertEquals(initialJson.get(createdField), afterJson.get(createdField)); + } + + @ParameterizedTest(name = "{0}: concurrent atomic writes form a single chain") + @ArgumentsSource(AllStoresProvider.class) + void testConcurrentAtomicWritesFormChain(String storeName) throws Exception { + String docId = generateDocId("atomic-concurrent"); + Key key = createKey(docId); + Collection collection = getCollection(storeName); + int threadCount = 20; + + ExecutorService executor = Executors.newFixedThreadPool(threadCount); + try { + List> tasks = new ArrayList<>(); + for (int seq = 0; seq < threadCount; seq++) { + final int value = seq; + tasks.add( + () -> { + ObjectNode node = OBJECT_MAPPER.createObjectNode(); + node.put("id", getKeyString(docId)); + node.put("price", value); + return atomicWrite(storeName, collection, key, new JSONDocument(node)); + }); + } + + List changes = new ArrayList<>(); + for (Future future : executor.invokeAll(tasks)) { + changes.add(future.get()); + } + + assertChangesFormSingleChain(changes, "price"); + } finally { + executor.shutdownNow(); + } + } + + private void assertChangesFormSingleChain(List changes, String field) + throws IOException { + long createCount = changes.stream().filter(c -> c.getBefore().isEmpty()).count(); + assertEquals(1, createCount, "expected exactly one insert in the chain"); + + Set afterValues = new HashSet<>(); + Set beforeValues = new HashSet<>(); + + for (DocumentChange change : changes) { + JsonNode afterField = OBJECT_MAPPER.readTree(change.getAfter().toJson()).get(field); + assertNotNull(afterField, "after document must include chain field: " + field); + afterValues.add(afterField.asInt()); + change + .getBefore() + .ifPresent( + before -> { + try { + beforeValues.add(OBJECT_MAPPER.readTree(before.toJson()).get(field).asInt()); + } catch (IOException e) { + throw new RuntimeException(e); + } + }); + } + + assertEquals(changes.size(), afterValues.size(), "after seq values must be distinct"); + assertEquals(changes.size() - 1, beforeValues.size()); + + for (Integer beforeValue : beforeValues) { + assertTrue( + afterValues.contains(beforeValue), + "before seq " + beforeValue + " must match some other write's after seq"); + } + + DocumentChange head = + changes.stream() + .filter(c -> c.getBefore().isEmpty()) + .findFirst() + .orElseThrow(() -> new AssertionError("missing head of chain")); + + int current = OBJECT_MAPPER.readTree(head.getAfter().toJson()).get(field).asInt(); + Set remaining = new HashSet<>(); + for (DocumentChange change : changes) { + if (change.getBefore().isPresent()) { + remaining.add(change); + } + } + + for (int step = 1; step < changes.size(); step++) { + int expectedBefore = current; + DocumentChange next = + remaining.stream() + .filter( + c -> { + try { + return OBJECT_MAPPER + .readTree(c.getBefore().get().toJson()) + .get(field) + .asInt() + == expectedBefore; + } catch (IOException e) { + throw new RuntimeException(e); + } + }) + .findFirst() + .orElse(null); + if (next == null) { + fail("chain gap after seq " + current); + } + remaining.remove(next); + current = OBJECT_MAPPER.readTree(next.getAfter().toJson()).get(field).asInt(); + } + assertTrue(remaining.isEmpty()); + } + } + @Nested @DisplayName("CreateOrReplace Consistency Tests") class CreateOrReplaceConsistencyTests { 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 ced562e3f..7da82f856 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 @@ -69,6 +69,18 @@ default Optional upsertAndReturnOlderDocument(Key key, Document docume throw new UnsupportedOperationException("upsertAndReturnOlderDocument is not supported"); } + /** + * Atomically perform the same merge-upsert as {@link #upsertAndReturn(Key, Document)} and return + * the document as it existed immediately before and immediately after this write. + * + *

The returned {@link DocumentChange#getBefore()} is empty if this call created the document. + * The before and after images always reflect the same atomic write, even under concurrent writers + * on the same key. + */ + default DocumentChange upsertAndReturnChange(Key key, Document document) throws IOException { + throw new UnsupportedOperationException("upsertAndReturnChange is not supported"); + } + /** * Update a sub document * @@ -319,6 +331,19 @@ default CloseableIterator bulkCreateOrReplaceReturnOlderDocuments( */ Document createOrReplaceAndReturn(final Key key, final Document document) throws IOException; + /** + * Atomically perform the same full-replace write as {@link #createOrReplace(Key, Document)} and + * return the document as it existed immediately before and immediately after this write. + * + *

The returned {@link DocumentChange#getBefore()} is empty if this call created the document. + * The before and after images always reflect the same atomic write, even under concurrent writers + * on the same key. + */ + default DocumentChange createOrReplaceAndReturnChange(Key key, Document document) + throws IOException { + throw new UnsupportedOperationException("createOrReplaceAndReturnChange is not supported"); + } + /** * Updates existing documents if the corresponding Filter condition evaluates to true * diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/DocumentChange.java b/document-store/src/main/java/org/hypertrace/core/documentstore/DocumentChange.java new file mode 100644 index 000000000..21e59ee95 --- /dev/null +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/DocumentChange.java @@ -0,0 +1,57 @@ +package org.hypertrace.core.documentstore; + +import static java.util.Objects.requireNonNull; + +import java.util.Objects; +import java.util.Optional; + +/** Before and after images from a single atomic write. */ +public final class DocumentChange { + + private final Document before; + private final Document after; + + private DocumentChange(final Document before, final Document after) { + this.before = before; + this.after = requireNonNull(after, "after"); + } + + public static DocumentChange created(final Document after) { + return new DocumentChange(null, after); + } + + public static DocumentChange updated(final Document before, final Document after) { + requireNonNull(before, "before"); + return new DocumentChange(before, after); + } + + public Optional getBefore() { + return Optional.ofNullable(before); + } + + public Document getAfter() { + return after; + } + + @Override + public boolean equals(final Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + DocumentChange that = (DocumentChange) o; + return Objects.equals(before, that.before) && after.equals(that.after); + } + + @Override + public int hashCode() { + return Objects.hash(before, after); + } + + @Override + public String toString() { + return "DocumentChange{before=" + before + ", after=" + after + '}'; + } +} 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 b9bbce3a9..316e5ce48 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 @@ -56,6 +56,7 @@ import org.hypertrace.core.documentstore.Collection; import org.hypertrace.core.documentstore.CreateResult; import org.hypertrace.core.documentstore.Document; +import org.hypertrace.core.documentstore.DocumentChange; import org.hypertrace.core.documentstore.Filter; import org.hypertrace.core.documentstore.Key; import org.hypertrace.core.documentstore.Query; @@ -320,6 +321,27 @@ public Optional upsertAndReturnOlderDocument(Key key, Document documen return Optional.of(dbObjectToDocument(upsertResult)); } + @Override + public DocumentChange upsertAndReturnChange(Key key, Document document) throws IOException { + final long now = System.currentTimeMillis(); + final BasicDBObject update = prepareUpsert(key, document, now); + final BasicDBObject beforeDb = + Failsafe.with(upsertRetryPolicy) + .get( + () -> + collection.findOneAndUpdate( + this.selectionCriteriaForKey(key), + update, + new FindOneAndUpdateOptions() + .upsert(true) + .returnDocument(ReturnDocument.BEFORE))); + final BasicDBObject afterDb = deriveAfterImageFromUpsert(beforeDb, update, now); + if (beforeDb == null) { + return DocumentChange.created(dbObjectToDocument(afterDb)); + } + return DocumentChange.updated(dbObjectToDocument(beforeDb), dbObjectToDocument(afterDb)); + } + @SuppressWarnings("MismatchedQueryAndUpdateOfCollection") private List prepareForCreateOrReplace(final Key key, final Document document) throws JsonProcessingException { @@ -341,13 +363,39 @@ private List prepareForCreateOrReplace(final Key key, final Docum } private BasicDBObject prepareUpsert(Key key, Document document) throws JsonProcessingException { - long now = System.currentTimeMillis(); + return prepareUpsert(key, document, System.currentTimeMillis()); + } + + private BasicDBObject prepareUpsert(Key key, Document document, long now) + throws JsonProcessingException { BasicDBObject setObject = prepareDocument(key, document, now); return new BasicDBObject("$set", setObject) .append("$currentDate", new BasicDBObject(LAST_UPDATE_TIME, true)) .append("$setOnInsert", new BasicDBObject(CREATED_TIME, now)); } + /** Derives the post-upsert document in memory. Must stay in sync with {@link #prepareUpsert}. */ + private BasicDBObject deriveAfterImageFromUpsert( + BasicDBObject beforeDb, BasicDBObject update, long now) { + BasicDBObject after = + beforeDb == null ? new BasicDBObject() : new BasicDBObject(beforeDb.toMap()); + + Object setClause = update.get("$set"); + if (setClause instanceof BasicDBObject) { + after.putAll(((BasicDBObject) setClause).toMap()); + } + + if (beforeDb == null) { + Object setOnInsertClause = update.get("$setOnInsert"); + if (setOnInsertClause instanceof BasicDBObject) { + after.putAll(((BasicDBObject) setOnInsertClause).toMap()); + } + } + + after.put(LAST_UPDATE_TIME, new Date(now)); + return after; + } + private BasicDBObject prepareInsert(Key key, Document document) throws JsonProcessingException { long now = System.currentTimeMillis(); BasicDBObject insertDbObject = prepareDocument(key, document, now); 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 1e7100157..e1ab24382 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 @@ -49,13 +49,17 @@ import org.hypertrace.core.documentstore.CreateStatus; import org.hypertrace.core.documentstore.DeepBulkUpdateResult; import org.hypertrace.core.documentstore.Document; +import org.hypertrace.core.documentstore.DocumentChange; import org.hypertrace.core.documentstore.DocumentType; import org.hypertrace.core.documentstore.Filter; import org.hypertrace.core.documentstore.Key; import org.hypertrace.core.documentstore.UpdateResult; import org.hypertrace.core.documentstore.commons.CommonUpdateValidator; import org.hypertrace.core.documentstore.commons.UpdateValidator; -import org.hypertrace.core.documentstore.expression.impl.KeyExpression; +import org.hypertrace.core.documentstore.expression.impl.ConstantExpression; +import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression; +import org.hypertrace.core.documentstore.expression.impl.RelationalExpression; +import org.hypertrace.core.documentstore.expression.operators.RelationalOperator; import org.hypertrace.core.documentstore.model.config.postgres.CollectionConfig; import org.hypertrace.core.documentstore.model.exception.DuplicateDocumentException; import org.hypertrace.core.documentstore.model.exception.SchemaMismatchException; @@ -97,6 +101,7 @@ public class FlatPostgresCollection extends PostgresCollection { private static final Logger LOGGER = LoggerFactory.getLogger(FlatPostgresCollection.class); private static final ObjectMapper MAPPER = new ObjectMapper(); + private static final int MAX_FIRST_INSERT_RACE_RETRIES = 3; private static final String WRITE_NOT_SUPPORTED = "Write operations are not supported for flat collections yet!"; private static final String MISSING_COLUMN_STRATEGY_CONFIG = "missingColumnStrategy"; @@ -239,13 +244,7 @@ public Optional upsertAndReturnOlderDocument(Key key, Document documen try (final Connection connection = client.getTransactionalConnection()) { try { - final Query query = - Query.builder() - .setFilter( - org.hypertrace.core.documentstore.query.Filter.builder() - .expression(KeyExpression.from(key.toString())) - .build()) - .build(); + final Query query = buildPrimaryKeyQuery(tableName, key); final Optional beforeDocument = selectFirstDocument(connection, query); upsertOnConnection(connection, key, document, tableName); @@ -265,6 +264,94 @@ public Optional upsertAndReturnOlderDocument(Key key, Document documen } } + @Override + public DocumentChange createOrReplaceAndReturnChange(Key key, Document document) + throws IOException { + return createOrReplaceAndReturnChangeWithRetry(key, document, false); + } + + private DocumentChange createOrReplaceAndReturnChangeWithRetry( + Key key, Document document, boolean isRetry) throws IOException { + final String tableName = tableIdentifier.getTableName(); + final Query keyQuery = buildPrimaryKeyQuery(tableName, key); + + try { + for (int attempt = 0; attempt < MAX_FIRST_INSERT_RACE_RETRIES; attempt++) { + try (Connection connection = client.getTransactionalConnection()) { + try { + final Optional beforeDocument = selectFirstDocument(connection, keyQuery); + + final List skippedFields = new ArrayList<>(); + final TypedDocument parsed = parseDocument(document, tableName, skippedFields); + final String pkColumn = getPKForTable(tableName); + final String quotedPkColumn = PostgresUtils.wrapFieldNamesWithDoubleQuotes(pkColumn); + final PostgresDataType pkType = getPrimaryKeyType(tableName, pkColumn); + parsed.add(quotedPkColumn, key.toString(), pkType, false); + + final List allColumns = getAllQuotedColumns(tableName); + final String sql = buildCreateOrReplaceSql(allColumns, quotedPkColumn); + final boolean isInsert = + executeCreateOrReplaceReturningIsInsertOnConnection( + connection, sql, allColumns, parsed); + + // FOR UPDATE locks nothing when the row does not exist; concurrent first writers can + // both observe an empty before-image while only one reports is_insert=true. + if (beforeDocument.isEmpty() && !isInsert) { + connection.rollback(); + continue; + } + + final Document afterDocument = + selectFirstDocument(connection, keyQuery) + .orElseThrow( + () -> + new IOException( + "Could not read document after createOrReplace for key: " + key)); + connection.commit(); + + if (beforeDocument.isEmpty()) { + return DocumentChange.created(afterDocument); + } + return DocumentChange.updated(beforeDocument.get(), afterDocument); + } catch (PSQLException e) { + connection.rollback(); + throw e; + } catch (Exception e) { + connection.rollback(); + if (e instanceof IOException) { + throw (IOException) e; + } + throw new IOException(e); + } + } + } + throw new IOException( + "Concurrent first insert on key " + + key + + " could not be resolved after " + + MAX_FIRST_INSERT_RACE_RETRIES + + " attempts"); + } catch (PSQLException e) { + if (!isRetry && shouldRefreshSchemaAndRetry(e.getSQLState())) { + LOGGER.info( + "Schema mismatch detected during createOrReplaceAndReturnChange (SQLState: {}), refreshing schema and retrying. key: {}", + e.getSQLState(), + key); + schemaRegistry.invalidate(tableName); + return createOrReplaceAndReturnChangeWithRetry(key, document, true); + } + LOGGER.error( + "SQLException in createOrReplaceAndReturnChange. key: {} content: {}", key, document, e); + throw new IOException(e); + } catch (IOException e) { + throw e; + } catch (Exception e) { + LOGGER.error( + "Exception in createOrReplaceAndReturnChange. key: {} content: {}", key, document, e); + throw new IOException(e); + } + } + @Override public boolean updateSubDoc(Key key, String subDocPath, Document subDocument) { throw new UnsupportedOperationException(WRITE_NOT_SUPPORTED); @@ -1825,6 +1912,35 @@ private void executeUpsertOnConnection( } } + /** Returns true if INSERT, false if UPDATE. */ + private boolean executeCreateOrReplaceReturningIsInsertOnConnection( + Connection connection, String sql, List allColumns, TypedDocument parsed) + throws SQLException { + try (PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(connection, sql)) { + int index = 1; + final Set parsedColumns = new HashSet<>(parsed.getColumns()); + for (String column : allColumns) { + if (parsedColumns.contains(column)) { + setParameter( + connection, + ps, + index++, + parsed.getValue(column), + parsed.getType(column), + parsed.isArray(column)); + } else { + ps.setObject(index++, null); + } + } + try (ResultSet rs = ps.executeQuery()) { + if (rs.next()) { + return rs.getBoolean("is_insert"); + } + } + return false; + } + } + private boolean executeUpsert(String sql, List allColumns, TypedDocument parsed) throws SQLException { long ta = System.nanoTime(); @@ -1938,6 +2054,20 @@ private boolean shouldRefreshSchemaAndRetry(String sqlState) { || PSQLState.DATATYPE_MISMATCH.getState().equals(sqlState); } + /** KeyExpression always targets an "id" column, but flat tables may use a different PK column. */ + private Query buildPrimaryKeyQuery(String tableName, Key key) { + return Query.builder() + .setFilter( + org.hypertrace.core.documentstore.query.Filter.builder() + .expression( + RelationalExpression.of( + IdentifierExpression.of(getPKForTable(tableName)), + RelationalOperator.EQ, + ConstantExpression.of(key.toString()))) + .build()) + .build(); + } + private String getPKForTable(String tableName) { return schemaRegistry.getPrimaryKeyColumn(tableName).orElse(DEFAULT_PRIMARY_KEY_COLUMN); } diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/DocumentChangeTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/DocumentChangeTest.java new file mode 100644 index 000000000..6f01e5115 --- /dev/null +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/DocumentChangeTest.java @@ -0,0 +1,30 @@ +package org.hypertrace.core.documentstore; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.io.IOException; +import org.junit.jupiter.api.Test; + +class DocumentChangeTest { + + @Test + void createdHasEmptyBefore() throws IOException { + Document after = new JSONDocument("{\"item\":\"after\"}"); + DocumentChange change = DocumentChange.created(after); + assertFalse(change.getBefore().isPresent()); + assertSame(after, change.getAfter()); + } + + @Test + void updatedHasBothImages() throws IOException { + Document before = new JSONDocument("{\"item\":\"before\"}"); + Document after = new JSONDocument("{\"item\":\"after\"}"); + DocumentChange change = DocumentChange.updated(before, after); + assertTrue(change.getBefore().isPresent()); + assertEquals(before, change.getBefore().get()); + assertSame(after, change.getAfter()); + } +} diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoCollectionTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoCollectionTest.java index afc5e9554..80ebe9404 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoCollectionTest.java +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoCollectionTest.java @@ -46,6 +46,7 @@ import org.bson.conversions.Bson; import org.hypertrace.core.documentstore.CloseableIterator; import org.hypertrace.core.documentstore.Document; +import org.hypertrace.core.documentstore.DocumentChange; import org.hypertrace.core.documentstore.Filter; import org.hypertrace.core.documentstore.JSONDocument; import org.hypertrace.core.documentstore.Key; @@ -499,4 +500,55 @@ void testUpsertAndReturnOlderDocumentReturnsBeforeImage() throws IOException { assertEquals(BEFORE, options.getValue().getReturnDocument()); } } + + @Nested + class UpsertAndReturnChangeTest { + + private final Key key = new SingleValueKey("default", "testKey"); + private final Document document; + + UpsertAndReturnChangeTest() throws IOException { + document = new JSONDocument("{\"item\":\"Soap\",\"price\":100}"); + } + + @Test + void testUpsertAndReturnChangeCreated() throws IOException { + final ArgumentCaptor options = + ArgumentCaptor.forClass(FindOneAndUpdateOptions.class); + + when(collection.findOneAndUpdate( + any(BasicDBObject.class), any(BasicDBObject.class), options.capture())) + .thenReturn(null); + + final DocumentChange result = mongoCollection.upsertAndReturnChange(key, document); + + assertFalse(result.getBefore().isPresent()); + assertTrue(result.getAfter().toJson().contains("Soap")); + assertTrue(result.getAfter().toJson().contains("100")); + assertTrue(options.getValue().isUpsert()); + assertEquals(BEFORE, options.getValue().getReturnDocument()); + } + + @Test + void testUpsertAndReturnChangeUpdated() throws IOException { + final ArgumentCaptor options = + ArgumentCaptor.forClass(FindOneAndUpdateOptions.class); + final BasicDBObject beforeDocument = + new BasicDBObject("item", "OldSoap").append("price", 50).append("quantity", 3); + + when(collection.findOneAndUpdate( + any(BasicDBObject.class), any(BasicDBObject.class), options.capture())) + .thenReturn(beforeDocument); + + final DocumentChange result = mongoCollection.upsertAndReturnChange(key, document); + + assertTrue(result.getBefore().isPresent()); + assertTrue(result.getBefore().get().toJson().contains("OldSoap")); + assertTrue(result.getAfter().toJson().contains("Soap")); + assertTrue(result.getAfter().toJson().contains("100")); + assertTrue(result.getAfter().toJson().contains("quantity")); + assertTrue(options.getValue().isUpsert()); + assertEquals(BEFORE, options.getValue().getReturnDocument()); + } + } } From 46ae7dbf699187a49465d595e98b33da5f95ad69 Mon Sep 17 00:00:00 2001 From: bhubam Date: Wed, 14 Oct 2026 12:58:42 +0530 Subject: [PATCH 5/6] AAP-13921: Select returned documents with UpsertOptions and drop upsertAndReturnOlderDocument Co-authored-by: Cursor --- .../MongoPostgresWriteConsistencyTest.java | 372 ++++++++++-------- .../postgres/PostgresDocStoreTest.java | 40 -- .../core/documentstore/Collection.java | 55 +-- .../core/documentstore/DocumentChange.java | 36 +- .../model/options/UpsertOptions.java | 20 + .../documentstore/mongo/MongoCollection.java | 152 ++++--- .../postgres/FlatPostgresCollection.java | 168 +++----- .../postgres/PostgresCollection.java | 58 +-- .../documentstore/DocumentChangeTest.java | 41 +- .../mongo/MongoCollectionTest.java | 190 ++++++--- 10 files changed, 603 insertions(+), 529 deletions(-) create mode 100644 document-store/src/main/java/org/hypertrace/core/documentstore/model/options/UpsertOptions.java 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 aacc42e95..57e505155 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 @@ -27,6 +27,8 @@ import java.util.stream.Stream; import org.hypertrace.core.documentstore.model.options.ReturnDocumentType; import org.hypertrace.core.documentstore.model.options.UpdateOptions; +import org.hypertrace.core.documentstore.model.options.UpsertOptions; +import org.hypertrace.core.documentstore.model.options.UpsertOptions.ReturnDocuments; import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate; import org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue; import org.hypertrace.core.documentstore.model.subdoc.UpdateOperator; @@ -38,6 +40,7 @@ import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtensionContext; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; @@ -130,6 +133,50 @@ public Stream provideArguments(final ExtensionContext context) { } } + private enum AtomicWrite { + MONGO_UPSERT(MONGO_STORE, true), + MONGO_CREATE_OR_REPLACE(MONGO_STORE, false), + POSTGRES_FLAT_UPSERT(POSTGRES_FLAT_STORE, true), + POSTGRES_FLAT_CREATE_OR_REPLACE(POSTGRES_FLAT_STORE, false); + + private final String storeName; + private final boolean merge; + + AtomicWrite(String storeName, boolean merge) { + this.storeName = storeName; + this.merge = merge; + } + + boolean supportsBoth() { + return this != MONGO_CREATE_OR_REPLACE; + } + + String createdTimeField() { + return MONGO_STORE.equals(storeName) ? "createdTime" : "created_at"; + } + } + + private static class AtomicWriteAndReturnDocumentsProvider implements ArgumentsProvider { + + @Override + public Stream provideArguments(final ExtensionContext context) { + return Stream.of(AtomicWrite.values()) + .flatMap( + write -> + Stream.of(ReturnDocuments.values()) + .filter(rd -> rd != ReturnDocuments.BOTH || write.supportsBoth()) + .map(rd -> Arguments.of(write, rd))); + } + } + + private static class BothCapableAtomicWritesProvider implements ArgumentsProvider { + + @Override + public Stream provideArguments(final ExtensionContext context) { + return Stream.of(AtomicWrite.values()).filter(AtomicWrite::supportsBoth).map(Arguments::of); + } + } + /** Inserts a test document into all collections (Mongo and PG) */ private void insertTestDocument(String docId) throws IOException { for (Collection collection : collectionMap.values()) { @@ -537,146 +584,168 @@ void testUpsertAndReturnExistingDoc(String storeName) throws Exception { } @Nested - @DisplayName("UpsertAndReturnOlderDocument Consistency Tests") - class UpsertAndReturnOlderDocumentConsistencyTests { - - @ParameterizedTest(name = "{0}: upsertAndReturnOlderDocument new document") - @ArgumentsSource(AllStoresProvider.class) - void testUpsertAndReturnOlderDocumentNewDoc(String storeName) throws Exception { - String docId = generateDocId("upsert-return-older-new"); - Key key = createKey(docId); - - Collection collection = getCollection(storeName); - Document document = createTestDocument(docId); + @DisplayName("Atomic write with UpsertOptions tests") + class AtomicWriteWithUpsertOptionsTests { - Optional returned = collection.upsertAndReturnOlderDocument(key, document); - - assertTrue(returned.isEmpty()); - - Query query = buildQueryById(docId); - try (CloseableIterator iterator = collection.find(query)) { - assertTrue(iterator.hasNext()); - Document retrievedDoc = iterator.next(); - JsonNode json = OBJECT_MAPPER.readTree(retrievedDoc.toJson()); - - assertEquals("TestItem", json.get("item").asText()); - assertEquals(100, json.get("price").asInt()); - assertEquals(50, json.get("quantity").asInt()); - } + private DocumentChange atomicWrite( + AtomicWrite write, Key key, Document document, ReturnDocuments returnDocuments) + throws IOException { + Collection collection = getCollection(write.storeName); + UpsertOptions options = UpsertOptions.builder().returnDocuments(returnDocuments).build(); + return write.merge + ? collection.upsertAndReturn(key, document, options) + : collection.createOrReplaceAndReturn(key, document, options); } - @ParameterizedTest(name = "{0}: upsertAndReturnOlderDocument existing document") - @ArgumentsSource(AllStoresProvider.class) - void testUpsertAndReturnOlderDocumentExistingDoc(String storeName) throws Exception { - String docId = generateDocId("upsert-return-older-existing"); - Key key = createKey(docId); - - Collection collection = getCollection(storeName); - Document initialDoc = createTestDocument(docId); - collection.upsert(key, initialDoc); - - ObjectNode partialNode = OBJECT_MAPPER.createObjectNode(); - partialNode.put("id", getKeyString(docId)); - partialNode.put("item", "UpdatedItem"); - partialNode.put("price", 999); - Document partialDoc = new JSONDocument(partialNode); - - Optional returned = collection.upsertAndReturnOlderDocument(key, partialDoc); - - assertTrue(returned.isPresent()); - JsonNode returnedJson = OBJECT_MAPPER.readTree(returned.get().toJson()); - assertEquals("TestItem", returnedJson.get("item").asText()); - assertEquals(100, returnedJson.get("price").asInt()); - assertEquals(50, returnedJson.get("quantity").asInt()); - - Query query = buildQueryById(docId); - try (CloseableIterator iterator = collection.find(query)) { + private Document readBack(Collection collection, String docId) throws IOException { + try (CloseableIterator iterator = collection.find(buildQueryById(docId))) { assertTrue(iterator.hasNext()); - Document retrievedDoc = iterator.next(); - JsonNode json = OBJECT_MAPPER.readTree(retrievedDoc.toJson()); - - assertEquals("UpdatedItem", json.get("item").asText()); - assertEquals(999, json.get("price").asInt()); - assertEquals(50, json.get("quantity").asInt()); - assertTrue(json.get("in_stock").asBoolean()); + return iterator.next(); } } - } - - @Nested - @DisplayName("DocumentChange atomic write tests") - class DocumentChangeConsistencyTests { - private DocumentChange atomicWrite( - String storeName, Collection collection, Key key, Document document) throws IOException { - if (MONGO_STORE.equals(storeName)) { - return collection.upsertAndReturnChange(key, document); - } - return collection.createOrReplaceAndReturnChange(key, document); + /** Flat Postgres images are read with SELECT ... FOR UPDATE, which also selects this alias. */ + private JsonNode withoutImplicitId(JsonNode json) { + ObjectNode copy = json.deepCopy(); + copy.remove("_implicit_id"); + return copy; } - private String createdTimestampField(String storeName) { - return MONGO_STORE.equals(storeName) ? "createdTime" : "created_at"; + private Document partialUpdate(String docId) { + ObjectNode updateNode = OBJECT_MAPPER.createObjectNode(); + updateNode.put("id", getKeyString(docId)); + updateNode.put("item", "UpdatedItem"); + updateNode.put("price", 999); + return new JSONDocument(updateNode); } - @ParameterizedTest(name = "{0}: atomic write on insert returns empty before") - @ArgumentsSource(AllStoresProvider.class) - void testAtomicWriteInsert(String storeName) throws Exception { + @ParameterizedTest(name = "{0} {1}: insert") + @ArgumentsSource(AtomicWriteAndReturnDocumentsProvider.class) + void testInsert(AtomicWrite write, ReturnDocuments returnDocuments) throws Exception { String docId = generateDocId("atomic-insert"); Key key = createKey(docId); - Collection collection = getCollection(storeName); + Collection collection = getCollection(write.storeName); - DocumentChange change = atomicWrite(storeName, collection, key, createTestDocument(docId)); + DocumentChange change = atomicWrite(write, key, createTestDocument(docId), returnDocuments); assertTrue(change.getBefore().isEmpty()); - JsonNode afterJson = OBJECT_MAPPER.readTree(change.getAfter().toJson()); - assertEquals("TestItem", afterJson.get("item").asText()); - assertEquals(100, afterJson.get("price").asInt()); - assertNotNull(afterJson.get(createdTimestampField(storeName))); + Document stored = readBack(collection, docId); + JsonNode storedJson = OBJECT_MAPPER.readTree(stored.toJson()); + assertEquals("TestItem", storedJson.get("item").asText()); + assertNotNull(storedJson.get(write.createdTimeField())); + + if (returnDocuments == ReturnDocuments.BEFORE) { + assertTrue(change.getAfter().isEmpty()); + } else { + Document after = change.getAfter().orElseThrow(); + JsonNode afterJson = OBJECT_MAPPER.readTree(after.toJson()); + assertEquals("TestItem", afterJson.get("item").asText()); + assertEquals(100, afterJson.get("price").asInt()); + assertEquals(50, afterJson.get("quantity").asInt()); + assertEquals(storedJson, withoutImplicitId(afterJson)); + } } - @ParameterizedTest(name = "{0}: atomic write on update preserves created timestamp") - @ArgumentsSource(AllStoresProvider.class) - void testAtomicWriteUpdate(String storeName) throws Exception { + @ParameterizedTest(name = "{0} {1}: update") + @ArgumentsSource(AtomicWriteAndReturnDocumentsProvider.class) + void testUpdate(AtomicWrite write, ReturnDocuments returnDocuments) throws Exception { String docId = generateDocId("atomic-update"); Key key = createKey(docId); - Collection collection = getCollection(storeName); - String createdField = createdTimestampField(storeName); + Collection collection = getCollection(write.storeName); + String createdField = write.createdTimeField(); collection.upsert(key, createTestDocument(docId)); + JsonNode initialJson = OBJECT_MAPPER.readTree(readBack(collection, docId).toJson()); + assertNotNull(initialJson.get(createdField)); - JsonNode initialJson; - Query query = buildQueryById(docId); - try (CloseableIterator iterator = collection.find(query)) { - initialJson = OBJECT_MAPPER.readTree(iterator.next().toJson()); + DocumentChange change = atomicWrite(write, key, partialUpdate(docId), returnDocuments); + + if (returnDocuments == ReturnDocuments.AFTER) { + assertTrue(change.getBefore().isEmpty()); + } else { + JsonNode beforeJson = OBJECT_MAPPER.readTree(change.getBefore().orElseThrow().toJson()); + assertEquals(initialJson, withoutImplicitId(beforeJson)); + assertEquals("TestItem", beforeJson.get("item").asText()); + assertEquals(100, beforeJson.get("price").asInt()); } - assertNotNull(initialJson.get(createdField)); - ObjectNode updateNode = OBJECT_MAPPER.createObjectNode(); - updateNode.put("id", getKeyString(docId)); - updateNode.put("item", "UpdatedItem"); - updateNode.put("price", 999); - Document updateDoc = new JSONDocument(updateNode); + JsonNode storedJson = OBJECT_MAPPER.readTree(readBack(collection, docId).toJson()); + assertEquals("UpdatedItem", storedJson.get("item").asText()); + assertEquals(initialJson.get(createdField), storedJson.get(createdField)); + + if (returnDocuments == ReturnDocuments.BEFORE) { + assertTrue(change.getAfter().isEmpty()); + } else { + JsonNode afterJson = OBJECT_MAPPER.readTree(change.getAfter().orElseThrow().toJson()); + assertEquals("UpdatedItem", afterJson.get("item").asText()); + assertEquals(999, afterJson.get("price").asInt()); + assertEquals(initialJson.get(createdField), afterJson.get(createdField)); + JsonNode quantity = afterJson.get("quantity"); + if (write.merge) { + assertEquals(50, quantity.asInt()); + } else { + assertTrue(quantity == null || quantity.isNull()); + } + assertEquals(storedJson, withoutImplicitId(afterJson)); + } + } - DocumentChange change = atomicWrite(storeName, collection, key, updateDoc); + @Test + @DisplayName("Mongo upsert BOTH: derived after-image equals the stored document") + void testMongoUpsertBothAfterImageIsExact() throws Exception { + String docId = generateDocId("atomic-mongo-exact"); + Key key = createKey(docId); + Collection collection = getCollection(MONGO_STORE); + + DocumentChange inserted = + atomicWrite( + AtomicWrite.MONGO_UPSERT, key, createTestDocument(docId), ReturnDocuments.BOTH); + assertTrue(inserted.getBefore().isEmpty()); + Document storedAfterInsert = readBack(collection, docId); + assertEquals(storedAfterInsert, inserted.getAfter().orElseThrow()); + assertTimestampFieldsPresent(storedAfterInsert); + + Thread.sleep(5); + DocumentChange updated = + atomicWrite(AtomicWrite.MONGO_UPSERT, key, partialUpdate(docId), ReturnDocuments.BOTH); + assertEquals(storedAfterInsert, updated.getBefore().orElseThrow()); + Document storedAfterUpdate = readBack(collection, docId); + assertEquals(storedAfterUpdate, updated.getAfter().orElseThrow()); + assertTimestampFieldsPresent(storedAfterUpdate); + + JsonNode insertJson = OBJECT_MAPPER.readTree(storedAfterInsert.toJson()); + JsonNode updateJson = OBJECT_MAPPER.readTree(storedAfterUpdate.toJson()); + assertEquals(insertJson.get("createdTime"), updateJson.get("createdTime")); + assertTrue( + updateJson.get("lastUpdatedTime").asLong() > insertJson.get("lastUpdatedTime").asLong()); + } - assertTrue(change.getBefore().isPresent()); - JsonNode beforeJson = OBJECT_MAPPER.readTree(change.getBefore().get().toJson()); - assertEquals("TestItem", beforeJson.get("item").asText()); - assertEquals(100, beforeJson.get("price").asInt()); + private void assertTimestampFieldsPresent(Document document) throws IOException { + JsonNode json = OBJECT_MAPPER.readTree(document.toJson()); + assertNotNull(json.get("_lastUpdateTime")); + assertNotNull(json.get("createdTime")); + assertNotNull(json.get("lastUpdatedTime")); + } - JsonNode afterJson = OBJECT_MAPPER.readTree(change.getAfter().toJson()); - assertEquals("UpdatedItem", afterJson.get("item").asText()); - assertEquals(999, afterJson.get("price").asInt()); - assertEquals(initialJson.get(createdField), afterJson.get(createdField)); + @Test + @DisplayName("Mongo createOrReplace BOTH is unsupported") + void testMongoCreateOrReplaceBothUnsupported() { + String docId = generateDocId("atomic-mongo-replace-both"); + assertThrows( + UnsupportedOperationException.class, + () -> + atomicWrite( + AtomicWrite.MONGO_CREATE_OR_REPLACE, + createKey(docId), + createTestDocument(docId), + ReturnDocuments.BOTH)); } - @ParameterizedTest(name = "{0}: concurrent atomic writes form a single chain") - @ArgumentsSource(AllStoresProvider.class) - void testConcurrentAtomicWritesFormChain(String storeName) throws Exception { + @ParameterizedTest(name = "{0}: concurrent BOTH writes form a single chain") + @ArgumentsSource(BothCapableAtomicWritesProvider.class) + void testConcurrentWritesFormChain(AtomicWrite write) throws Exception { String docId = generateDocId("atomic-concurrent"); Key key = createKey(docId); - Collection collection = getCollection(storeName); int threadCount = 20; ExecutorService executor = Executors.newFixedThreadPool(threadCount); @@ -689,7 +758,7 @@ void testConcurrentAtomicWritesFormChain(String storeName) throws Exception { ObjectNode node = OBJECT_MAPPER.createObjectNode(); node.put("id", getKeyString(docId)); node.put("price", value); - return atomicWrite(storeName, collection, key, new JSONDocument(node)); + return atomicWrite(write, key, new JSONDocument(node), ReturnDocuments.BOTH); }); } @@ -704,78 +773,47 @@ void testConcurrentAtomicWritesFormChain(String storeName) throws Exception { } } + private int fieldValue(Document document, String field) throws IOException { + JsonNode value = OBJECT_MAPPER.readTree(document.toJson()).get(field); + assertNotNull(value, "document must include chain field: " + field); + return value.asInt(); + } + private void assertChangesFormSingleChain(List changes, String field) throws IOException { - long createCount = changes.stream().filter(c -> c.getBefore().isEmpty()).count(); - assertEquals(1, createCount, "expected exactly one insert in the chain"); + assertEquals( + 1, + changes.stream().filter(c -> c.getBefore().isEmpty()).count(), + "expected exactly one insert in the chain"); + Map afterByBefore = new HashMap<>(); Set afterValues = new HashSet<>(); - Set beforeValues = new HashSet<>(); - + Integer head = null; for (DocumentChange change : changes) { - JsonNode afterField = OBJECT_MAPPER.readTree(change.getAfter().toJson()).get(field); - assertNotNull(afterField, "after document must include chain field: " + field); - afterValues.add(afterField.asInt()); - change - .getBefore() - .ifPresent( - before -> { - try { - beforeValues.add(OBJECT_MAPPER.readTree(before.toJson()).get(field).asInt()); - } catch (IOException e) { - throw new RuntimeException(e); - } - }); + int after = fieldValue(change.getAfter().orElseThrow(), field); + assertTrue(afterValues.add(after), "after values must be distinct"); + if (change.getBefore().isPresent()) { + int before = fieldValue(change.getBefore().get(), field); + assertNull(afterByBefore.put(before, after), "before values must be distinct"); + } else { + head = after; + } } - assertEquals(changes.size(), afterValues.size(), "after seq values must be distinct"); - assertEquals(changes.size() - 1, beforeValues.size()); - - for (Integer beforeValue : beforeValues) { + for (Integer before : afterByBefore.keySet()) { assertTrue( - afterValues.contains(beforeValue), - "before seq " + beforeValue + " must match some other write's after seq"); - } - - DocumentChange head = - changes.stream() - .filter(c -> c.getBefore().isEmpty()) - .findFirst() - .orElseThrow(() -> new AssertionError("missing head of chain")); - - int current = OBJECT_MAPPER.readTree(head.getAfter().toJson()).get(field).asInt(); - Set remaining = new HashSet<>(); - for (DocumentChange change : changes) { - if (change.getBefore().isPresent()) { - remaining.add(change); - } + afterValues.contains(before), + "before value " + before + " must match some write's after value"); } + Integer current = head; for (int step = 1; step < changes.size(); step++) { - int expectedBefore = current; - DocumentChange next = - remaining.stream() - .filter( - c -> { - try { - return OBJECT_MAPPER - .readTree(c.getBefore().get().toJson()) - .get(field) - .asInt() - == expectedBefore; - } catch (IOException e) { - throw new RuntimeException(e); - } - }) - .findFirst() - .orElse(null); + Integer next = afterByBefore.get(current); if (next == null) { - fail("chain gap after seq " + current); + fail("chain gap after value " + current); } - remaining.remove(next); - current = OBJECT_MAPPER.readTree(next.getAfter().toJson()).get(field).asInt(); + current = next; } - assertTrue(remaining.isEmpty()); } } diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresDocStoreTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresDocStoreTest.java index 02aff3da2..2d32e5b52 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresDocStoreTest.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/postgres/PostgresDocStoreTest.java @@ -18,15 +18,12 @@ import java.util.Iterator; import java.util.List; import java.util.Map; -import java.util.Optional; import java.util.Properties; -import org.hypertrace.core.documentstore.CloseableIterator; import org.hypertrace.core.documentstore.Collection; import org.hypertrace.core.documentstore.Datastore; import org.hypertrace.core.documentstore.DatastoreProvider; import org.hypertrace.core.documentstore.Document; import org.hypertrace.core.documentstore.Filter; -import org.hypertrace.core.documentstore.JSONDocument; import org.hypertrace.core.documentstore.Key; import org.hypertrace.core.documentstore.Query; import org.hypertrace.core.documentstore.SingleValueKey; @@ -119,43 +116,6 @@ public void testUpsertAndReturn() throws IOException { assertEquals(document.toJson(), resultDocument.toJson()); } - @Test - public void testUpsertAndReturnOlderDocument() throws IOException { - Collection collection = datastore.getCollection(COLLECTION_NAME); - Key key = new SingleValueKey("default", "olderDocKey"); - - ObjectMapper objectMapper = new ObjectMapper(); - ObjectNode initialNode = objectMapper.createObjectNode(); - initialNode.put("foo1", "bar1"); - initialNode.put("foo2", "bar2"); - Document initialDoc = new JSONDocument(initialNode); - collection.upsert(key, initialDoc); - - ObjectNode partialNode = objectMapper.createObjectNode(); - partialNode.put("foo1", "updated"); - Document partialDoc = new JSONDocument(partialNode); - - Optional olderDocument = collection.upsertAndReturnOlderDocument(key, partialDoc); - - assertTrue(olderDocument.isPresent()); - JsonNode olderJson = objectMapper.readTree(olderDocument.get().toJson()); - assertEquals("bar1", olderJson.get("foo1").asText()); - assertEquals("bar2", olderJson.get("foo2").asText()); - - Query query = new Query(); - query.setFilter(Filter.eq("id", key.toString())); - try (CloseableIterator iterator = collection.search(query)) { - assertTrue(iterator.hasNext()); - JsonNode storedJson = objectMapper.readTree(iterator.next().toJson()); - assertEquals("updated", storedJson.get("foo1").asText()); - } - - Optional createdDocument = - collection.upsertAndReturnOlderDocument( - new SingleValueKey("default", "newOlderDocKey"), Utils.createDocument("foo1", "bar1")); - assertFalse(createdDocument.isPresent()); - } - @Test public void test_getJsonNodeAtPath() throws Exception { ObjectMapper objectMapper = new ObjectMapper(); 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 7da82f856..7e02f24e1 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 @@ -11,6 +11,7 @@ import org.hypertrace.core.documentstore.model.exception.DuplicateDocumentException; import org.hypertrace.core.documentstore.model.options.QueryOptions; import org.hypertrace.core.documentstore.model.options.UpdateOptions; +import org.hypertrace.core.documentstore.model.options.UpsertOptions; import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate; /** Interface spec for common operations on a collection of documents */ @@ -58,27 +59,23 @@ public interface Collection { Document upsertAndReturn(Key key, Document document) throws IOException; /** - * Atomically perform the same merge-upsert as {@link #upsertAndReturn(Key, Document)} but return - * the document as it existed BEFORE the write. + * Atomically perform the same merge-upsert as {@link #upsert(Key, Document)} and return the + * document images selected by {@link UpsertOptions#getReturnDocuments()}. * - * @return the pre-write document, or {@link Optional#empty()} if the document did not exist (i.e. - * this call created it) - */ - default Optional upsertAndReturnOlderDocument(Key key, Document document) - throws IOException { - throw new UnsupportedOperationException("upsertAndReturnOlderDocument is not supported"); - } - - /** - * Atomically perform the same merge-upsert as {@link #upsertAndReturn(Key, Document)} and return - * the document as it existed immediately before and immediately after this write. + *

Whichever images are requested come from the same atomic write, even under concurrent + * writers on the same key: the before-image is exactly the document this write replaced and the + * after-image is exactly the document this write produced. * - *

The returned {@link DocumentChange#getBefore()} is empty if this call created the document. - * The before and after images always reflect the same atomic write, even under concurrent writers - * on the same key. + * @param key Unique key of the document in the collection. + * @param document Document to be upserted. + * @param options Selects which images (BEFORE, AFTER or BOTH) to return. + * @return {@link DocumentChange#getBefore()} is empty if not requested or if this call created + * the document; {@link DocumentChange#getAfter()} is empty if not requested. + * @throws IOException If the operation could not be performed */ - default DocumentChange upsertAndReturnChange(Key key, Document document) throws IOException { - throw new UnsupportedOperationException("upsertAndReturnChange is not supported"); + default DocumentChange upsertAndReturn(Key key, Document document, UpsertOptions options) + throws IOException { + throw new UnsupportedOperationException("upsertAndReturn with UpsertOptions is not supported"); } /** @@ -333,15 +330,23 @@ default CloseableIterator bulkCreateOrReplaceReturnOlderDocuments( /** * Atomically perform the same full-replace write as {@link #createOrReplace(Key, Document)} and - * return the document as it existed immediately before and immediately after this write. + * return the document images selected by {@link UpsertOptions#getReturnDocuments()}. + * + *

Whichever images are requested come from the same atomic write, even under concurrent + * writers on the same key: the before-image is exactly the document this write replaced and the + * after-image is exactly the document this write produced. * - *

The returned {@link DocumentChange#getBefore()} is empty if this call created the document. - * The before and after images always reflect the same atomic write, even under concurrent writers - * on the same key. + * @param key Unique key of the document in the collection. + * @param document The document to be created/replaced + * @param options Selects which images (BEFORE, AFTER or BOTH) to return. + * @return {@link DocumentChange#getBefore()} is empty if not requested or if this call created + * the document; {@link DocumentChange#getAfter()} is empty if not requested. + * @throws IOException If the operation could not be performed */ - default DocumentChange createOrReplaceAndReturnChange(Key key, Document document) - throws IOException { - throw new UnsupportedOperationException("createOrReplaceAndReturnChange is not supported"); + default DocumentChange createOrReplaceAndReturn( + final Key key, final Document document, final UpsertOptions options) throws IOException { + throw new UnsupportedOperationException( + "createOrReplaceAndReturn with UpsertOptions is not supported"); } /** diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/DocumentChange.java b/document-store/src/main/java/org/hypertrace/core/documentstore/DocumentChange.java index 21e59ee95..c78b8da49 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/DocumentChange.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/DocumentChange.java @@ -1,36 +1,40 @@ package org.hypertrace.core.documentstore; -import static java.util.Objects.requireNonNull; - import java.util.Objects; import java.util.Optional; +import javax.annotation.Nullable; -/** Before and after images from a single atomic write. */ +/** + * The document images returned by a single atomic write. Which images are present depends on the + * requested {@link org.hypertrace.core.documentstore.model.options.UpsertOptions.ReturnDocuments}. + */ public final class DocumentChange { - private final Document before; - private final Document after; + @Nullable private final Document before; + @Nullable private final Document after; - private DocumentChange(final Document before, final Document after) { + private DocumentChange(@Nullable final Document before, @Nullable final Document after) { this.before = before; - this.after = requireNonNull(after, "after"); - } - - public static DocumentChange created(final Document after) { - return new DocumentChange(null, after); + this.after = after; } - public static DocumentChange updated(final Document before, final Document after) { - requireNonNull(before, "before"); + /** + * @param before the document before the write; null if not requested or if the write created the + * document + * @param after the document after the write; null if not requested + */ + public static DocumentChange of(@Nullable final Document before, @Nullable final Document after) { return new DocumentChange(before, after); } + /** Empty when the before-image was not requested, or when the write created the document. */ public Optional getBefore() { return Optional.ofNullable(before); } - public Document getAfter() { - return after; + /** Empty when the after-image was not requested. */ + public Optional getAfter() { + return Optional.ofNullable(after); } @Override @@ -42,7 +46,7 @@ public boolean equals(final Object o) { return false; } DocumentChange that = (DocumentChange) o; - return Objects.equals(before, that.before) && after.equals(that.after); + return Objects.equals(before, that.before) && Objects.equals(after, that.after); } @Override diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/model/options/UpsertOptions.java b/document-store/src/main/java/org/hypertrace/core/documentstore/model/options/UpsertOptions.java new file mode 100644 index 000000000..e11ee7cb4 --- /dev/null +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/model/options/UpsertOptions.java @@ -0,0 +1,20 @@ +package org.hypertrace.core.documentstore.model.options; + +import lombok.Builder; +import lombok.Builder.Default; +import lombok.Value; + +@Value +@Builder +public class UpsertOptions { + public static UpsertOptions DEFAULT_UPSERT_OPTIONS = + UpsertOptions.builder().returnDocuments(ReturnDocuments.AFTER).build(); + + @Default ReturnDocuments returnDocuments = ReturnDocuments.AFTER; + + public enum ReturnDocuments { + BEFORE, + AFTER, + BOTH, + } +} 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 316e5ce48..6499d697f 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 @@ -42,6 +42,7 @@ import java.util.Optional; import java.util.Set; import java.util.stream.Collectors; +import javax.annotation.Nullable; import net.jodah.failsafe.Failsafe; import net.jodah.failsafe.RetryPolicy; import org.bson.BsonInt64; @@ -63,6 +64,8 @@ import org.hypertrace.core.documentstore.model.config.ConnectionConfig; import org.hypertrace.core.documentstore.model.exception.DuplicateDocumentException; import org.hypertrace.core.documentstore.model.options.QueryOptions; +import org.hypertrace.core.documentstore.model.options.UpsertOptions; +import org.hypertrace.core.documentstore.model.options.UpsertOptions.ReturnDocuments; import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate; import org.hypertrace.core.documentstore.mongo.query.MongoQueryExecutor; import org.hypertrace.core.documentstore.mongo.update.MongoUpdateExecutor; @@ -301,45 +304,91 @@ public Document upsertAndReturn(Key key, Document document) throws IOException { return dbObjectToDocument(upsertResult); } + /** + * Performs one atomic server-side write (a single {@code findOneAndUpdate} with upsert). + * + *

    + *
  • BEFORE / AFTER: the server returns the requested image from that write. + *
  • BOTH: the server returns the before-image from that write; the after-image is computed by + * applying that same update to the before-image (no second read). This is exact because the + * update only sets top-level fields to literal values, including {@code _lastUpdateTime}, + * which is set from the client clock rather than {@code $currentDate}. + *
+ */ @Override - public Optional upsertAndReturnOlderDocument(Key key, Document document) + public DocumentChange upsertAndReturn(Key key, Document document, UpsertOptions options) throws IOException { - BasicDBObject upsertResult = + final ReturnDocuments returnDocuments = options.getReturnDocuments(); + final BasicDBObject update = + prepareUpsertWithExplicitLastUpdateTime(key, document, System.currentTimeMillis()); + final ReturnDocument serverImage = + returnDocuments == ReturnDocuments.AFTER ? ReturnDocument.AFTER : ReturnDocument.BEFORE; + final BasicDBObject result = Failsafe.with(upsertRetryPolicy) .get( () -> collection.findOneAndUpdate( this.selectionCriteriaForKey(key), - this.prepareUpsert(key, document), - new FindOneAndUpdateOptions() - .upsert(true) - .returnDocument(ReturnDocument.BEFORE))); - if (upsertResult == null) { - return Optional.empty(); + update, + new FindOneAndUpdateOptions().upsert(true).returnDocument(serverImage))); + + switch (returnDocuments) { + case BEFORE: + return DocumentChange.of(toDocumentOrNull(result), null); + case AFTER: + if (result == null) { + throw new IOException("Could not upsert the document with key: " + key); + } + return DocumentChange.of(null, dbObjectToDocument(result)); + case BOTH: + final BasicDBObject after = applyUpsertToBeforeImage(key, result, update); + return DocumentChange.of(toDocumentOrNull(result), dbObjectToDocument(after)); + default: + throw new UnsupportedOperationException("Unhandled return documents: " + returnDocuments); } - - return Optional.of(dbObjectToDocument(upsertResult)); } + /** + * Performs one atomic server-side write (a single {@code findOneAndUpdate} with upsert) and + * returns the BEFORE or AFTER image from that write. + * + *

BOTH is not supported: the replace pipeline evaluates document values as aggregation + * expressions (e.g. a string value starting with {@code $} is read as a field path), so the + * stored document cannot always be derived exactly from the before-image without a second read. + */ @Override - public DocumentChange upsertAndReturnChange(Key key, Document document) throws IOException { - final long now = System.currentTimeMillis(); - final BasicDBObject update = prepareUpsert(key, document, now); - final BasicDBObject beforeDb = - Failsafe.with(upsertRetryPolicy) - .get( - () -> - collection.findOneAndUpdate( - this.selectionCriteriaForKey(key), - update, - new FindOneAndUpdateOptions() - .upsert(true) - .returnDocument(ReturnDocument.BEFORE))); - final BasicDBObject afterDb = deriveAfterImageFromUpsert(beforeDb, update, now); - if (beforeDb == null) { - return DocumentChange.created(dbObjectToDocument(afterDb)); + public DocumentChange createOrReplaceAndReturn( + final Key key, final Document document, final UpsertOptions options) throws IOException { + final ReturnDocuments returnDocuments = options.getReturnDocuments(); + if (returnDocuments == ReturnDocuments.BOTH) { + throw new UnsupportedOperationException( + "createOrReplaceAndReturn with ReturnDocuments.BOTH is not supported on Mongo"); + } + final boolean returnBefore = returnDocuments == ReturnDocuments.BEFORE; + + try { + final FindOneAndUpdateOptions findOptions = + new FindOneAndUpdateOptions() + .upsert(true) + .returnDocument(returnBefore ? ReturnDocument.BEFORE : ReturnDocument.AFTER); + final BasicDBObject result = + collection.findOneAndUpdate( + this.selectionCriteriaForKey(key), + this.prepareForCreateOrReplace(key, document), + findOptions); + LOGGER.debug("Create or replace result: {}", result); + return returnBefore + ? DocumentChange.of(toDocumentOrNull(result), null) + : DocumentChange.of(null, dbObjectToDocument(requireNonNull(result))); + } catch (final Exception e) { + LOGGER.error("Exception creating/replacing document. key: {} content:{}", key, document, e); + throw e; } - return DocumentChange.updated(dbObjectToDocument(beforeDb), dbObjectToDocument(afterDb)); + } + + @Nullable + private static Document toDocumentOrNull(@Nullable final BasicDBObject dbObject) { + return dbObject == null ? null : dbObjectToDocument(dbObject); } @SuppressWarnings("MismatchedQueryAndUpdateOfCollection") @@ -363,36 +412,41 @@ private List prepareForCreateOrReplace(final Key key, final Docum } private BasicDBObject prepareUpsert(Key key, Document document) throws JsonProcessingException { - return prepareUpsert(key, document, System.currentTimeMillis()); - } - - private BasicDBObject prepareUpsert(Key key, Document document, long now) - throws JsonProcessingException { + long now = System.currentTimeMillis(); BasicDBObject setObject = prepareDocument(key, document, now); return new BasicDBObject("$set", setObject) .append("$currentDate", new BasicDBObject(LAST_UPDATE_TIME, true)) .append("$setOnInsert", new BasicDBObject(CREATED_TIME, now)); } - /** Derives the post-upsert document in memory. Must stay in sync with {@link #prepareUpsert}. */ - private BasicDBObject deriveAfterImageFromUpsert( - BasicDBObject beforeDb, BasicDBObject update, long now) { - BasicDBObject after = - beforeDb == null ? new BasicDBObject() : new BasicDBObject(beforeDb.toMap()); - - Object setClause = update.get("$set"); - if (setClause instanceof BasicDBObject) { - after.putAll(((BasicDBObject) setClause).toMap()); - } + /** + * Same as {@link #prepareUpsert} except {@code _lastUpdateTime} is a literal client-clock value, + * so the stored document is fully determined by the before-image and this update. + */ + private BasicDBObject prepareUpsertWithExplicitLastUpdateTime( + Key key, Document document, long now) throws JsonProcessingException { + BasicDBObject setObject = prepareDocument(key, document, now); + setObject.put(LAST_UPDATE_TIME, new Date(now)); + return new BasicDBObject("$set", setObject) + .append("$setOnInsert", new BasicDBObject(CREATED_TIME, now)); + } - if (beforeDb == null) { - Object setOnInsertClause = update.get("$setOnInsert"); - if (setOnInsertClause instanceof BasicDBObject) { - after.putAll(((BasicDBObject) setOnInsertClause).toMap()); - } + /** + * Applies an update from {@link #prepareUpsertWithExplicitLastUpdateTime} to the before-image the + * server returned for that same write, yielding the stored document. Exact because the update + * only sets top-level fields to literal values ({@code $setOnInsert} applies only on insert). + */ + @SuppressWarnings("unchecked") + private static BasicDBObject applyUpsertToBeforeImage( + Key key, @Nullable BasicDBObject before, BasicDBObject update) { + final BasicDBObject after = + before == null + ? new BasicDBObject(ID_KEY, key.toString()) + : new BasicDBObject(before.toMap()); + after.putAll((Map) update.get("$set")); + if (before == null) { + after.putAll((Map) update.get("$setOnInsert")); } - - after.put(LAST_UPDATE_TIME, new Date(now)); return after; } 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 e1ab24382..5b4b600a5 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 @@ -39,6 +39,7 @@ import java.util.Optional; import java.util.Set; import java.util.TreeMap; +import java.util.function.BiFunction; import java.util.stream.Collectors; import org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest; import org.hypertrace.core.documentstore.BulkDeleteResult; @@ -67,6 +68,8 @@ import org.hypertrace.core.documentstore.model.options.QueryOptions; import org.hypertrace.core.documentstore.model.options.ReturnDocumentType; import org.hypertrace.core.documentstore.model.options.UpdateOptions; +import org.hypertrace.core.documentstore.model.options.UpsertOptions; +import org.hypertrace.core.documentstore.model.options.UpsertOptions.ReturnDocuments; import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate; import org.hypertrace.core.documentstore.model.subdoc.UpdateOperator; import org.hypertrace.core.documentstore.postgres.model.PostgresColumnMetadata; @@ -237,41 +240,45 @@ public Document upsertAndReturn(Key key, Document document) throws IOException { } } + /** + * Merge-upserts like {@link #upsert(Key, Document)} and returns the requested images from the + * same transaction. See {@link #writeAndReturnChange}. + */ @Override - public Optional upsertAndReturnOlderDocument(Key key, Document document) + public DocumentChange upsertAndReturn(Key key, Document document, UpsertOptions options) throws IOException { - final String tableName = tableIdentifier.getTableName(); - - try (final Connection connection = client.getTransactionalConnection()) { - try { - final Query query = buildPrimaryKeyQuery(tableName, key); - - final Optional beforeDocument = selectFirstDocument(connection, query); - upsertOnConnection(connection, key, document, tableName); - connection.commit(); - return beforeDocument; - } catch (final Exception e) { - connection.rollback(); - throw e; - } - } catch (final Exception e) { - LOGGER.error( - "Exception upserting document and returning older version. key: {} content:{}", - key, - document, - e); - throw new IOException(e); - } + return writeAndReturnChange( + key, document, options, this::buildUpsertSql, "upsertAndReturn", false); } + /** + * Full-replaces like {@link #createOrReplace(Key, Document)} and returns the requested images + * from the same transaction. See {@link #writeAndReturnChange}. + */ @Override - public DocumentChange createOrReplaceAndReturnChange(Key key, Document document) + public DocumentChange createOrReplaceAndReturn(Key key, Document document, UpsertOptions options) throws IOException { - return createOrReplaceAndReturnChangeWithRetry(key, document, false); + return writeAndReturnChange( + key, document, options, this::buildCreateOrReplaceSql, "createOrReplaceAndReturn", false); } - private DocumentChange createOrReplaceAndReturnChangeWithRetry( - Key key, Document document, boolean isRetry) throws IOException { + /** + * Executes a single-row {@code INSERT ... ON CONFLICT DO UPDATE ... RETURNING is_insert} built by + * {@code sqlBuilder} and reads the requested images in one transaction on one connection. The + * before-image is read with {@code SELECT ... FOR UPDATE}, so no other writer can change the row + * between that read, the write and the after-image read. + */ + private DocumentChange writeAndReturnChange( + Key key, + Document document, + UpsertOptions options, + BiFunction, String, String> sqlBuilder, + String operation, + boolean isRetry) + throws IOException { + final ReturnDocuments returnDocuments = options.getReturnDocuments(); + final boolean readBefore = returnDocuments != ReturnDocuments.AFTER; + final boolean readAfter = returnDocuments != ReturnDocuments.BEFORE; final String tableName = tableIdentifier.getTableName(); final Query keyQuery = buildPrimaryKeyQuery(tableName, key); @@ -279,49 +286,42 @@ private DocumentChange createOrReplaceAndReturnChangeWithRetry( for (int attempt = 0; attempt < MAX_FIRST_INSERT_RACE_RETRIES; attempt++) { try (Connection connection = client.getTransactionalConnection()) { try { - final Optional beforeDocument = selectFirstDocument(connection, keyQuery); + final Document before = + readBefore ? selectFirstDocument(connection, keyQuery).orElse(null) : null; - final List skippedFields = new ArrayList<>(); - final TypedDocument parsed = parseDocument(document, tableName, skippedFields); + final TypedDocument parsed = parseDocument(document, tableName, new ArrayList<>()); final String pkColumn = getPKForTable(tableName); final String quotedPkColumn = PostgresUtils.wrapFieldNamesWithDoubleQuotes(pkColumn); - final PostgresDataType pkType = getPrimaryKeyType(tableName, pkColumn); - parsed.add(quotedPkColumn, key.toString(), pkType, false); - + parsed.add( + quotedPkColumn, key.toString(), getPrimaryKeyType(tableName, pkColumn), false); final List allColumns = getAllQuotedColumns(tableName); - final String sql = buildCreateOrReplaceSql(allColumns, quotedPkColumn); final boolean isInsert = - executeCreateOrReplaceReturningIsInsertOnConnection( - connection, sql, allColumns, parsed); + executeReturningIsInsertOnConnection( + connection, sqlBuilder.apply(allColumns, quotedPkColumn), allColumns, parsed); // FOR UPDATE locks nothing when the row does not exist; concurrent first writers can // both observe an empty before-image while only one reports is_insert=true. - if (beforeDocument.isEmpty() && !isInsert) { + if (readBefore && before == null && !isInsert) { connection.rollback(); continue; } - final Document afterDocument = - selectFirstDocument(connection, keyQuery) - .orElseThrow( - () -> - new IOException( - "Could not read document after createOrReplace for key: " + key)); + final Document after = + readAfter + ? selectFirstDocument(connection, keyQuery) + .orElseThrow( + () -> + new IOException( + "Could not read document after " + + operation + + " for key: " + + key)) + : null; connection.commit(); - - if (beforeDocument.isEmpty()) { - return DocumentChange.created(afterDocument); - } - return DocumentChange.updated(beforeDocument.get(), afterDocument); - } catch (PSQLException e) { - connection.rollback(); - throw e; + return DocumentChange.of(before, after); } catch (Exception e) { connection.rollback(); - if (e instanceof IOException) { - throw (IOException) e; - } - throw new IOException(e); + throw e; } } } @@ -334,20 +334,19 @@ private DocumentChange createOrReplaceAndReturnChangeWithRetry( } catch (PSQLException e) { if (!isRetry && shouldRefreshSchemaAndRetry(e.getSQLState())) { LOGGER.info( - "Schema mismatch detected during createOrReplaceAndReturnChange (SQLState: {}), refreshing schema and retrying. key: {}", + "Schema mismatch detected during {} (SQLState: {}), refreshing schema and retrying. key: {}", + operation, e.getSQLState(), key); schemaRegistry.invalidate(tableName); - return createOrReplaceAndReturnChangeWithRetry(key, document, true); + return writeAndReturnChange(key, document, options, sqlBuilder, operation, true); } - LOGGER.error( - "SQLException in createOrReplaceAndReturnChange. key: {} content: {}", key, document, e); + LOGGER.error("SQLException in {}. key: {} content: {}", operation, key, document, e); throw new IOException(e); } catch (IOException e) { throw e; } catch (Exception e) { - LOGGER.error( - "Exception in createOrReplaceAndReturnChange. key: {} content: {}", key, document, e); + LOGGER.error("Exception in {}. key: {} content: {}", operation, key, document, e); throw new IOException(e); } } @@ -1867,53 +1866,8 @@ private String buildCreateOrReplaceSql(List allTableColumns, String pkCo tableIdentifier, columnList, placeholders, pkColumn, setClause); } - private void upsertOnConnection( - Connection connection, Key key, Document document, String tableName) throws IOException { - final List skippedFields = new ArrayList<>(); - - try { - final TypedDocument parsed = parseDocument(document, tableName, skippedFields); - final String pkColumn = getPKForTable(tableName); - final String quotedPkColumn = PostgresUtils.wrapFieldNamesWithDoubleQuotes(pkColumn); - final PostgresDataType pkType = getPrimaryKeyType(tableName, pkColumn); - parsed.add(quotedPkColumn, key.toString(), pkType, false); - - final List allColumns = getAllQuotedColumns(tableName); - final String sql = buildUpsertSql(allColumns, quotedPkColumn); - executeUpsertOnConnection(connection, sql, allColumns, parsed); - } catch (SQLException e) { - LOGGER.error("SQLException in upsert. key: {} content: {}", key, document, e); - throw new IOException(e); - } - } - - private void executeUpsertOnConnection( - Connection connection, String sql, List allColumns, TypedDocument parsed) - throws SQLException { - try (PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(connection, sql)) { - int index = 1; - final Set parsedColumns = new HashSet<>(parsed.getColumns()); - for (String column : allColumns) { - if (parsedColumns.contains(column)) { - setParameter( - connection, - ps, - index++, - parsed.getValue(column), - parsed.getType(column), - parsed.isArray(column)); - } else { - ps.setObject(index++, null); - } - } - try (ResultSet rs = ps.executeQuery()) { - rs.next(); - } - } - } - /** Returns true if INSERT, false if UPDATE. */ - private boolean executeCreateOrReplaceReturningIsInsertOnConnection( + private boolean executeReturningIsInsertOnConnection( Connection connection, String sql, List allColumns, TypedDocument parsed) throws SQLException { try (PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(connection, sql)) { 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 8c43890d2..686228aab 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 @@ -299,47 +299,6 @@ public Document upsertAndReturn(Key key, Document document) throws IOException { } } - @Override - public Optional upsertAndReturnOlderDocument(Key key, Document document) - throws IOException { - final String selectQuery = - String.format("SELECT * FROM %s WHERE %s = ? FOR UPDATE", tableIdentifier, ID); - - try (final Connection connection = client.getTransactionalConnection()) { - try { - Optional beforeDocument = empty(); - try (final PreparedStatement selectStatement = connection.prepareStatement(selectQuery)) { - selectStatement.setString(1, key.toString()); - beforeDocument = - getFirstDocument(new PostgresResultIterator(selectStatement.executeQuery())); - } - - try (final PreparedStatement upsertStatement = - connection.prepareStatement(getUpsertSQL(), Statement.RETURN_GENERATED_KEYS)) { - // createdTime must come from this locked row; a second connection would deadlock. - final String jsonString = prepareDocument(key, document, createdTimeFrom(beforeDocument)); - upsertStatement.setString(1, key.toString()); - upsertStatement.setString(2, jsonString); - upsertStatement.setString(3, jsonString); - upsertStatement.executeUpdate(); - } - - connection.commit(); - return beforeDocument; - } catch (final Exception e) { - connection.rollback(); - throw e; - } - } catch (final Exception e) { - LOGGER.error( - "Exception upserting document and returning older version. key: {} content:{}", - key, - document, - e); - throw new IOException(e); - } - } - /** * Update the sub document based on subDocPath based on longest key match. * @@ -1190,11 +1149,6 @@ private long bulkUpdateRequestsWithoutFilter(List requestsWit } private String prepareDocument(Key key, Document document) throws IOException { - return prepareDocument(key, document, getCreatedTime(key)); - } - - private String prepareDocument(Key key, Document document, Optional existingCreatedTime) - throws IOException { String jsonString = document.toJson(); ObjectNode jsonNode = (ObjectNode) MAPPER.readTree(jsonString); @@ -1202,23 +1156,13 @@ private String prepareDocument(Key key, Document document, Optional existi // update time fields long now = System.currentTimeMillis(); + Optional existingCreatedTime = getCreatedTime(key); jsonNode.put(DocStoreConstants.CREATED_TIME, existingCreatedTime.orElse(now)); jsonNode.put(DocStoreConstants.LAST_UPDATED_TIME, now); return MAPPER.writeValueAsString(jsonNode); } - private Optional createdTimeFrom(Optional document) throws IOException { - if (document.isEmpty()) { - return empty(); - } - JsonNode existingDocument = getDocAsJSON(document.get()); - if (existingDocument.has(DocStoreConstants.CREATED_TIME)) { - return Optional.of(existingDocument.get(DocStoreConstants.CREATED_TIME).asLong()); - } - return empty(); - } - private String prepareDocumentForCreation(final Key key, final Document document) throws IOException { final String jsonString = document.toJson(); diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/DocumentChangeTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/DocumentChangeTest.java index 6f01e5115..436be3cb8 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/DocumentChangeTest.java +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/DocumentChangeTest.java @@ -2,8 +2,8 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertSame; -import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.IOException; import org.junit.jupiter.api.Test; @@ -11,20 +11,37 @@ class DocumentChangeTest { @Test - void createdHasEmptyBefore() throws IOException { + void bothImagesPresent() throws IOException { + Document before = new JSONDocument("{\"item\":\"before\"}"); Document after = new JSONDocument("{\"item\":\"after\"}"); - DocumentChange change = DocumentChange.created(after); - assertFalse(change.getBefore().isPresent()); - assertSame(after, change.getAfter()); + DocumentChange change = DocumentChange.of(before, after); + assertSame(before, change.getBefore().orElseThrow()); + assertSame(after, change.getAfter().orElseThrow()); } @Test - void updatedHasBothImages() throws IOException { - Document before = new JSONDocument("{\"item\":\"before\"}"); - Document after = new JSONDocument("{\"item\":\"after\"}"); - DocumentChange change = DocumentChange.updated(before, after); - assertTrue(change.getBefore().isPresent()); - assertEquals(before, change.getBefore().get()); - assertSame(after, change.getAfter()); + void nullImagesAreEmpty() throws IOException { + Document document = new JSONDocument("{\"item\":\"doc\"}"); + + DocumentChange afterOnly = DocumentChange.of(null, document); + assertFalse(afterOnly.getBefore().isPresent()); + assertSame(document, afterOnly.getAfter().orElseThrow()); + + DocumentChange beforeOnly = DocumentChange.of(document, null); + assertSame(document, beforeOnly.getBefore().orElseThrow()); + assertFalse(beforeOnly.getAfter().isPresent()); + + DocumentChange none = DocumentChange.of(null, null); + assertFalse(none.getBefore().isPresent()); + assertFalse(none.getAfter().isPresent()); + } + + @Test + void equalityUsesBothImages() throws IOException { + Document document = new JSONDocument("{\"item\":\"doc\"}"); + assertEquals(DocumentChange.of(null, document), DocumentChange.of(null, document)); + assertEquals( + DocumentChange.of(null, document).hashCode(), DocumentChange.of(null, document).hashCode()); + assertNotEquals(DocumentChange.of(null, document), DocumentChange.of(document, null)); } } diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoCollectionTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoCollectionTest.java index 80ebe9404..a55088ee2 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoCollectionTest.java +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoCollectionTest.java @@ -39,6 +39,7 @@ import com.mongodb.client.result.UpdateResult; import java.io.IOException; import java.time.Clock; +import java.util.Date; import java.util.List; import java.util.Optional; import java.util.concurrent.TimeUnit; @@ -60,6 +61,8 @@ import org.hypertrace.core.documentstore.model.config.ConnectionConfig; import org.hypertrace.core.documentstore.model.options.UpdateOptions; import org.hypertrace.core.documentstore.model.options.UpdateOptions.MissingDocumentStrategy; +import org.hypertrace.core.documentstore.model.options.UpsertOptions; +import org.hypertrace.core.documentstore.model.options.UpsertOptions.ReturnDocuments; import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate; import org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue; import org.hypertrace.core.documentstore.query.SortingSpec; @@ -457,98 +460,173 @@ void testBulkUpdateWithoutUpdates() { } @Nested - class UpsertAndReturnOlderDocumentTest { + class UpsertAndReturnWithOptionsTest { private final Key key = new SingleValueKey("default", "testKey"); private final Document document; + private final ArgumentCaptor update = + ArgumentCaptor.forClass(BasicDBObject.class); + private final ArgumentCaptor options = + ArgumentCaptor.forClass(FindOneAndUpdateOptions.class); - UpsertAndReturnOlderDocumentTest() throws IOException { + UpsertAndReturnWithOptionsTest() throws IOException { document = new JSONDocument("{\"item\":\"Soap\",\"price\":100}"); } - @Test - void testUpsertAndReturnOlderDocumentReturnsEmptyWhenCreated() throws IOException { - final ArgumentCaptor options = - ArgumentCaptor.forClass(FindOneAndUpdateOptions.class); - + private void stubFindOneAndUpdate(final BasicDBObject result) { when(collection.findOneAndUpdate( - any(BasicDBObject.class), any(BasicDBObject.class), options.capture())) - .thenReturn(null); + any(BasicDBObject.class), update.capture(), options.capture())) + .thenReturn(result); + } - final Optional result = mongoCollection.upsertAndReturnOlderDocument(key, document); + private UpsertOptions optionsFor(final ReturnDocuments returnDocuments) { + return UpsertOptions.builder().returnDocuments(returnDocuments).build(); + } - assertFalse(result.isPresent()); - assertTrue(options.getValue().isUpsert()); - assertEquals(BEFORE, options.getValue().getReturnDocument()); + private BasicDBObject existingDocument() { + return new BasicDBObject("_id", key.toString()) + .append("item", "OldSoap") + .append("price", 50) + .append("quantity", 3) + .append("createdTime", 123L) + .append("lastUpdatedTime", 456L); } - @Test - void testUpsertAndReturnOlderDocumentReturnsBeforeImage() throws IOException { - final ArgumentCaptor options = - ArgumentCaptor.forClass(FindOneAndUpdateOptions.class); - final BasicDBObject beforeDocument = new BasicDBObject("item", "OldSoap"); + /** Asserts the update uses a literal _lastUpdateTime and returns its $set clause. */ + private BasicDBObject assertExplicitLastUpdateTime() { + final BasicDBObject updateDocument = update.getValue(); + assertFalse(updateDocument.containsField("$currentDate")); + final BasicDBObject set = (BasicDBObject) updateDocument.get("$set"); + final Object lastUpdateTime = set.get("_lastUpdateTime"); + assertTrue(lastUpdateTime instanceof Date); + assertEquals(set.getLong("lastUpdatedTime"), ((Date) lastUpdateTime).getTime()); + assertEquals( + set.getLong("lastUpdatedTime"), + ((BasicDBObject) updateDocument.get("$setOnInsert")).getLong("createdTime")); + return set; + } - when(collection.findOneAndUpdate( - any(BasicDBObject.class), any(BasicDBObject.class), options.capture())) - .thenReturn(beforeDocument); + @Test + void testBeforeOnInsertReturnsNothing() throws IOException { + stubFindOneAndUpdate(null); - final Optional result = mongoCollection.upsertAndReturnOlderDocument(key, document); + final DocumentChange result = + mongoCollection.upsertAndReturn(key, document, optionsFor(ReturnDocuments.BEFORE)); - assertTrue(result.isPresent()); - assertTrue(result.get().toJson().contains("OldSoap")); + assertFalse(result.getBefore().isPresent()); + assertFalse(result.getAfter().isPresent()); assertTrue(options.getValue().isUpsert()); assertEquals(BEFORE, options.getValue().getReturnDocument()); + assertExplicitLastUpdateTime(); } - } - @Nested - class UpsertAndReturnChangeTest { + @Test + void testBeforeOnUpdateReturnsBeforeImage() throws IOException { + stubFindOneAndUpdate(existingDocument()); - private final Key key = new SingleValueKey("default", "testKey"); - private final Document document; + final DocumentChange result = + mongoCollection.upsertAndReturn(key, document, optionsFor(ReturnDocuments.BEFORE)); - UpsertAndReturnChangeTest() throws IOException { - document = new JSONDocument("{\"item\":\"Soap\",\"price\":100}"); + assertEquals( + MongoUtils.dbObjectToDocument(existingDocument()), result.getBefore().orElseThrow()); + assertFalse(result.getAfter().isPresent()); + assertEquals(BEFORE, options.getValue().getReturnDocument()); } @Test - void testUpsertAndReturnChangeCreated() throws IOException { - final ArgumentCaptor options = - ArgumentCaptor.forClass(FindOneAndUpdateOptions.class); - - when(collection.findOneAndUpdate( - any(BasicDBObject.class), any(BasicDBObject.class), options.capture())) - .thenReturn(null); + void testAfterReturnsServerAfterImage() throws IOException { + final BasicDBObject serverAfter = + new BasicDBObject("_id", key.toString()).append("item", "Soap").append("price", 100); + stubFindOneAndUpdate(serverAfter); - final DocumentChange result = mongoCollection.upsertAndReturnChange(key, document); + final DocumentChange result = + mongoCollection.upsertAndReturn(key, document, optionsFor(ReturnDocuments.AFTER)); assertFalse(result.getBefore().isPresent()); - assertTrue(result.getAfter().toJson().contains("Soap")); - assertTrue(result.getAfter().toJson().contains("100")); + assertEquals( + MongoUtils.dbObjectToDocument(new BasicDBObject("item", "Soap").append("price", 100)), + result.getAfter().orElseThrow()); assertTrue(options.getValue().isUpsert()); + assertEquals(AFTER, options.getValue().getReturnDocument()); + assertExplicitLastUpdateTime(); + } + + @Test + void testAfterThrowsWhenServerReturnsNothing() { + stubFindOneAndUpdate(null); + assertThrows( + IOException.class, + () -> mongoCollection.upsertAndReturn(key, document, optionsFor(ReturnDocuments.AFTER))); + } + + @Test + void testBothOnInsertDerivesAfterImage() throws IOException { + stubFindOneAndUpdate(null); + + final DocumentChange result = + mongoCollection.upsertAndReturn(key, document, optionsFor(ReturnDocuments.BOTH)); + assertEquals(BEFORE, options.getValue().getReturnDocument()); + final BasicDBObject set = assertExplicitLastUpdateTime(); + assertFalse(result.getBefore().isPresent()); + final BasicDBObject expectedAfter = + new BasicDBObject("item", "Soap") + .append("price", 100) + .append("lastUpdatedTime", set.getLong("lastUpdatedTime")) + .append("_lastUpdateTime", set.get("_lastUpdateTime")) + .append("createdTime", set.getLong("lastUpdatedTime")); + assertEquals(MongoUtils.dbObjectToDocument(expectedAfter), result.getAfter().orElseThrow()); } @Test - void testUpsertAndReturnChangeUpdated() throws IOException { - final ArgumentCaptor options = - ArgumentCaptor.forClass(FindOneAndUpdateOptions.class); - final BasicDBObject beforeDocument = - new BasicDBObject("item", "OldSoap").append("price", 50).append("quantity", 3); + void testBothOnUpdateDerivesAfterImage() throws IOException { + stubFindOneAndUpdate(existingDocument()); - when(collection.findOneAndUpdate( - any(BasicDBObject.class), any(BasicDBObject.class), options.capture())) - .thenReturn(beforeDocument); + final DocumentChange result = + mongoCollection.upsertAndReturn(key, document, optionsFor(ReturnDocuments.BOTH)); - final DocumentChange result = mongoCollection.upsertAndReturnChange(key, document); + assertEquals(BEFORE, options.getValue().getReturnDocument()); + final BasicDBObject set = assertExplicitLastUpdateTime(); + assertEquals( + MongoUtils.dbObjectToDocument(existingDocument()), result.getBefore().orElseThrow()); + final BasicDBObject expectedAfter = + new BasicDBObject("item", "Soap") + .append("price", 100) + .append("quantity", 3) + .append("createdTime", 123L) + .append("lastUpdatedTime", set.getLong("lastUpdatedTime")) + .append("_lastUpdateTime", set.get("_lastUpdateTime")); + assertEquals(MongoUtils.dbObjectToDocument(expectedAfter), result.getAfter().orElseThrow()); + } - assertTrue(result.getBefore().isPresent()); - assertTrue(result.getBefore().get().toJson().contains("OldSoap")); - assertTrue(result.getAfter().toJson().contains("Soap")); - assertTrue(result.getAfter().toJson().contains("100")); - assertTrue(result.getAfter().toJson().contains("quantity")); - assertTrue(options.getValue().isUpsert()); + @Test + void testCreateOrReplaceBothIsUnsupported() { + assertThrows( + UnsupportedOperationException.class, + () -> + mongoCollection.createOrReplaceAndReturn( + key, document, optionsFor(ReturnDocuments.BOTH))); + } + + @Test + void testCreateOrReplaceBeforeAndAfter() throws IOException { + when(collection.findOneAndUpdate(any(BasicDBObject.class), anyList(), options.capture())) + .thenReturn(existingDocument()); + + final DocumentChange before = + mongoCollection.createOrReplaceAndReturn( + key, document, optionsFor(ReturnDocuments.BEFORE)); assertEquals(BEFORE, options.getValue().getReturnDocument()); + assertTrue(options.getValue().isUpsert()); + assertTrue(before.getBefore().isPresent()); + assertFalse(before.getAfter().isPresent()); + + final DocumentChange after = + mongoCollection.createOrReplaceAndReturn( + key, document, optionsFor(ReturnDocuments.AFTER)); + assertEquals(AFTER, options.getValue().getReturnDocument()); + assertFalse(after.getBefore().isPresent()); + assertTrue(after.getAfter().isPresent()); } } } From b96ad4a0bf44dba623fe4292927fa7e63c2108e6 Mon Sep 17 00:00:00 2001 From: bhubam Date: Wed, 14 Oct 2026 18:15:15 +0530 Subject: [PATCH 6/6] AAP-13921: Address review comments on UpsertOptions write-and-return - Postgres: replace first-insert retry loop with pg_advisory_xact_lock on the key, take the after-image from INSERT ... RETURNING, strip _implicit_id from images - Mongo: BEFORE/AFTER use $currentDate like upsert(); explicit _lastUpdateTime only for BOTH - Mongo: support BOTH in createOrReplaceAndReturn by deriving the after-image from a fully $literal-wrapped replace - Document non-atomicity on Collection Co-authored-by: Cursor --- .../MongoPostgresWriteConsistencyTest.java | 78 ++++---- .../core/documentstore/Collection.java | 8 +- .../documentstore/mongo/MongoCollection.java | 117 ++++++++---- .../postgres/FlatPostgresCollection.java | 172 ++++++++++++------ .../mongo/MongoCollectionTest.java | 163 ++++++++++++----- 5 files changed, 368 insertions(+), 170 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 57e505155..2d7b18165 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 @@ -147,10 +147,6 @@ private enum AtomicWrite { this.merge = merge; } - boolean supportsBoth() { - return this != MONGO_CREATE_OR_REPLACE; - } - String createdTimeField() { return MONGO_STORE.equals(storeName) ? "createdTime" : "created_at"; } @@ -161,19 +157,15 @@ private static class AtomicWriteAndReturnDocumentsProvider implements ArgumentsP @Override public Stream provideArguments(final ExtensionContext context) { return Stream.of(AtomicWrite.values()) - .flatMap( - write -> - Stream.of(ReturnDocuments.values()) - .filter(rd -> rd != ReturnDocuments.BOTH || write.supportsBoth()) - .map(rd -> Arguments.of(write, rd))); + .flatMap(write -> Stream.of(ReturnDocuments.values()).map(rd -> Arguments.of(write, rd))); } } - private static class BothCapableAtomicWritesProvider implements ArgumentsProvider { + private static class AtomicWritesProvider implements ArgumentsProvider { @Override public Stream provideArguments(final ExtensionContext context) { - return Stream.of(AtomicWrite.values()).filter(AtomicWrite::supportsBoth).map(Arguments::of); + return Stream.of(AtomicWrite.values()).map(Arguments::of); } } @@ -604,13 +596,6 @@ private Document readBack(Collection collection, String docId) throws IOExceptio } } - /** Flat Postgres images are read with SELECT ... FOR UPDATE, which also selects this alias. */ - private JsonNode withoutImplicitId(JsonNode json) { - ObjectNode copy = json.deepCopy(); - copy.remove("_implicit_id"); - return copy; - } - private Document partialUpdate(String docId) { ObjectNode updateNode = OBJECT_MAPPER.createObjectNode(); updateNode.put("id", getKeyString(docId)); @@ -642,7 +627,7 @@ void testInsert(AtomicWrite write, ReturnDocuments returnDocuments) throws Excep assertEquals("TestItem", afterJson.get("item").asText()); assertEquals(100, afterJson.get("price").asInt()); assertEquals(50, afterJson.get("quantity").asInt()); - assertEquals(storedJson, withoutImplicitId(afterJson)); + assertEquals(storedJson, afterJson); } } @@ -664,7 +649,7 @@ void testUpdate(AtomicWrite write, ReturnDocuments returnDocuments) throws Excep assertTrue(change.getBefore().isEmpty()); } else { JsonNode beforeJson = OBJECT_MAPPER.readTree(change.getBefore().orElseThrow().toJson()); - assertEquals(initialJson, withoutImplicitId(beforeJson)); + assertEquals(initialJson, beforeJson); assertEquals("TestItem", beforeJson.get("item").asText()); assertEquals(100, beforeJson.get("price").asInt()); } @@ -686,7 +671,7 @@ void testUpdate(AtomicWrite write, ReturnDocuments returnDocuments) throws Excep } else { assertTrue(quantity == null || quantity.isNull()); } - assertEquals(storedJson, withoutImplicitId(afterJson)); + assertEquals(storedJson, afterJson); } } @@ -728,21 +713,48 @@ private void assertTimestampFieldsPresent(Document document) throws IOException } @Test - @DisplayName("Mongo createOrReplace BOTH is unsupported") - void testMongoCreateOrReplaceBothUnsupported() { - String docId = generateDocId("atomic-mongo-replace-both"); - assertThrows( - UnsupportedOperationException.class, - () -> - atomicWrite( - AtomicWrite.MONGO_CREATE_OR_REPLACE, - createKey(docId), - createTestDocument(docId), - ReturnDocuments.BOTH)); + @DisplayName("Mongo createOrReplace BOTH: derived after-image equals the stored document") + void testMongoCreateOrReplaceBothAfterImageIsExact() throws Exception { + String docId = generateDocId("atomic-mongo-replace-exact"); + Key key = createKey(docId); + Collection collection = getCollection(MONGO_STORE); + + DocumentChange inserted = + atomicWrite( + AtomicWrite.MONGO_CREATE_OR_REPLACE, + key, + createTestDocument(docId), + ReturnDocuments.BOTH); + assertTrue(inserted.getBefore().isEmpty()); + Document storedAfterInsert = readBack(collection, docId); + assertEquals(storedAfterInsert, inserted.getAfter().orElseThrow()); + assertTimestampFieldsPresent(storedAfterInsert); + + Thread.sleep(5); + ObjectNode replacement = OBJECT_MAPPER.createObjectNode(); + replacement.put("id", getKeyString(docId)); + replacement.put("item", "$price"); + replacement.putObject("props"); + DocumentChange replaced = + atomicWrite( + AtomicWrite.MONGO_CREATE_OR_REPLACE, + key, + new JSONDocument(replacement), + ReturnDocuments.BOTH); + assertEquals(storedAfterInsert, replaced.getBefore().orElseThrow()); + Document storedAfterReplace = readBack(collection, docId); + assertEquals(storedAfterReplace, replaced.getAfter().orElseThrow()); + + JsonNode insertJson = OBJECT_MAPPER.readTree(storedAfterInsert.toJson()); + JsonNode replaceJson = OBJECT_MAPPER.readTree(storedAfterReplace.toJson()); + assertEquals("$price", replaceJson.get("item").asText()); + assertTrue(replaceJson.get("props").isObject() && replaceJson.get("props").isEmpty()); + assertNull(replaceJson.get("price")); + assertEquals(insertJson.get("createdTime"), replaceJson.get("createdTime")); } @ParameterizedTest(name = "{0}: concurrent BOTH writes form a single chain") - @ArgumentsSource(BothCapableAtomicWritesProvider.class) + @ArgumentsSource(AtomicWritesProvider.class) void testConcurrentWritesFormChain(AtomicWrite write) throws Exception { String docId = generateDocId("atomic-concurrent"); Key key = createKey(docId); 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 7e02f24e1..fc67595cc 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 @@ -14,7 +14,13 @@ import org.hypertrace.core.documentstore.model.options.UpsertOptions; import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate; -/** Interface spec for common operations on a collection of documents */ +/** + * Interface spec for common operations on a collection of documents. + * + *

Methods are not atomic, and give no isolation from concurrent writers, unless their javadoc + * explicitly says so. In particular, reading a document before or after a write with separate calls + * can observe other writers' changes (or stale replicas, depending on read preference). + */ public interface Collection { /** 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 6499d697f..f71593f7c 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 @@ -308,11 +308,13 @@ public Document upsertAndReturn(Key key, Document document) throws IOException { * Performs one atomic server-side write (a single {@code findOneAndUpdate} with upsert). * *

    - *
  • BEFORE / AFTER: the server returns the requested image from that write. + *
  • BEFORE / AFTER: the server returns the requested image from that write; the update is the + * same as {@link #upsert(Key, Document)}. *
  • BOTH: the server returns the before-image from that write; the after-image is computed by * applying that same update to the before-image (no second read). This is exact because the - * update only sets top-level fields to literal values, including {@code _lastUpdateTime}, - * which is set from the client clock rather than {@code $currentDate}. + * update only sets top-level fields to literal values. For that reason {@code + * _lastUpdateTime} is set from the client clock instead of {@code $currentDate} in this + * mode only. *
*/ @Override @@ -320,10 +322,12 @@ public DocumentChange upsertAndReturn(Key key, Document document, UpsertOptions throws IOException { final ReturnDocuments returnDocuments = options.getReturnDocuments(); final BasicDBObject update = - prepareUpsertWithExplicitLastUpdateTime(key, document, System.currentTimeMillis()); + returnDocuments == ReturnDocuments.BOTH + ? prepareUpsertWithExplicitLastUpdateTime(key, document, System.currentTimeMillis()) + : prepareUpsert(key, document); final ReturnDocument serverImage = returnDocuments == ReturnDocuments.AFTER ? ReturnDocument.AFTER : ReturnDocument.BEFORE; - final BasicDBObject result = + final BasicDBObject serverResult = Failsafe.with(upsertRetryPolicy) .get( () -> @@ -334,58 +338,109 @@ public DocumentChange upsertAndReturn(Key key, Document document, UpsertOptions switch (returnDocuments) { case BEFORE: - return DocumentChange.of(toDocumentOrNull(result), null); + return DocumentChange.of(toDocumentOrNull(serverResult), null); case AFTER: - if (result == null) { + if (serverResult == null) { throw new IOException("Could not upsert the document with key: " + key); } - return DocumentChange.of(null, dbObjectToDocument(result)); + return DocumentChange.of(null, dbObjectToDocument(serverResult)); case BOTH: - final BasicDBObject after = applyUpsertToBeforeImage(key, result, update); - return DocumentChange.of(toDocumentOrNull(result), dbObjectToDocument(after)); + final BasicDBObject serverBefore = serverResult; + final BasicDBObject derivedAfter = applyUpsertToBeforeImage(key, serverBefore, update); + return DocumentChange.of(toDocumentOrNull(serverBefore), dbObjectToDocument(derivedAfter)); default: throw new UnsupportedOperationException("Unhandled return documents: " + returnDocuments); } } /** - * Performs one atomic server-side write (a single {@code findOneAndUpdate} with upsert) and - * returns the BEFORE or AFTER image from that write. + * Performs one atomic server-side write (a single {@code findOneAndUpdate} with upsert). * - *

BOTH is not supported: the replace pipeline evaluates document values as aggregation - * expressions (e.g. a string value starting with {@code $} is read as a field path), so the - * stored document cannot always be derived exactly from the before-image without a second read. + *

    + *
  • BEFORE / AFTER: the server returns the requested image from that write. + *
  • BOTH: the server returns the before-image from that write; the after-image is computed by + * applying the same replace to the before-image (no second read). + *
+ * + *

Unlike {@link #createOrReplace(Key, Document)}, every top-level value is wrapped in {@code + * $literal}, so values are stored exactly as given (e.g. a string starting with {@code $} is not + * read as a field path). That is what makes the derived after-image exact. */ @Override public DocumentChange createOrReplaceAndReturn( final Key key, final Document document, final UpsertOptions options) throws IOException { final ReturnDocuments returnDocuments = options.getReturnDocuments(); - if (returnDocuments == ReturnDocuments.BOTH) { - throw new UnsupportedOperationException( - "createOrReplaceAndReturn with ReturnDocuments.BOTH is not supported on Mongo"); - } - final boolean returnBefore = returnDocuments == ReturnDocuments.BEFORE; + final long now = System.currentTimeMillis(); + final BasicDBObject replacement = prepareDocumentForLiteralReplace(key, document, now); + final ReturnDocument serverImage = + returnDocuments == ReturnDocuments.AFTER ? ReturnDocument.AFTER : ReturnDocument.BEFORE; try { - final FindOneAndUpdateOptions findOptions = - new FindOneAndUpdateOptions() - .upsert(true) - .returnDocument(returnBefore ? ReturnDocument.BEFORE : ReturnDocument.AFTER); - final BasicDBObject result = + final BasicDBObject serverResult = collection.findOneAndUpdate( this.selectionCriteriaForKey(key), - this.prepareForCreateOrReplace(key, document), - findOptions); - LOGGER.debug("Create or replace result: {}", result); - return returnBefore - ? DocumentChange.of(toDocumentOrNull(result), null) - : DocumentChange.of(null, dbObjectToDocument(requireNonNull(result))); + prepareLiteralCreateOrReplace(replacement, now), + new FindOneAndUpdateOptions().upsert(true).returnDocument(serverImage)); + LOGGER.debug("Create or replace result: {}", serverResult); + + switch (returnDocuments) { + case BEFORE: + return DocumentChange.of(toDocumentOrNull(serverResult), null); + case AFTER: + return DocumentChange.of(null, dbObjectToDocument(requireNonNull(serverResult))); + case BOTH: + final BasicDBObject serverBefore = serverResult; + final BasicDBObject derivedAfter = + applyReplaceToBeforeImage(serverBefore, replacement, now); + return DocumentChange.of( + toDocumentOrNull(serverBefore), dbObjectToDocument(derivedAfter)); + default: + throw new UnsupportedOperationException("Unhandled return documents: " + returnDocuments); + } } catch (final Exception e) { LOGGER.error("Exception creating/replacing document. key: {} content:{}", key, document, e); throw e; } } + private BasicDBObject prepareDocumentForLiteralReplace(Key key, Document document, long now) + throws JsonProcessingException { + final BasicDBObject replacement = prepareDocument(key, document, now); + replacement.put(LAST_UPDATE_TIMESTAMP_ISO_8601, new Date(now)); + return replacement; + } + + /** Same stages as {@link #prepareForCreateOrReplace}, with every value wrapped in $literal. */ + private List prepareLiteralCreateOrReplace(BasicDBObject replacement, long now) { + final BasicDBObject project = + new BasicDBObject( + PROJECT_CLAUSE, + new BasicDBObject( + CREATED_TIME, + new BasicDBObject( + IF_NULL_CLAUSE, + new BsonValue[] {new BsonString(PREFIX + CREATED_TIME), new BsonInt64(now)}))); + final BasicDBObject literalSet = new BasicDBObject(); + replacement.forEach( + (field, value) -> literalSet.put(field, new BasicDBObject(PREFIX + "literal", value))); + return List.of(project, new BasicDBObject(SET_CLAUSE, literalSet)); + } + + /** + * Applies {@link #prepareLiteralCreateOrReplace} to the before-image the server returned for that + * same write: {@code $project} keeps only {@code _id} and {@code createdTime} (defaulting to + * {@code now} when missing or null), then {@code $set} adds the replacement fields. + */ + private static BasicDBObject applyReplaceToBeforeImage( + @Nullable BasicDBObject before, BasicDBObject replacement, long now) { + final BasicDBObject after = + new BasicDBObject(ID_KEY, before == null ? replacement.get(ID_KEY) : before.get(ID_KEY)); + final Object createdTime = before == null ? null : before.get(CREATED_TIME); + after.put(CREATED_TIME, createdTime == null ? now : createdTime); + after.putAll(replacement.toMap()); + return after; + } + @Nullable private static Document toDocumentOrNull(@Nullable final BasicDBObject dbObject) { return dbObject == null ? null : dbObjectToDocument(dbObject); 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 5b4b600a5..3109e0edc 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 @@ -1,6 +1,7 @@ package org.hypertrace.core.documentstore.postgres; import static java.util.Map.entry; +import static org.hypertrace.core.documentstore.commons.DocStoreConstants.IMPLICIT_ID; import static org.hypertrace.core.documentstore.model.options.ReturnDocumentType.AFTER_UPDATE; import static org.hypertrace.core.documentstore.model.options.ReturnDocumentType.BEFORE_UPDATE; import static org.hypertrace.core.documentstore.model.subdoc.UpdateOperator.ADD; @@ -10,8 +11,10 @@ import static org.hypertrace.core.documentstore.model.subdoc.UpdateOperator.SET; import static org.hypertrace.core.documentstore.model.subdoc.UpdateOperator.UNSET; +import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; import java.io.IOException; @@ -53,6 +56,7 @@ import org.hypertrace.core.documentstore.DocumentChange; import org.hypertrace.core.documentstore.DocumentType; import org.hypertrace.core.documentstore.Filter; +import org.hypertrace.core.documentstore.JSONDocument; import org.hypertrace.core.documentstore.Key; import org.hypertrace.core.documentstore.UpdateResult; import org.hypertrace.core.documentstore.commons.CommonUpdateValidator; @@ -104,7 +108,8 @@ public class FlatPostgresCollection extends PostgresCollection { private static final Logger LOGGER = LoggerFactory.getLogger(FlatPostgresCollection.class); private static final ObjectMapper MAPPER = new ObjectMapper(); - private static final int MAX_FIRST_INSERT_RACE_RETRIES = 3; + private static final String IS_INSERT_COLUMN = "is_insert"; + private static final String RETURNING_ROW_SUFFIX = ", *"; private static final String WRITE_NOT_SUPPORTED = "Write operations are not supported for flat collections yet!"; private static final String MISSING_COLUMN_STRATEGY_CONFIG = "missingColumnStrategy"; @@ -263,10 +268,18 @@ public DocumentChange createOrReplaceAndReturn(Key key, Document document, Upser } /** - * Executes a single-row {@code INSERT ... ON CONFLICT DO UPDATE ... RETURNING is_insert} built by - * {@code sqlBuilder} and reads the requested images in one transaction on one connection. The - * before-image is read with {@code SELECT ... FOR UPDATE}, so no other writer can change the row - * between that read, the write and the after-image read. + * Executes a single-row {@code INSERT ... ON CONFLICT DO UPDATE} built by {@code sqlBuilder} in + * one transaction on one connection and returns the requested images. + * + *

    + *
  • Before-image: {@code SELECT ... FOR UPDATE}, so no other writer can change the row + * between that read and the write. When the row does not exist, FOR UPDATE locks nothing, + * so a transaction-scoped advisory lock on the key is taken and the row is re-read; this + * serialises concurrent first writers that go through this method. + *
  • After-image: {@code RETURNING *} of the write itself (same columns as {@code find()}). + *
+ * + *

Statements per call: AFTER = 1; BEFORE / BOTH = 2 when the row exists, 4 on first insert. */ private DocumentChange writeAndReturnChange( Key key, @@ -280,57 +293,38 @@ private DocumentChange writeAndReturnChange( final boolean readBefore = returnDocuments != ReturnDocuments.AFTER; final boolean readAfter = returnDocuments != ReturnDocuments.BEFORE; final String tableName = tableIdentifier.getTableName(); - final Query keyQuery = buildPrimaryKeyQuery(tableName, key); - - try { - for (int attempt = 0; attempt < MAX_FIRST_INSERT_RACE_RETRIES; attempt++) { - try (Connection connection = client.getTransactionalConnection()) { - try { - final Document before = - readBefore ? selectFirstDocument(connection, keyQuery).orElse(null) : null; - - final TypedDocument parsed = parseDocument(document, tableName, new ArrayList<>()); - final String pkColumn = getPKForTable(tableName); - final String quotedPkColumn = PostgresUtils.wrapFieldNamesWithDoubleQuotes(pkColumn); - parsed.add( - quotedPkColumn, key.toString(), getPrimaryKeyType(tableName, pkColumn), false); - final List allColumns = getAllQuotedColumns(tableName); - final boolean isInsert = - executeReturningIsInsertOnConnection( - connection, sqlBuilder.apply(allColumns, quotedPkColumn), allColumns, parsed); - - // FOR UPDATE locks nothing when the row does not exist; concurrent first writers can - // both observe an empty before-image while only one reports is_insert=true. - if (readBefore && before == null && !isInsert) { - connection.rollback(); - continue; - } - final Document after = - readAfter - ? selectFirstDocument(connection, keyQuery) - .orElseThrow( - () -> - new IOException( - "Could not read document after " - + operation - + " for key: " - + key)) - : null; - connection.commit(); - return DocumentChange.of(before, after); - } catch (Exception e) { - connection.rollback(); - throw e; - } + try (Connection connection = client.getTransactionalConnection()) { + try { + final Document before = + readBefore + ? selectBeforeImageLockingKey(connection, tableName, key).orElse(null) + : null; + + final TypedDocument parsed = parseDocument(document, tableName, new ArrayList<>()); + final String pkColumn = getPKForTable(tableName); + final String quotedPkColumn = PostgresUtils.wrapFieldNamesWithDoubleQuotes(pkColumn); + parsed.add(quotedPkColumn, key.toString(), getPrimaryKeyType(tableName, pkColumn), false); + final List allColumns = getAllQuotedColumns(tableName); + final String sql = sqlBuilder.apply(allColumns, quotedPkColumn); + final WriteResult result = + executeWriteOnConnection( + connection, readAfter ? sql + RETURNING_ROW_SUFFIX : sql, allColumns, parsed); + + // Only possible when a writer that does not take the advisory lock (e.g. upsert()) + // inserted the row after our read; the true before-image is unknown, so nothing is + // written and the caller may retry. + if (readBefore && before == null && !result.isInsert()) { + throw new IOException( + "Concurrent insert detected on key " + key + " during " + operation + "; retry"); } + + connection.commit(); + return DocumentChange.of(before, readAfter ? result.getRow() : null); + } catch (Exception e) { + connection.rollback(); + throw e; } - throw new IOException( - "Concurrent first insert on key " - + key - + " could not be resolved after " - + MAX_FIRST_INSERT_RACE_RETRIES - + " attempts"); } catch (PSQLException e) { if (!isRetry && shouldRefreshSchemaAndRetry(e.getSQLState())) { LOGGER.info( @@ -1866,10 +1860,32 @@ private String buildCreateOrReplaceSql(List allTableColumns, String pkCo tableIdentifier, columnList, placeholders, pkColumn, setClause); } - /** Returns true if INSERT, false if UPDATE. */ - private boolean executeReturningIsInsertOnConnection( + private static final class WriteResult { + private final boolean insert; + private final Document row; + + private WriteResult(boolean insert, Document row) { + this.insert = insert; + this.row = row; + } + + boolean isInsert() { + return insert; + } + + Document getRow() { + return row; + } + } + + /** + * Executes a single-row write ending in {@code RETURNING (xmax = 0) AS is_insert}, optionally + * followed by {@link #RETURNING_ROW_SUFFIX}, in which case the returned row is the stored + * document. + */ + private WriteResult executeWriteOnConnection( Connection connection, String sql, List allColumns, TypedDocument parsed) - throws SQLException { + throws SQLException, IOException { try (PreparedStatement ps = queryExecutor.prepareStatementWithTimeout(connection, sql)) { int index = 1; final Set parsedColumns = new HashSet<>(parsed.getColumns()); @@ -1887,11 +1903,49 @@ private boolean executeReturningIsInsertOnConnection( } } try (ResultSet rs = ps.executeQuery()) { - if (rs.next()) { - return rs.getBoolean("is_insert"); + if (!sql.endsWith(RETURNING_ROW_SUFFIX)) { + return new WriteResult(rs.next() && rs.getBoolean(IS_INSERT_COLUMN), null); } + final PostgresResultIteratorWithBasicTypes iterator = + new PostgresResultIteratorWithBasicTypes(rs, DocumentType.FLAT); + if (!iterator.hasNext()) { + throw new IOException("Write returned no row"); + } + final Document row = iterator.next(); + return new WriteResult( + rs.getBoolean(IS_INSERT_COLUMN), withoutFields(row, IS_INSERT_COLUMN, IMPLICIT_ID)); } - return false; + } + } + + /** + * Reads the row for {@code key} with {@code SELECT ... FOR UPDATE}. If absent, takes a + * transaction-scoped advisory lock on the key and reads again, so concurrent first writers are + * serialised and the later one sees the earlier one's row. + */ + private Optional selectBeforeImageLockingKey( + Connection connection, String tableName, Key key) throws SQLException, IOException { + final Query keyQuery = buildPrimaryKeyQuery(tableName, key); + Optional before = selectFirstDocument(connection, keyQuery); + if (before.isEmpty()) { + try (PreparedStatement ps = + queryExecutor.prepareStatementWithTimeout( + connection, "SELECT pg_advisory_xact_lock(hashtextextended(?, 0))")) { + ps.setString(1, tableIdentifier + SHAPE_KEY_DELIMITER + key); + ps.execute(); + } + before = selectFirstDocument(connection, keyQuery); + } + return before.map(document -> withoutFields(document, IMPLICIT_ID)); + } + + private static Document withoutFields(Document document, String... fields) { + try { + final ObjectNode node = (ObjectNode) MAPPER.readTree(document.toJson()); + node.remove(Arrays.asList(fields)); + return new JSONDocument(node, DocumentType.FLAT); + } catch (JsonProcessingException e) { + throw new IllegalStateException(e); } } diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoCollectionTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoCollectionTest.java index a55088ee2..5b56043f1 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoCollectionTest.java +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/mongo/MongoCollectionTest.java @@ -469,43 +469,14 @@ class UpsertAndReturnWithOptionsTest { private final ArgumentCaptor options = ArgumentCaptor.forClass(FindOneAndUpdateOptions.class); + @SuppressWarnings("unchecked") + private final ArgumentCaptor> pipeline = + ArgumentCaptor.forClass(List.class); + UpsertAndReturnWithOptionsTest() throws IOException { document = new JSONDocument("{\"item\":\"Soap\",\"price\":100}"); } - private void stubFindOneAndUpdate(final BasicDBObject result) { - when(collection.findOneAndUpdate( - any(BasicDBObject.class), update.capture(), options.capture())) - .thenReturn(result); - } - - private UpsertOptions optionsFor(final ReturnDocuments returnDocuments) { - return UpsertOptions.builder().returnDocuments(returnDocuments).build(); - } - - private BasicDBObject existingDocument() { - return new BasicDBObject("_id", key.toString()) - .append("item", "OldSoap") - .append("price", 50) - .append("quantity", 3) - .append("createdTime", 123L) - .append("lastUpdatedTime", 456L); - } - - /** Asserts the update uses a literal _lastUpdateTime and returns its $set clause. */ - private BasicDBObject assertExplicitLastUpdateTime() { - final BasicDBObject updateDocument = update.getValue(); - assertFalse(updateDocument.containsField("$currentDate")); - final BasicDBObject set = (BasicDBObject) updateDocument.get("$set"); - final Object lastUpdateTime = set.get("_lastUpdateTime"); - assertTrue(lastUpdateTime instanceof Date); - assertEquals(set.getLong("lastUpdatedTime"), ((Date) lastUpdateTime).getTime()); - assertEquals( - set.getLong("lastUpdatedTime"), - ((BasicDBObject) updateDocument.get("$setOnInsert")).getLong("createdTime")); - return set; - } - @Test void testBeforeOnInsertReturnsNothing() throws IOException { stubFindOneAndUpdate(null); @@ -517,7 +488,7 @@ void testBeforeOnInsertReturnsNothing() throws IOException { assertFalse(result.getAfter().isPresent()); assertTrue(options.getValue().isUpsert()); assertEquals(BEFORE, options.getValue().getReturnDocument()); - assertExplicitLastUpdateTime(); + assertServerLastUpdateTime(); } @Test @@ -531,6 +502,7 @@ void testBeforeOnUpdateReturnsBeforeImage() throws IOException { MongoUtils.dbObjectToDocument(existingDocument()), result.getBefore().orElseThrow()); assertFalse(result.getAfter().isPresent()); assertEquals(BEFORE, options.getValue().getReturnDocument()); + assertServerLastUpdateTime(); } @Test @@ -548,7 +520,7 @@ void testAfterReturnsServerAfterImage() throws IOException { result.getAfter().orElseThrow()); assertTrue(options.getValue().isUpsert()); assertEquals(AFTER, options.getValue().getReturnDocument()); - assertExplicitLastUpdateTime(); + assertServerLastUpdateTime(); } @Test @@ -599,19 +571,9 @@ void testBothOnUpdateDerivesAfterImage() throws IOException { assertEquals(MongoUtils.dbObjectToDocument(expectedAfter), result.getAfter().orElseThrow()); } - @Test - void testCreateOrReplaceBothIsUnsupported() { - assertThrows( - UnsupportedOperationException.class, - () -> - mongoCollection.createOrReplaceAndReturn( - key, document, optionsFor(ReturnDocuments.BOTH))); - } - @Test void testCreateOrReplaceBeforeAndAfter() throws IOException { - when(collection.findOneAndUpdate(any(BasicDBObject.class), anyList(), options.capture())) - .thenReturn(existingDocument()); + stubFindOneAndReplacePipeline(existingDocument()); final DocumentChange before = mongoCollection.createOrReplaceAndReturn( @@ -628,5 +590,114 @@ void testCreateOrReplaceBeforeAndAfter() throws IOException { assertFalse(after.getBefore().isPresent()); assertTrue(after.getAfter().isPresent()); } + + @Test + void testCreateOrReplaceWrapsEveryValueInLiteral() throws IOException { + stubFindOneAndReplacePipeline(null); + + mongoCollection.createOrReplaceAndReturn(key, document, optionsFor(ReturnDocuments.BEFORE)); + + final BasicDBObject set = (BasicDBObject) pipeline.getValue().get(1).get("$set"); + for (final String field : set.keySet()) { + assertTrue( + ((BasicDBObject) set.get(field)).containsField("$literal"), + field + " is not wrapped in $literal"); + } + assertEquals("Soap", ((BasicDBObject) set.get("item")).get("$literal")); + } + + @Test + void testCreateOrReplaceBothOnInsertDerivesAfterImage() throws IOException { + stubFindOneAndReplacePipeline(null); + + final DocumentChange result = + mongoCollection.createOrReplaceAndReturn(key, document, optionsFor(ReturnDocuments.BOTH)); + + assertEquals(BEFORE, options.getValue().getReturnDocument()); + final BasicDBObject set = literalSetValues(); + assertFalse(result.getBefore().isPresent()); + final BasicDBObject expectedAfter = + new BasicDBObject("createdTime", set.getLong("lastUpdatedTime")) + .append("item", "Soap") + .append("price", 100) + .append("lastUpdatedTime", set.getLong("lastUpdatedTime")) + .append("_lastUpdateTime", set.get("_lastUpdateTime")); + assertEquals(MongoUtils.dbObjectToDocument(expectedAfter), result.getAfter().orElseThrow()); + } + + @Test + void testCreateOrReplaceBothOnUpdateDropsOldFieldsAndKeepsCreatedTime() throws IOException { + stubFindOneAndReplacePipeline(existingDocument()); + + final DocumentChange result = + mongoCollection.createOrReplaceAndReturn(key, document, optionsFor(ReturnDocuments.BOTH)); + + assertEquals(BEFORE, options.getValue().getReturnDocument()); + final BasicDBObject set = literalSetValues(); + assertEquals( + MongoUtils.dbObjectToDocument(existingDocument()), result.getBefore().orElseThrow()); + final BasicDBObject expectedAfter = + new BasicDBObject("createdTime", 123L) + .append("item", "Soap") + .append("price", 100) + .append("lastUpdatedTime", set.getLong("lastUpdatedTime")) + .append("_lastUpdateTime", set.get("_lastUpdateTime")); + assertEquals(MongoUtils.dbObjectToDocument(expectedAfter), result.getAfter().orElseThrow()); + } + + private void stubFindOneAndUpdate(final BasicDBObject result) { + when(collection.findOneAndUpdate( + any(BasicDBObject.class), update.capture(), options.capture())) + .thenReturn(result); + } + + private void stubFindOneAndReplacePipeline(final BasicDBObject result) { + when(collection.findOneAndUpdate( + any(BasicDBObject.class), pipeline.capture(), options.capture())) + .thenReturn(result); + } + + private UpsertOptions optionsFor(final ReturnDocuments returnDocuments) { + return UpsertOptions.builder().returnDocuments(returnDocuments).build(); + } + + private BasicDBObject existingDocument() { + return new BasicDBObject("_id", key.toString()) + .append("item", "OldSoap") + .append("price", 50) + .append("quantity", 3) + .append("createdTime", 123L) + .append("lastUpdatedTime", 456L); + } + + /** Asserts the update uses $currentDate for _lastUpdateTime, like upsert(). */ + private void assertServerLastUpdateTime() { + final BasicDBObject updateDocument = update.getValue(); + assertTrue( + ((BasicDBObject) updateDocument.get("$currentDate")).containsField("_lastUpdateTime")); + assertFalse(((BasicDBObject) updateDocument.get("$set")).containsField("_lastUpdateTime")); + } + + /** Asserts the update uses a literal _lastUpdateTime and returns its $set clause. */ + private BasicDBObject assertExplicitLastUpdateTime() { + final BasicDBObject updateDocument = update.getValue(); + assertFalse(updateDocument.containsField("$currentDate")); + final BasicDBObject set = (BasicDBObject) updateDocument.get("$set"); + final Object lastUpdateTime = set.get("_lastUpdateTime"); + assertTrue(lastUpdateTime instanceof Date); + assertEquals(set.getLong("lastUpdatedTime"), ((Date) lastUpdateTime).getTime()); + assertEquals( + set.getLong("lastUpdatedTime"), + ((BasicDBObject) updateDocument.get("$setOnInsert")).getLong("createdTime")); + return set; + } + + /** Returns the replace pipeline's $set values with the $literal wrappers removed. */ + private BasicDBObject literalSetValues() { + final BasicDBObject set = (BasicDBObject) pipeline.getValue().get(1).get("$set"); + final BasicDBObject values = new BasicDBObject(); + set.forEach((field, value) -> values.put(field, ((BasicDBObject) value).get("$literal"))); + return values; + } } }