Introduce new createOrReplaceAndReturn/bulkCreateOrReplaceAndReturn APIs with WriteAndReturnOptions in Collection

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.WriteAndReturnOptions;
 import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate;
 
 /** Interface spec for common operations on a collection of documents */
@@ -307,6 +308,56 @@ default CloseableIterator<Document> bulkCreateOrReplaceReturnOlderDocuments(
    */
   Document createOrReplaceAndReturn(final Key key, final Document document) throws IOException;
 
+  /**
+   * Create-or-replace and return BEFORE/AFTER/BOTH document image(s) for the given key.
+   *
+   * <p>Atomicity is controlled by {@link WriteAndReturnOptions#isAtomic()}:
+   *
+   * <ul>
+   *   <li>{@code atomic = true} (default): the write and the BEFORE/AFTER capture occur in a single
+   *       round-trip; returned images reflect exactly the row this call wrote.
+   *   <li>{@code atomic = false}: implementations MAY use separate SELECT and INSERT statements.
+   *       Under concurrent writers the BEFORE image may not correspond to the row replaced by this
+   *       call and the AFTER image may reflect a later writer's value.
+   * </ul>
+   *
+   * @param key Unique key of the document in the collection
+   * @param document The document to be created/replaced
+   * @param options Return-image selection and atomicity preference. Must not be {@code null}.
+   * @return Optional result; empty for {@link ReturnOptions#NONE}.
+   * @throws IOException If the operation could not be performed
+   */
+  default Optional<WriteAndReturnResult> createOrReplaceAndReturn(
+      final Key key, final Document document, final WriteAndReturnOptions options)
+      throws IOException {
+    throw new UnsupportedOperationException(
+        "createOrReplaceAndReturn with options is not supported");
+  }
+
+  /**
+   * Batch create-or-replace returning BEFORE/AFTER/BOTH document image(s) per input key.
+   *
+   * <p>Atomicity is controlled by {@link WriteAndReturnOptions#isAtomic()}:
+   *
+   * <ul>
+   *   <li>{@code atomic = true} (default): the write and the BEFORE/AFTER capture occur in a single
+   *       round-trip; returned images reflect exactly the rows this call wrote.
+   *   <li>{@code atomic = false}: implementations MAY use separate SELECT and INSERT statements.
+   *       Per-row consistency between BEFORE/AFTER and what this call wrote is not guaranteed under
+   *       concurrent writers.
+   * </ul>
+   *
+   * @param documents Map of key to document to create or replace
+   * @param options Return-image selection and atomicity preference. Must not be {@code null}.
+   * @return One result per upserted input key (empty for {@link ReturnOptions#NONE})
+   * @throws IOException If the operation could not be performed
+   */
+  default List<WriteAndReturnResult> bulkCreateOrReplaceAndReturn(
+      Map<Key, Document> documents, WriteAndReturnOptions options) throws IOException {
+    throw new UnsupportedOperationException(
+        "bulkCreateOrReplaceAndReturn with options is not supported");
+  }
+
   /**
    * Updates existing documents if the corresponding Filter condition evaluates to true
    *

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

@@ -0,0 +1,8 @@
+package org.hypertrace.core.documentstore;
+
+public enum ReturnOptions {
+  BEFORE,
+  AFTER,
+  NONE,
+  BOTH
+}

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

@@ -0,0 +1,23 @@
+package org.hypertrace.core.documentstore;
+
+import javax.annotation.Nullable;
+import lombok.Value;
+
+@Value
+public class WriteAndReturnResult {
+
+  /** The input key this result corresponds to. Never {@code null}. */
+  Key key;
+
+  /**
+   * The pre-image of the row. Non-null iff BEFORE was requested AND a row existed for {@link
+   * #getKey()} prior to this call; {@code null} otherwise.
+   */
+  @Nullable Document before;
+
+  /**
+   * The post-image of the row. Non-null iff AFTER was requested (always present in that case);
+   * {@code null} otherwise.
+   */
+  @Nullable Document after;
+}

document-store/src/main/java/org/hypertrace/core/documentstore/model/options/WriteAndReturnOptions.java

@@ -0,0 +1,28 @@
+package org.hypertrace.core.documentstore.model.options;
+
+import lombok.Builder;
+import lombok.Builder.Default;
+import lombok.Value;
+import org.hypertrace.core.documentstore.ReturnOptions;
+
+@Value
+@Builder
+public class WriteAndReturnOptions {
+
+  public static final WriteAndReturnOptions DEFAULT_OPTIONS =
+      WriteAndReturnOptions.builder().build();
+
+  /** Which document image(s) to return. Defaults to {@link ReturnOptions#NONE}. */
+  @Default ReturnOptions returnOptions = ReturnOptions.NONE;
+
+  /**
+   * If {@code true} (default), the write and the BEFORE/AFTER capture occur in a single round-trip
+   * such that the returned BEFORE/AFTER images reflect exactly the row(s) this call wrote.
+   *
+   * <p>If {@code false}, implementations MAY use separate SELECT and INSERT statements. In that
+   * case, under concurrent writers, the returned BEFORE image may not correspond to the row
+   * actually replaced by this call, and the returned AFTER image may reflect a later writer's
+   * value. Per-row consistency between BEFORE/AFTER and what this call wrote is not guaranteed.
+   */
+  @Default boolean atomic = true;
+}