Support encrypted table operations and server-side purge #5060
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,160 @@ | ||
| /* | ||
| * Licensed to the Apache Software Foundation (ASF) under one | ||
| * or more contributor license agreements. See the NOTICE file | ||
| * distributed with this work for additional information | ||
| * regarding copyright ownership. The ASF licenses this file | ||
| * to you under the Apache License, Version 2.0 (the | ||
| * "License"); you may not use this file except in compliance | ||
| * with the License. You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, | ||
| * software distributed under the License is distributed on an | ||
| * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | ||
| * KIND, either express or implied. See the License for the | ||
| * specific language governing permissions and limitations | ||
| * under the License. | ||
| */ | ||
| package org.apache.polaris.service.catalog.io; | ||
|
|
||
| import java.io.Closeable; | ||
| import java.io.IOException; | ||
| import java.util.List; | ||
| import java.util.Map; | ||
| import org.apache.iceberg.CatalogProperties; | ||
| import org.apache.iceberg.EncryptedKeyParser; | ||
| import org.apache.iceberg.TableMetadata; | ||
| import org.apache.iceberg.TableProperties; | ||
| import org.apache.iceberg.encryption.EncryptedKey; | ||
| import org.apache.iceberg.encryption.EncryptingFileIO; | ||
| import org.apache.iceberg.encryption.EncryptionManager; | ||
| import org.apache.iceberg.encryption.EncryptionUtil; | ||
| import org.apache.iceberg.encryption.KeyManagementClient; | ||
| import org.apache.iceberg.encryption.StandardEncryptionManager; | ||
| import org.apache.iceberg.io.FileIO; | ||
| import org.apache.iceberg.util.PropertyUtil; | ||
| import org.apache.polaris.core.entity.PolarisTaskConstants; | ||
| import org.apache.polaris.core.persistence.PolarisObjectMapperUtil; | ||
|
|
||
| /** Utilities for connecting Polaris file operations to Iceberg table encryption. */ | ||
| public final class PolarisEncryptionUtil { | ||
|
|
||
| private PolarisEncryptionUtil() {} | ||
|
|
||
| /** | ||
| * Copies the catalog KMS configuration and the table's wrapped encryption keys into a cleanup | ||
| * task. Cleanup tasks run after the catalog request has ended, so they cannot reconstruct these | ||
| * values from the live table operations object. | ||
| */ | ||
| public static void setupCleanupTaskEncryptionProperties( | ||
| Map<String, String> taskProperties, | ||
| Map<String, String> catalogProperties, | ||
| TableMetadata metadata) { | ||
| // Task properties start with table metadata. Always discard a table-controlled value for this | ||
| // reserved task field, including when the table is not encrypted. | ||
| taskProperties.remove(PolarisTaskConstants.ENCRYPTION_CONTEXT); | ||
|
|
||
| if (metadata == null | ||
| || metadata.properties().get(TableProperties.ENCRYPTION_TABLE_KEY) == null) { | ||
| return; | ||
| } | ||
|
|
||
| // Keep trusted catalog KMS settings and wrapped table keys in one task-only envelope. A custom | ||
| // KeyManagementClient may use additional catalog properties, but these values must not be | ||
| // merged into the FileIO and storage properties already resolved for the cleanup task. | ||
| CleanupTaskEncryptionContext context = | ||
| new CleanupTaskEncryptionContext( | ||
| catalogProperties, | ||
| metadata.encryptionKeys().stream().map(EncryptedKeyParser::toJson).toList()); | ||
| taskProperties.put( | ||
| PolarisTaskConstants.ENCRYPTION_CONTEXT, PolarisObjectMapperUtil.serialize(context)); | ||
| } | ||
|
|
||
| /** | ||
| * Wraps a task FileIO with Iceberg's EncryptingFileIO when the task represents an encrypted | ||
| * table. The KMS client is owned by the returned FileIO and is closed with it. | ||
| */ | ||
| public static FileIO encryptTaskFileIO(FileIO fileIO, Map<String, String> taskProperties) { | ||
| if (taskProperties.get(TableProperties.ENCRYPTION_TABLE_KEY) == null) { | ||
| return fileIO; | ||
| } | ||
|
|
||
| CleanupTaskEncryptionContext context = cleanupTaskEncryptionContext(taskProperties); | ||
| Map<String, String> catalogKmsProperties = context.kmsProperties(); | ||
| if (!catalogKmsProperties.containsKey(CatalogProperties.ENCRYPTION_KMS_TYPE) | ||
| && !catalogKmsProperties.containsKey(CatalogProperties.ENCRYPTION_KMS_IMPL)) { | ||
| throw new IllegalArgumentException( | ||
| "Cannot purge an encrypted table without encryption.kms-type or encryption.kms-impl"); | ||
| } | ||
|
|
||
| // Like CatalogUtil.loadFileIO, createKmsClient dynamically loads the configured implementation | ||
| // and initializes it with the supplied catalog properties. Credential resolution is delegated | ||
| // to the configured KeyManagementClient. | ||
| KeyManagementClient keyManagementClient = EncryptionUtil.createKmsClient(catalogKmsProperties); | ||
| boolean kmsOwnershipTransferred = false; | ||
| try { | ||
| EncryptionManager encryptionManager = | ||
| new CloseableStandardEncryptionManager( | ||
| encryptedKeys(context), | ||
| taskProperties.get(TableProperties.ENCRYPTION_TABLE_KEY), | ||
| PropertyUtil.propertyAsInt( | ||
| taskProperties, | ||
| TableProperties.ENCRYPTION_DEK_LENGTH, | ||
| TableProperties.ENCRYPTION_DEK_LENGTH_DEFAULT), | ||
| keyManagementClient); | ||
|
|
||
| // EncryptingFileIO closes an EncryptionManager when it is Closeable. The subclass preserves | ||
| // StandardEncryptionManager's type because Iceberg uses that type to decrypt manifest-list | ||
| // key metadata, while also closing the task-owned KMS client. | ||
| FileIO encryptingFileIO = EncryptingFileIO.combine(fileIO, encryptionManager); | ||
|
Contributor
There was a problem hiding this comment. The KMS client leaks if the raw FileIO throws during
Contributor
Author
There was a problem hiding this comment. Thanks for catching this. I agree that this is an Iceberg lifecycle issue rather than something specific to the Polaris cleanup path. I opened the following upstream fix: It updates EncryptingFileIO.close() to use try-with-resources so that a closeable EncryptionManager is still closed when the delegate FileIO throws. If both close operations fail, the EncryptionManager exception is retained as a suppressed exception. The upstream CI checks are passing. Since this was marked as non-blocking and the lifecycle is owned by EncryptingFileIO, I have not added a Polaris-specific workaround. |
||
| kmsOwnershipTransferred = true; | ||
| return encryptingFileIO; | ||
| } finally { | ||
| if (!kmsOwnershipTransferred) { | ||
| keyManagementClient.close(); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| private static CleanupTaskEncryptionContext cleanupTaskEncryptionContext( | ||
| Map<String, String> taskProperties) { | ||
| String contextJson = taskProperties.get(PolarisTaskConstants.ENCRYPTION_CONTEXT); | ||
| if (contextJson == null) { | ||
| throw new IllegalArgumentException("Missing encryption context in cleanup task properties"); | ||
| } | ||
| return PolarisObjectMapperUtil.deserialize(contextJson, CleanupTaskEncryptionContext.class); | ||
| } | ||
|
|
||
| private static List<EncryptedKey> encryptedKeys(CleanupTaskEncryptionContext context) { | ||
| return context.encryptedKeys().stream().map(EncryptedKeyParser::fromJson).toList(); | ||
| } | ||
|
|
||
| record CleanupTaskEncryptionContext( | ||
| Map<String, String> kmsProperties, List<String> encryptedKeys) { | ||
| CleanupTaskEncryptionContext { | ||
| kmsProperties = Map.copyOf(kmsProperties); | ||
| encryptedKeys = List.copyOf(encryptedKeys); | ||
| } | ||
| } | ||
|
|
||
| /** Makes the task-owned KMS client closeable without hiding Iceberg's standard manager type. */ | ||
| private static final class CloseableStandardEncryptionManager extends StandardEncryptionManager | ||
| implements Closeable { | ||
| private final KeyManagementClient keyManagementClient; | ||
|
|
||
| private CloseableStandardEncryptionManager( | ||
| List<EncryptedKey> keys, | ||
| String tableKeyId, | ||
| int dataKeyLength, | ||
| KeyManagementClient keyManagementClient) { | ||
| super(keys, tableKeyId, dataKeyLength, keyManagementClient); | ||
| this.keyManagementClient = keyManagementClient; | ||
| } | ||
|
|
||
| @Override | ||
| public void close() throws IOException { | ||
| keyManagementClient.close(); | ||
| } | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If we return here after the
removecall on line 56, it will contradict the logic implied by the method name 🤷If the remove is truly necessary, please rename the method. Suggestion:
setupCleanupTaskEncryptionProperties()There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The remove is necessary because
taskPropertiesstarts with user-configurable table metadata properties, whileENCRYPTION_CONTEXTis a reserved internal task field.We need to discard any table-controlled value even when the table is not encrypted, so that the task field is either absent or populated by Polaris from trusted catalog settings and wrapped keys.
Since the method may remove the field without adding it again, I agree that the current name is misleading. I renamed it to
setupCleanupTaskEncryptionProperties().