Fix replace transactions using Hive approach

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

@@ -164,6 +164,7 @@ public void testConcurrentAppendTransactions() {
     assertThat(currentDataFiles(table)).hasSize(dataFiles.size() + 2);
   }
 
+  // See CatalogTests#testConcurrentReplaceTransactions
   @TestTemplate
   public void testConcurrentReplaceTransactions() {
     validationCatalog.initialize(catalogName, catalogConfig);
@@ -174,18 +175,18 @@ public void testConcurrentReplaceTransactions() {
 
     // 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();
+        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();
+        validationCatalog
+            .buildTable(tableIdent, schema)
+            .withProperty("encryption.key-id", UnitestKMS.MASTER_KEY_NAME1)
+            .replaceTransaction();
     firstReplace.newFastAppend().appendFile(file).commit();
     firstReplace.commitTransaction();