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..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 @@ -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,9 +20,15 @@ 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; +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; @@ -31,12 +40,16 @@ 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; 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 +63,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 +82,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(); @@ -89,6 +133,42 @@ 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; + } + + 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()).map(rd -> Arguments.of(write, rd))); + } + } + + private static class AtomicWritesProvider implements ArgumentsProvider { + + @Override + public Stream provideArguments(final ExtensionContext context) { + return Stream.of(AtomicWrite.values()).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()) { @@ -495,6 +575,260 @@ void testUpsertAndReturnExistingDoc(String storeName) throws Exception { } } + @Nested + @DisplayName("Atomic write with UpsertOptions tests") + class AtomicWriteWithUpsertOptionsTests { + + 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); + } + + private Document readBack(Collection collection, String docId) throws IOException { + try (CloseableIterator iterator = collection.find(buildQueryById(docId))) { + assertTrue(iterator.hasNext()); + return iterator.next(); + } + } + + 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} {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(write.storeName); + + DocumentChange change = atomicWrite(write, key, createTestDocument(docId), returnDocuments); + + assertTrue(change.getBefore().isEmpty()); + 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, afterJson); + } + } + + @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(write.storeName); + String createdField = write.createdTimeField(); + + collection.upsert(key, createTestDocument(docId)); + JsonNode initialJson = OBJECT_MAPPER.readTree(readBack(collection, docId).toJson()); + assertNotNull(initialJson.get(createdField)); + + 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, beforeJson); + assertEquals("TestItem", beforeJson.get("item").asText()); + assertEquals(100, beforeJson.get("price").asInt()); + } + + 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, afterJson); + } + } + + @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()); + } + + 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")); + } + + @Test + @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(AtomicWritesProvider.class) + void testConcurrentWritesFormChain(AtomicWrite write) throws Exception { + String docId = generateDocId("atomic-concurrent"); + Key key = createKey(docId); + 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(write, key, new JSONDocument(node), ReturnDocuments.BOTH); + }); + } + + List changes = new ArrayList<>(); + for (Future future : executor.invokeAll(tasks)) { + changes.add(future.get()); + } + + assertChangesFormSingleChain(changes, "price"); + } finally { + executor.shutdownNow(); + } + } + + 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 { + 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<>(); + Integer head = null; + for (DocumentChange change : changes) { + 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; + } + } + + for (Integer before : afterByBefore.keySet()) { + assertTrue( + afterValues.contains(before), + "before value " + before + " must match some write's after value"); + } + + Integer current = head; + for (int step = 1; step < changes.size(); step++) { + Integer next = afterByBefore.get(current); + if (next == null) { + fail("chain gap after value " + current); + } + current = next; + } + } + } + @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 a3cb7d8ca..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 @@ -11,9 +11,16 @@ 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 */ +/** + * 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 { /** @@ -57,6 +64,26 @@ public interface Collection { @Deprecated(forRemoval = true) Document upsertAndReturn(Key key, Document document) throws IOException; + /** + * Atomically perform the same merge-upsert as {@link #upsert(Key, Document)} and 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. + * + * @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 upsertAndReturn(Key key, Document document, UpsertOptions options) + throws IOException { + throw new UnsupportedOperationException("upsertAndReturn with UpsertOptions is not supported"); + } + /** * Update a sub document * @@ -307,6 +334,27 @@ 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 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. + * + * @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 createOrReplaceAndReturn( + final Key key, final Document document, final UpsertOptions options) throws IOException { + throw new UnsupportedOperationException( + "createOrReplaceAndReturn with UpsertOptions 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..c78b8da49 --- /dev/null +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/DocumentChange.java @@ -0,0 +1,61 @@ +package org.hypertrace.core.documentstore; + +import java.util.Objects; +import java.util.Optional; +import javax.annotation.Nullable; + +/** + * 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 { + + @Nullable private final Document before; + @Nullable private final Document after; + + private DocumentChange(@Nullable final Document before, @Nullable final Document after) { + this.before = before; + this.after = after; + } + + /** + * @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); + } + + /** Empty when the after-image was not requested. */ + public Optional getAfter() { + return Optional.ofNullable(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) && Objects.equals(after, 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/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 f80ef9e0a..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 @@ -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; @@ -56,12 +57,15 @@ 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; 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; @@ -300,6 +304,148 @@ 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; 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. For that reason {@code + * _lastUpdateTime} is set from the client clock instead of {@code $currentDate} in this + * mode only. + *
+ */ + @Override + public DocumentChange upsertAndReturn(Key key, Document document, UpsertOptions options) + throws IOException { + final ReturnDocuments returnDocuments = options.getReturnDocuments(); + final BasicDBObject update = + returnDocuments == ReturnDocuments.BOTH + ? prepareUpsertWithExplicitLastUpdateTime(key, document, System.currentTimeMillis()) + : prepareUpsert(key, document); + final ReturnDocument serverImage = + returnDocuments == ReturnDocuments.AFTER ? ReturnDocument.AFTER : ReturnDocument.BEFORE; + final BasicDBObject serverResult = + Failsafe.with(upsertRetryPolicy) + .get( + () -> + collection.findOneAndUpdate( + this.selectionCriteriaForKey(key), + update, + new FindOneAndUpdateOptions().upsert(true).returnDocument(serverImage))); + + switch (returnDocuments) { + case BEFORE: + return DocumentChange.of(toDocumentOrNull(serverResult), null); + case AFTER: + if (serverResult == null) { + throw new IOException("Could not upsert the document with key: " + key); + } + return DocumentChange.of(null, dbObjectToDocument(serverResult)); + case BOTH: + 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). + * + *
    + *
  • 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(); + final long now = System.currentTimeMillis(); + final BasicDBObject replacement = prepareDocumentForLiteralReplace(key, document, now); + final ReturnDocument serverImage = + returnDocuments == ReturnDocuments.AFTER ? ReturnDocument.AFTER : ReturnDocument.BEFORE; + + try { + final BasicDBObject serverResult = + collection.findOneAndUpdate( + this.selectionCriteriaForKey(key), + 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); + } + @SuppressWarnings("MismatchedQueryAndUpdateOfCollection") private List prepareForCreateOrReplace(final Key key, final Document document) throws JsonProcessingException { @@ -328,6 +474,37 @@ private BasicDBObject prepareUpsert(Key key, Document document) throws JsonProce .append("$setOnInsert", new BasicDBObject(CREATED_TIME, now)); } + /** + * 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)); + } + + /** + * 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")); + } + 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 100449189..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; @@ -39,6 +42,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; @@ -49,12 +53,18 @@ 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.JSONDocument; 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.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; @@ -62,6 +72,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; @@ -96,6 +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 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"; @@ -231,6 +245,106 @@ 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 DocumentChange upsertAndReturn(Key key, Document document, UpsertOptions options) + throws IOException { + 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 createOrReplaceAndReturn(Key key, Document document, UpsertOptions options) + throws IOException { + return writeAndReturnChange( + key, document, options, this::buildCreateOrReplaceSql, "createOrReplaceAndReturn", false); + } + + /** + * 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, + 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(); + + 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; + } + } catch (PSQLException e) { + if (!isRetry && shouldRefreshSchemaAndRetry(e.getSQLState())) { + LOGGER.info( + "Schema mismatch detected during {} (SQLState: {}), refreshing schema and retrying. key: {}", + operation, + e.getSQLState(), + key); + schemaRegistry.invalidate(tableName); + return writeAndReturnChange(key, document, options, sqlBuilder, operation, true); + } + 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 {}. key: {} content: {}", operation, 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 +1860,95 @@ private String buildCreateOrReplaceSql(List allTableColumns, String pkCo tableIdentifier, columnList, placeholders, pkColumn, setClause); } + 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, IOException { + 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 (!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)); + } + } + } + + /** + * 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); + } + } + private boolean executeUpsert(String sql, List allColumns, TypedDocument parsed) throws SQLException { long ta = System.nanoTime(); @@ -1859,6 +2062,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..436be3cb8 --- /dev/null +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/DocumentChangeTest.java @@ -0,0 +1,47 @@ +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.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertSame; + +import java.io.IOException; +import org.junit.jupiter.api.Test; + +class DocumentChangeTest { + + @Test + void bothImagesPresent() throws IOException { + Document before = new JSONDocument("{\"item\":\"before\"}"); + Document after = new JSONDocument("{\"item\":\"after\"}"); + DocumentChange change = DocumentChange.of(before, after); + assertSame(before, change.getBefore().orElseThrow()); + assertSame(after, change.getAfter().orElseThrow()); + } + + @Test + 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 172d71b13..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 @@ -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; @@ -38,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; @@ -45,10 +47,12 @@ 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; 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; @@ -57,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; @@ -452,4 +458,246 @@ void testBulkUpdateWithoutUpdates() { UpdateOptions.DEFAULT_UPDATE_OPTIONS)); } } + + @Nested + 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); + + @SuppressWarnings("unchecked") + private final ArgumentCaptor> pipeline = + ArgumentCaptor.forClass(List.class); + + UpsertAndReturnWithOptionsTest() throws IOException { + document = new JSONDocument("{\"item\":\"Soap\",\"price\":100}"); + } + + @Test + void testBeforeOnInsertReturnsNothing() throws IOException { + stubFindOneAndUpdate(null); + + final DocumentChange result = + mongoCollection.upsertAndReturn(key, document, optionsFor(ReturnDocuments.BEFORE)); + + assertFalse(result.getBefore().isPresent()); + assertFalse(result.getAfter().isPresent()); + assertTrue(options.getValue().isUpsert()); + assertEquals(BEFORE, options.getValue().getReturnDocument()); + assertServerLastUpdateTime(); + } + + @Test + void testBeforeOnUpdateReturnsBeforeImage() throws IOException { + stubFindOneAndUpdate(existingDocument()); + + final DocumentChange result = + mongoCollection.upsertAndReturn(key, document, optionsFor(ReturnDocuments.BEFORE)); + + assertEquals( + MongoUtils.dbObjectToDocument(existingDocument()), result.getBefore().orElseThrow()); + assertFalse(result.getAfter().isPresent()); + assertEquals(BEFORE, options.getValue().getReturnDocument()); + assertServerLastUpdateTime(); + } + + @Test + void testAfterReturnsServerAfterImage() throws IOException { + final BasicDBObject serverAfter = + new BasicDBObject("_id", key.toString()).append("item", "Soap").append("price", 100); + stubFindOneAndUpdate(serverAfter); + + final DocumentChange result = + mongoCollection.upsertAndReturn(key, document, optionsFor(ReturnDocuments.AFTER)); + + assertFalse(result.getBefore().isPresent()); + assertEquals( + MongoUtils.dbObjectToDocument(new BasicDBObject("item", "Soap").append("price", 100)), + result.getAfter().orElseThrow()); + assertTrue(options.getValue().isUpsert()); + assertEquals(AFTER, options.getValue().getReturnDocument()); + assertServerLastUpdateTime(); + } + + @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 testBothOnUpdateDerivesAfterImage() throws IOException { + stubFindOneAndUpdate(existingDocument()); + + final DocumentChange result = + mongoCollection.upsertAndReturn(key, document, optionsFor(ReturnDocuments.BOTH)); + + 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()); + } + + @Test + void testCreateOrReplaceBeforeAndAfter() throws IOException { + stubFindOneAndReplacePipeline(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()); + } + + @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; + } + } }