Fix and test concurrent replace transactions

Author: smaheshwar-pltr

core/src/main/java/org/apache/iceberg/rest/RESTTableOperations.java

@@ -21,6 +21,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
@@ -77,7 +78,13 @@ enum UpdateType {
   private EncryptingFileIO encryptingFileIO;
   private String tableKeyId;
   private int encryptionDekLength;
-  private List<EncryptedKey> encryptedKeysFromMetadata;
+
+  // keys loaded from the latest metadata
+  private Optional<List<EncryptedKey>> encryptedKeysFromMetadata = Optional.empty();
+
+  // keys added to EM (e.g. as a result of a FileAppend) but not committed into the latest metadata
+  // yet
+  private Optional<List<EncryptedKey>> encryptedKeysPending = Optional.empty();
 
   RESTTableOperations(
       RESTClient client,
@@ -290,9 +297,12 @@ public EncryptionManager encryption() {
               TableProperties.ENCRYPTION_DEK_LENGTH,
               String.valueOf(encryptionDekLength));
 
+      List<EncryptedKey> keys = Lists.newLinkedList();
+      encryptedKeysFromMetadata.ifPresent(keys::addAll);
+      encryptedKeysPending.ifPresent(keys::addAll);
+
       encryptionManager =
-          EncryptionUtil.createEncryptionManager(
-              encryptedKeysFromMetadata, encryptionProperties, kmsClient);
+          EncryptionUtil.createEncryptionManager(keys, encryptionProperties, kmsClient);
     } else {
       return PlaintextEncryptionManager.instance();
     }
@@ -342,7 +352,25 @@ private void encryptionPropsFromMetadata(TableMetadata metadata) {
       return;
     }
 
-    encryptedKeysFromMetadata = metadata.encryptionKeys();
+    encryptedKeysFromMetadata = Optional.ofNullable(metadata.encryptionKeys());
+
+    if (encryptionManager != null) {
+      encryptedKeysPending = Optional.of(Lists.newLinkedList());
+
+      Set<String> keyIdsFromMetadata =
+          encryptedKeysFromMetadata.orElseGet(Lists::newLinkedList).stream()
+              .map(EncryptedKey::keyId)
+              .collect(Collectors.toSet());
+
+      for (EncryptedKey keyFromEM : EncryptionUtil.encryptionKeys(encryptionManager).values()) {
+        if (!keyIdsFromMetadata.contains(keyFromEM.keyId())) {
+          encryptedKeysPending.get().add(keyFromEM);
+        }
+      }
+
+    } else {
+      encryptedKeysPending = Optional.empty();
+    }
 
     // Refresh encryption-related table properties on new/refreshed metadata
     Map<String, String> tableProperties = metadata.properties();
@@ -368,7 +396,7 @@ private TableMetadata updateCurrentMetadata(LoadTableResponse response) {
     if (current == null
         || !Objects.equals(current.metadataFileLocation(), response.metadataLocation())) {
       this.current = checkUUID(current, response.tableMetadata());
-      encryptionPropsFromMetadata(this.current);
+      encryptionPropsFromMetadata(current);
     }
 
     return current;

spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/sql/TestTableEncryption.java

@@ -127,7 +127,7 @@ public void testRefresh() {
   }
 
   @TestTemplate
-  public void testTransaction() {
+  public void testAppendTransaction() {
     validationCatalog.initialize(catalogName, catalogConfig);
     Table table = validationCatalog.loadTable(tableIdent);
 
@@ -140,11 +140,11 @@ public void testTransaction() {
     append.commit();
     transaction.commitTransaction();
 
-    assertThat(currentDataFiles(table).size()).isEqualTo(dataFiles.size() + 1);
+    assertThat(currentDataFiles(table)).hasSize(dataFiles.size() + 1);
   }
 
   @TestTemplate
-  public void testTransactionRetry() {
+  public void testConcurrentAppendTransactions() {
     validationCatalog.initialize(catalogName, catalogConfig);
     Table table = validationCatalog.loadTable(tableIdent);
 
@@ -161,7 +161,39 @@ public void testTransactionRetry() {
     append.commit();
     transaction.commitTransaction();
 
-    assertThat(currentDataFiles(table).size()).isEqualTo(dataFiles.size() + 2);
+    assertThat(currentDataFiles(table)).hasSize(dataFiles.size() + 2);
+  }
+
+  // See CatalogTests#testConcurrentReplaceTransactions
+  @TestTemplate
+  public void testConcurrentReplaceTransactions() {
+    validationCatalog.initialize(catalogName, catalogConfig);
+
+    Table table = validationCatalog.loadTable(tableIdent);
+    DataFile file = currentDataFiles(table).get(0);
+    Schema schema = table.schema();
+
+    // Write data for a replace transaction that will be committed later
+    Transaction secondReplace =
+        validationCatalog
+            .buildTable(tableIdent, schema)
+            .withProperty("encryption.key-id", UnitestKMS.MASTER_KEY_NAME1)
+            .replaceTransaction();
+    secondReplace.newFastAppend().appendFile(file).commit();
+
+    // Commit another replace transaction first
+    Transaction firstReplace =
+        validationCatalog
+            .buildTable(tableIdent, schema)
+            .withProperty("encryption.key-id", UnitestKMS.MASTER_KEY_NAME1)
+            .replaceTransaction();
+    firstReplace.newFastAppend().appendFile(file).commit();
+    firstReplace.commitTransaction();
+
+    secondReplace.commitTransaction();
+
+    Table afterSecondReplace = validationCatalog.loadTable(tableIdent);
+    assertThat(currentDataFiles(afterSecondReplace)).hasSize(1);
   }
 
   @TestTemplate