Rebase properly

Author: Sreesh Maheshwar

core/src/main/java/org/apache/iceberg/encryption/EncryptionUtil.java

@@ -75,7 +75,7 @@ public static KeyManagementClient createKmsClient(Map<String, String> catalogPro
     return kmsClient;
   }
 
-  static EncryptionManager createEncryptionManager(
+  public static EncryptionManager createEncryptionManager(
       List<EncryptedKey> keys, Map<String, String> tableProperties, KeyManagementClient kmsClient) {
     Preconditions.checkArgument(kmsClient != null, "Invalid KMS client: null");
     String tableKeyId = tableProperties.get(TableProperties.ENCRYPTION_TABLE_KEY);
@@ -96,7 +96,7 @@ static EncryptionManager createEncryptionManager(
         "Invalid data key length: %s (must be 16, 24, or 32)",
         dataKeyLength);
 
-    return new StandardEncryptionManager(tableKeyId, dataKeyLength, kmsClient);
+    return new StandardEncryptionManager(keys, tableKeyId, dataKeyLength, kmsClient);
   }
 
   public static EncryptedOutputFile plainAsEncryptedOutput(OutputFile encryptingOutputFile) {
@@ -145,10 +145,10 @@ static ByteBuffer encryptManifestListKeyMetadata(
     return ByteBuffer.wrap(encryptedKeyMetadata);
   }
 
-    public static Map<String, EncryptedKey> encryptionKeys(EncryptionManager em) {
-        Preconditions.checkState(
-                em instanceof StandardEncryptionManager,
-                "Encryption keys are only available for StandardEncryptionManager");
-        return ((StandardEncryptionManager) em).encryptionKeys();
-    }
+  public static Map<String, EncryptedKey> encryptionKeys(EncryptionManager em) {
+    Preconditions.checkState(
+        em instanceof StandardEncryptionManager,
+        "Encryption keys are only available for StandardEncryptionManager");
+    return ((StandardEncryptionManager) em).encryptionKeys();
+  }
 }

core/src/main/java/org/apache/iceberg/encryption/StandardEncryptionManager.java

@@ -23,6 +23,7 @@
 import java.nio.ByteBuffer;
 import java.security.SecureRandom;
 import java.util.Base64;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import org.apache.iceberg.TableProperties;
@@ -46,9 +47,16 @@ private class TransientEncryptionState {
     private final Map<String, EncryptedKey> encryptionKeys;
     private final LoadingCache<String, ByteBuffer> unwrappedKeyCache;
 
-    private TransientEncryptionState(KeyManagementClient kmsClient) {
+    private TransientEncryptionState(List<EncryptedKey> keys, KeyManagementClient kmsClient) {
       this.kmsClient = kmsClient;
       this.encryptionKeys = Maps.newLinkedHashMap();
+
+      for (EncryptedKey key : keys) {
+        Preconditions.checkArgument(
+            key.keyId() != null, "Key id cannot be null"); // Required by spec.
+        encryptionKeys.put(key.keyId(), key);
+      }
+
       this.unwrappedKeyCache =
           Caffeine.newBuilder()
               .expireAfterWrite(1, TimeUnit.HOURS)
@@ -70,14 +78,28 @@ private TransientEncryptionState(KeyManagementClient kmsClient) {
    */
   public StandardEncryptionManager(
       String tableKeyId, int dataKeyLength, KeyManagementClient kmsClient) {
+    this(List.of(), tableKeyId, dataKeyLength, kmsClient);
+  }
+
+  /**
+   * @param keys a list of existing {@link EncryptedKey}s for this {@link EncryptionManager} to use
+   * @param tableKeyId table encryption key id
+   * @param dataKeyLength length of data encryption key (16/24/32 bytes)
+   * @param kmsClient Client of KMS used to wrap/unwrap keys in envelope encryption
+   */
+  public StandardEncryptionManager(
+      List<EncryptedKey> keys,
+      String tableKeyId,
+      int dataKeyLength,
+      KeyManagementClient kmsClient) {
     Preconditions.checkNotNull(tableKeyId, "Invalid encryption key ID: null");
     Preconditions.checkArgument(
         dataKeyLength == 16 || dataKeyLength == 24 || dataKeyLength == 32,
         "Invalid data key length: %s (must be 16, 24, or 32)",
         dataKeyLength);
     Preconditions.checkNotNull(kmsClient, "Invalid KMS client: null");
     this.tableKeyId = tableKeyId;
-    this.transientState = new TransientEncryptionState(kmsClient);
+    this.transientState = new TransientEncryptionState(keys, kmsClient);
     this.dataKeyLength = dataKeyLength;
   }
 
@@ -199,6 +221,14 @@ public String addManifestListKeyMetadata(NativeEncryptionKeyMetadata keyMetadata
     return manifestListKeyID;
   }
 
+  public Map<String, EncryptedKey> encryptionKeys() {
+    if (transientState == null) {
+      throw new IllegalStateException("Cannot return encryption keys after serialization");
+    }
+
+    return transientState.encryptionKeys;
+  }
+
   private String generateKeyId() {
     byte[] idBytes = new byte[16];
     workerRNG().nextBytes(idBytes);

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

@@ -49,7 +49,6 @@
 import org.apache.iceberg.catalog.TableIdentifier;
 import org.apache.iceberg.encryption.EncryptionUtil;
 import org.apache.iceberg.encryption.KeyManagementClient;
-import org.apache.iceberg.encryption.EncryptionUtil;
 import org.apache.iceberg.exceptions.AlreadyExistsException;
 import org.apache.iceberg.exceptions.NoSuchNamespaceException;
 import org.apache.iceberg.exceptions.NoSuchTableException;

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

@@ -152,7 +152,7 @@ public void commit(TableMetadata base, TableMetadata metadata) {
         builder.addEncryptionKey(entry.getValue());
       }
       metadataToCommit = builder.build();
-      // TODO(smaheshwar): Think about requirements.
+      // TODO(smaheshwar-pltr): Think about requirements
     }
 
     switch (updateType) {
@@ -246,7 +246,7 @@ public EncryptionManager encryption() {
   }
 
   private void encryptionPropsFromMetadata(TableMetadata metadata) {
-    // TODO(smaheshwar): Check generally for changed encryption-related properties!
+    // TODO(smaheshwar-pltr): Check generally for changed encryption-related properties
     if (metadata == null || metadata.properties() == null) {
       return;
     }

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

@@ -51,10 +51,10 @@
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.TestTemplate;
 
-// TODO(smaheshwar): This test is taken from https://github.com/apache/iceberg/pull/13066, with the
-// exception of testCtas, but adapted for the REST catalog. When that merges, we can directly use
-// those tests for the REST catalog as well by  adding to the parameters method there, to have a
-// single test class for table encryption.
+// TODO(smaheshwar-pltr): This test is taken from https://github.com/apache/iceberg/pull/13066, with
+// the exception of testCtas, but adapted for the REST catalog. When that merges, we can directly
+// use those tests for the REST catalog as well by  adding to the parameters method there, to have
+// a single test class for table encryption.
 public class TestRestTableEncryption extends CatalogTestBase {
   private static Map<String, String> appendCatalogEncryptionProperties(Map<String, String> props) {
     Map<String, String> newProps = Maps.newHashMap();