Allow DomainMetadata to be set in Transaction instead of TransactionBuilder (#4245)

<!--
Thanks for sending a pull request!  Here are some tips for you:
1. If this is your first time, please read our contributor guidelines:
https://github.com/delta-io/delta/blob/master/CONTRIBUTING.md
2. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP]
Your PR title ...'.
  3. Be sure to keep the PR description updated to reflect all changes.
  4. Please write your PR title to summarize what this PR proposes.
5. If possible, provide a concise example to reproduce the issue for a
faster review.
6. If applicable, include the corresponding issue number in the PR title
and link it in the body.
-->

#### Which Delta project/connector is this regarding?
<!--
Please add the component selected below to the beginning of the pull
request title
For example: [Spark] Title of my pull request
-->

- [ ] Spark
- [ ] Standalone
- [ ] Flink
- [x] Kernel
- [ ] Other (fill in here)

## Description

<!--
- Describe what this PR changes.
- Describe why we need the change.
 
If this PR resolves an issue be sure to include "Resolves #XXX" to
correctly link and close the issue upon merge.
-->
This PR allows connectors to specify domain metadata in the
`Transaction` using `addDomainMetadata` and `removeDomainMetadata`
instead of having to specify the domain metadata in the
`TransactionBuilder`.

The `TransactionBuilder` has the method `withDomainMetadataSupported`,
which allows the user to enable domain metadata support.


## How was this patch tested?

<!--
If tests were added, say they were added here. Please make sure to test
the changes thoroughly including negative and positive cases if
possible.
If the changes were tested in any way other than unit tests, please
clarify how you tested step by step (ideally copy and paste-able, so
that other reviewers can test and check, and descendants can verify in
the future).
If the changes were not tested, please explain why.
-->
- All tests that check that the protocol is automatically updated when
passing domain metadata are removed. This is now done explicitly, so the
tests are no longer needed.
## Does this PR introduce _any_ user-facing changes?

<!--
If yes, please clarify the previous behavior and the change this PR
proposes - provide the console output, description and/or an example to
show the behavior difference if possible.
If possible, please also clarify if this is a user-facing change
compared to the released Delta Lake versions or within the unreleased
branches such as master.
If no, write 'No'.
-->

---------

Co-authored-by: Oussama Saoudi <oussama.saoudi@databricks.com>

Author: OussamaSaoudi

kernel/kernel-api/src/main/java/io/delta/kernel/Transaction.java

@@ -28,6 +28,7 @@
 import io.delta.kernel.data.*;
 import io.delta.kernel.engine.Engine;
 import io.delta.kernel.exceptions.ConcurrentWriteException;
+import io.delta.kernel.exceptions.DomainDoesNotExistException;
 import io.delta.kernel.expressions.Literal;
 import io.delta.kernel.internal.DataWriteContextImpl;
 import io.delta.kernel.internal.actions.AddFile;
@@ -97,6 +98,31 @@ public interface Transaction {
   TransactionCommitResult commit(Engine engine, CloseableIterable<Row> dataActions)
       throws ConcurrentWriteException;
 
+  /**
+   * Commit the provided domain metadata as part of this transaction. If this is called more than
+   * once with the same {@code domain} the latest provided {@code config} will be committed in the
+   * transaction. Only user-controlled domains are allowed (aka. domains with a `delta.` prefix are
+   * not allowed). Adding and removing a domain with the same identifier in the same txn is not
+   * allowed. Adding domain metadata to a table that does not support the table feature is not
+   * allowed. To enable the table feature, make sure to call {@link
+   * TransactionBuilder#withDomainMetadataSupported}
+   *
+   * @param domain the domain identifier
+   * @param config configuration string for this domain
+   */
+  void addDomainMetadata(String domain, String config);
+
+  /**
+   * Mark the domain metadata with identifier {@code domain} as removed in this transaction. If this
+   * domain does not exist in the latest version of the table, calling {@link
+   * Transaction#commit(Engine, CloseableIterable)} will throw a {@link
+   * DomainDoesNotExistException}. Adding and removing a domain with the same identifier in one txn
+   * is not allowed.
+   *
+   * @param domain the domain identifier for the domain to remove
+   */
+  void removeDomainMetadata(String domain);
+
   /**
    * Given the logical data that needs to be written to the table, convert it into the required
    * physical data depending upon the table Delta protocol and features enabled on the table. Kernel

kernel/kernel-api/src/main/java/io/delta/kernel/TransactionBuilder.java

@@ -93,38 +93,19 @@ TransactionBuilder withTransactionId(
   TransactionBuilder withMaxRetries(int maxRetries);
 
   /**
-   * Commit the provided domain metadata as part of this transaction. If this is called more than
-   * once with the same {@code domain} the latest provided {@code config} will be committed in the
-   * transaction. Only user-controlled domains are allowed (aka. domains with a `delta.` prefix are
-   * not allowed). Adding and removing a domain with the same identifier in the same txn is not
-   * allowed.
-   *
-   * <p>See the Delta protocol for more information on how to use domain metadata <a
-   * href="https://github.com/delta-io/delta/blob/master/PROTOCOL.md#domain-metadata">Domain
-   * Metadata</a>.
-   *
-   * <p>Please note using this API will automatically upgrade the protocol of the table to support
-   * Domain Metadata if it is not already supported. See <a
+   * Enables support for Domain Metadata on this table if it is not supported already. The table
+   * feature _must_ be supported on the table to add or remove domain metadata using {@link
+   * Transaction#addDomainMetadata} or {@link Transaction#removeDomainMetadata}. See <a
    * href="https://docs.delta.io/latest/versioning.html#how-does-delta-lake-manage-feature-compatibility">
-   * How does Delta Lake manage feature compatibility?</a> for more details. This may break existing
-   * writers that do not support the Domain Metadata feature; readers will be unaffected.
-   *
-   * @param domain the domain identifier
-   * @param config configuration string for this domain
-   * @return updated {@link TransactionBuilder} instance
-   */
-  TransactionBuilder withDomainMetadata(String domain, String config);
-
-  /**
-   * Mark the domain metadata with identifier {@code domain} as removed in this transaction. If this
-   * domain does not exist in the latest version of the table will throw a {@link
-   * DomainDoesNotExistException} upon calling {@link TransactionBuilder#build(Engine)}. Adding and
-   * removing a domain with the same identifier in one txn is not allowed.
+   * How does Delta Lake manage feature compatibility?</a> for more details on table feature
+   * support.
    *
-   * @param domain the domain identifier for the domain to remove
-   * @return updated {@link TransactionBuilder} instance
+   * <p>See the Delta protocol for more information on how to use <a
+   * href="https://github.com/delta-io/delta/blob/master/PROTOCOL.md#domain-metadata">Domain
+   * Metadata</a>. This may break existing writers that do not support the Domain Metadata feature;
+   * readers will be unaffected.
    */
-  TransactionBuilder withDomainMetadataRemoved(String domain);
+  TransactionBuilder withDomainMetadataSupported();
 
   /**
    * Build the transaction. Also validates the given info to ensure that a valid transaction can be

kernel/kernel-api/src/main/java/io/delta/kernel/internal/TransactionBuilderImpl.java

@@ -29,7 +29,6 @@
 
 import io.delta.kernel.*;
 import io.delta.kernel.engine.Engine;
-import io.delta.kernel.exceptions.DomainDoesNotExistException;
 import io.delta.kernel.exceptions.TableNotFoundException;
 import io.delta.kernel.internal.actions.*;
 import io.delta.kernel.internal.fs.Path;
@@ -58,12 +57,11 @@ public class TransactionBuilderImpl implements TransactionBuilder {
   private final TableImpl table;
   private final String engineInfo;
   private final Operation operation;
-  private final Map<String, DomainMetadata> domainMetadatasAdded = new HashMap<>();
-  private final Set<String> domainMetadatasRemoved = new HashSet<>();
   private Optional<StructType> schema = Optional.empty();
   private Optional<List<String>> partitionColumns = Optional.empty();
   private Optional<SetTransaction> setTxnOpt = Optional.empty();
   private Optional<Map<String, String>> tableProperties = Optional.empty();
+  private boolean needDomainMetadataSupport = false;
 
   /**
    * Number of retries for concurrent write exceptions to resolve conflicts and retry commit. In
@@ -119,27 +117,8 @@ public TransactionBuilder withMaxRetries(int maxRetries) {
   }
 
   @Override
-  public TransactionBuilder withDomainMetadata(String domain, String config) {
-    checkArgument(
-        DomainMetadata.isUserControlledDomain(domain),
-        "Setting a system-controlled domain is not allowed: " + domain);
-    checkArgument(
-        !domainMetadatasRemoved.contains(domain),
-        "Cannot add a domain that is removed in this transaction");
-    // we override any existing value
-    domainMetadatasAdded.put(domain, new DomainMetadata(domain, config, false /* removed */));
-    return this;
-  }
-
-  @Override
-  public TransactionBuilder withDomainMetadataRemoved(String domain) {
-    checkArgument(
-        DomainMetadata.isUserControlledDomain(domain),
-        "Removing a system-controlled domain is not allowed: " + domain);
-    checkArgument(
-        !domainMetadatasAdded.containsKey(domain),
-        "Cannot remove a domain that is added in this transaction");
-    domainMetadatasRemoved.add(domain);
+  public TransactionBuilder withDomainMetadataSupported() {
+    needDomainMetadataSupport = true;
     return this;
   }
 
@@ -189,9 +168,7 @@ public Transaction build(Engine engine) {
     // Ex: We enable feature `icebergCompatV2` plus dependent features `columnMapping`
     Optional<Tuple2<Protocol, Set<TableFeature>>> newProtocolAndFeatures =
         TableFeatures.autoUpgradeProtocolBasedOnMetadata(
-            newMetadata.orElse(snapshotMetadata),
-            !domainMetadatasAdded.isEmpty(),
-            snapshotProtocol);
+            newMetadata.orElse(snapshotMetadata), needDomainMetadataSupport, snapshotProtocol);
     if (newProtocolAndFeatures.isPresent()) {
       logger.info(
           "Automatically enabling table features: {}",
@@ -244,8 +221,7 @@ public Transaction build(Engine engine) {
         newMetadata.isPresent() /* shouldUpdateMetadata */,
         newProtocol.isPresent() /* shouldUpdateProtocol */,
         maxRetries,
-        table.getClock(),
-        getDomainMetadatasToCommit(snapshot));
+        table.getClock());
   }
 
   /**
@@ -386,45 +362,4 @@ private Metadata getInitialMetadata() {
   private Protocol getInitialProtocol() {
     return new Protocol(DEFAULT_READ_VERSION, DEFAULT_WRITE_VERSION);
   }
-
-  /**
-   * Returns a list of the domain metadatas to commit. This consists of the domain metadatas added
-   * in the transaction using {@link TransactionBuilder#withDomainMetadata(String, String)} and the
-   * tombstones for the domain metadatas removed in the transaction using {@link
-   * TransactionBuilder#withDomainMetadataRemoved(String)}.
-   */
-  private List<DomainMetadata> getDomainMetadatasToCommit(SnapshotImpl snapshot) {
-    // Add all domain metadatas added in the transaction
-    List<DomainMetadata> finalDomainMetadatas = new ArrayList<>(domainMetadatasAdded.values());
-
-    // Generate the tombstones for the removed domain metadatas
-    Map<String, DomainMetadata> snapshotDomainMetadataMap = snapshot.getDomainMetadataMap();
-    for (String domainName : domainMetadatasRemoved) {
-      // Note: we know domainName is not already in finalDomainMetadatas because we do not allow
-      // removing and adding a domain with the same identifier in a single txn!
-      if (snapshotDomainMetadataMap.containsKey(domainName)) {
-        DomainMetadata domainToRemove = snapshotDomainMetadataMap.get(domainName);
-        if (domainToRemove.isRemoved()) {
-          // If the domain is already removed we throw an error to avoid any inconsistencies or
-          // ambiguity. The snapshot read by the connector is inconsistent with the snapshot
-          // loaded here as the domain to remove no longer exists.
-          throw new DomainDoesNotExistException(
-              table.getDataPath().toString(), domainName, snapshot.getVersion());
-        }
-        finalDomainMetadatas.add(domainToRemove.removed());
-      } else {
-        // We must throw an error if the domain does not exist. Otherwise, there could be unexpected
-        // behavior within conflict resolution. For example, consider the following
-        // 1. Table has no domains set in V0
-        // 2. txnA is started and wants to remove domain "foo"
-        // 3. txnB is started and adds domain "foo" and commits V1 before txnA
-        // 4. txnA needs to perform conflict resolution against the V1 commit from txnB
-        // Conflict resolution should fail but since the domain does not exist we cannot create
-        // a tombstone to mark it as removed and correctly perform conflict resolution.
-        throw new DomainDoesNotExistException(
-            table.getDataPath().toString(), domainName, snapshot.getVersion());
-      }
-    }
-    return finalDomainMetadatas;
-  }
 }

kernel/kernel-api/src/main/java/io/delta/kernel/internal/TransactionImpl.java

@@ -29,6 +29,7 @@
 import io.delta.kernel.data.Row;
 import io.delta.kernel.engine.Engine;
 import io.delta.kernel.exceptions.ConcurrentWriteException;
+import io.delta.kernel.exceptions.DomainDoesNotExistException;
 import io.delta.kernel.expressions.Column;
 import io.delta.kernel.hook.PostCommitHook;
 import io.delta.kernel.internal.actions.*;
@@ -81,7 +82,9 @@ public class TransactionImpl implements Transaction {
   private final Optional<SetTransaction> setTxnOpt;
   private final boolean shouldUpdateProtocol;
   private final Clock clock;
-  private List<DomainMetadata> domainMetadatas;
+  private final Map<String, DomainMetadata> domainMetadatasAdded = new HashMap<>();
+  private final Set<String> domainMetadatasRemoved = new HashSet<>();
+  private Optional<List<DomainMetadata>> domainMetadatas = Optional.empty();
   private Metadata metadata;
   private boolean shouldUpdateMetadata;
   private int maxRetries;
@@ -101,8 +104,7 @@ public TransactionImpl(
       boolean shouldUpdateMetadata,
       boolean shouldUpdateProtocol,
       int maxRetries,
-      Clock clock,
-      List<DomainMetadata> domainMetadatas) {
+      Clock clock) {
     this.isNewTable = isNewTable;
     this.dataPath = dataPath;
     this.logPath = logPath;
@@ -116,7 +118,6 @@ public TransactionImpl(
     this.shouldUpdateProtocol = shouldUpdateProtocol;
     this.maxRetries = maxRetries;
     this.clock = clock;
-    this.domainMetadatas = domainMetadatas;
   }
 
   @Override
@@ -150,18 +151,94 @@ public Optional<SetTransaction> getSetTxnOpt() {
     return setTxnOpt;
   }
 
-  /**
-   * Internal API to add domain metadata actions for this transaction. Visible for testing.
-   *
-   * @param domainMetadatas List of domain metadata to be added to the transaction.
-   */
   @VisibleForTesting
-  public void addDomainMetadatas(List<DomainMetadata> domainMetadatas) {
-    this.domainMetadatas.addAll(domainMetadatas);
+  public void addDomainMetadataInternal(String domain, String config) {
+    checkArgument(
+        !domainMetadatasRemoved.contains(domain),
+        "Cannot add a domain that is removed in this transaction");
+    checkState(!closed, "Cannot add a domain metadata after the transaction has completed");
+    // we override any existing value
+    domainMetadatasAdded.put(domain, new DomainMetadata(domain, config, false /* removed */));
+  }
+
+  @Override
+  public void addDomainMetadata(String domain, String config) {
+    checkState(
+        TableFeatures.isDomainMetadataSupported(protocol),
+        "Unable to add domain metadata when the domain metadata table feature is disabled");
+    checkArgument(
+        DomainMetadata.isUserControlledDomain(domain),
+        "Setting a system-controlled domain is not allowed: " + domain);
+    addDomainMetadataInternal(domain, config);
+  }
+
+  @VisibleForTesting
+  public void removeDomainMetadataInternal(String domain) {
+    checkArgument(
+        !domainMetadatasAdded.containsKey(domain),
+        "Cannot remove a domain that is added in this transaction");
+    checkState(!closed, "Cannot remove a domain after the transaction has completed");
+    domainMetadatasRemoved.add(domain);
   }
 
+  @Override
+  public void removeDomainMetadata(String domain) {
+    checkState(
+        TableFeatures.isDomainMetadataSupported(protocol),
+        "Unable to add domain metadata when the domain metadata table feature is disabled");
+    checkArgument(
+        DomainMetadata.isUserControlledDomain(domain),
+        "Removing a system-controlled domain is not allowed: " + domain);
+    removeDomainMetadataInternal(domain);
+  }
+
+  /**
+   * Returns a list of the domain metadatas to commit. This consists of the domain metadatas added
+   * in the transaction using {@link Transaction#addDomainMetadata(String, String)} and the
+   * tombstones for the domain metadatas removed in the transaction using {@link
+   * Transaction#removeDomainMetadata(String)}. The result is stored in {@code domainMetadatas}.
+   *
+   * @return A list of {@link DomainMetadata} containing domain metadata to be committed in this
+   *     transaction.
+   */
   public List<DomainMetadata> getDomainMetadatas() {
-    return domainMetadatas;
+    // If we have already processed the domain metadatas, then return the list.
+    if (domainMetadatas.isPresent()) {
+      return domainMetadatas.get();
+    }
+    // Add all domain metadatas added in the transaction
+    List<DomainMetadata> finalDomainMetadatas = new ArrayList<>(domainMetadatasAdded.values());
+
+    // Generate the tombstones for the removed domain metadatas
+    Map<String, DomainMetadata> snapshotDomainMetadataMap = readSnapshot.getDomainMetadataMap();
+    for (String domainName : domainMetadatasRemoved) {
+      // Note: we know domainName is not already in finalDomainMetadatas because we do not allow
+      // removing and adding a domain with the same identifier in a single txn!
+      if (snapshotDomainMetadataMap.containsKey(domainName)) {
+        DomainMetadata domainToRemove = snapshotDomainMetadataMap.get(domainName);
+        if (domainToRemove.isRemoved()) {
+          // If the domain is already removed we throw an error to avoid any inconsistencies or
+          // ambiguity. The snapshot read by the connector is inconsistent with the snapshot
+          // loaded here as the domain to remove no longer exists.
+          throw new DomainDoesNotExistException(
+              dataPath.toString(), domainName, readSnapshot.getVersion());
+        }
+        finalDomainMetadatas.add(domainToRemove.removed());
+      } else {
+        // We must throw an error if the domain does not exist. Otherwise, there could be unexpected
+        // behavior within conflict resolution. For example, consider the following
+        // 1. Table has no domains set in V0
+        // 2. txnA is started and wants to remove domain "foo"
+        // 3. txnB is started and adds domain "foo" and commits V1 before txnA
+        // 4. txnA needs to perform conflict resolution against the V1 commit from txnB
+        // Conflict resolution should fail but since the domain does not exist we cannot create
+        // a tombstone to mark it as removed and correctly perform conflict resolution.
+        throw new DomainDoesNotExistException(
+            dataPath.toString(), domainName, readSnapshot.getVersion());
+      }
+    }
+    domainMetadatas = Optional.of(finalDomainMetadatas);
+    return finalDomainMetadatas;
   }
 
   public Protocol getProtocol() {
@@ -201,18 +278,20 @@ private TransactionCommitResult commitWithRetry(
       CommitInfo attemptCommitInfo = generateCommitAction(engine);
       updateMetadataWithICTIfRequired(
           engine, attemptCommitInfo.getInCommitTimestamp(), readSnapshot.getVersion());
+      List<DomainMetadata> resolvedDomainMetadatas = getDomainMetadatas();
 
       // If row tracking is supported, assign base row IDs and default row commit versions to any
       // AddFile actions that do not yet have them. If the row ID high watermark changes, emit a
       // DomainMetadata action to update it.
       if (TableFeatures.isRowTrackingSupported(protocol)) {
-        domainMetadatas =
+        List<DomainMetadata> updatedDomainMetadata =
             RowTracking.updateRowIdHighWatermarkIfNeeded(
                 readSnapshot,
                 protocol,
                 Optional.empty() /* winningTxnRowIdHighWatermark */,
                 dataActions,
-                domainMetadatas);
+                resolvedDomainMetadatas);
+        domainMetadatas = Optional.of(updatedDomainMetadata);
         dataActions =
             RowTracking.assignBaseRowIdAndDefaultRowCommitVersion(
                 readSnapshot,
@@ -239,7 +318,7 @@ private TransactionCommitResult commitWithRetry(
                 resolveConflicts(engine, commitAsVersion, attemptCommitInfo, numTries, dataActions);
             commitAsVersion = rebaseState.getLatestVersion() + 1;
             dataActions = rebaseState.getUpdatedDataActions();
-            domainMetadatas = rebaseState.getUpdatedDomainMetadatas();
+            domainMetadatas = Optional.of(rebaseState.getUpdatedDomainMetadatas());
           }
         }
         numTries++;
@@ -329,10 +408,12 @@ private TransactionCommitResult doCommit(
     }
     setTxnOpt.ifPresent(setTxn -> metadataActions.add(createTxnSingleAction(setTxn.toRow())));
 
+    List<DomainMetadata> resolvedDomainMetadatas = getDomainMetadatas();
+
     // Check for duplicate domain metadata and if the protocol supports
-    DomainMetadataUtils.validateDomainMetadatas(domainMetadatas, protocol);
+    DomainMetadataUtils.validateDomainMetadatas(resolvedDomainMetadatas, protocol);
 
-    domainMetadatas.forEach(
+    resolvedDomainMetadatas.forEach(
         dm -> metadataActions.add(createDomainMetadataSingleAction(dm.toRow())));
 
     try (CloseableIterator<Row> stageDataIter = dataActions.iterator()) {