GH-3683: Support programmatic KMS client factories - #3785
GH-3683: Support programmatic KMS client factories#3785stevenwarejones wants to merge 6 commits into
Conversation
|
Hi @ggershinsky, @wgtmac, and @shangxinli — would you be willing to review this change? It implements the KmsClientFactory direction discussed in #3683, including explicit registration cleanup and isolation of the KMS and KEK caches. |
|
Sure, thanks |
| Objects.requireNonNull(configuration, "configuration"); | ||
| Objects.requireNonNull(kmsClientFactory, "kmsClientFactory"); | ||
| KmsClientCacheContext previous = | ||
| KMS_CLIENT_FACTORY_REGISTRATIONS.put(configuration, new KmsClientCacheContext(kmsClientFactory)); |
There was a problem hiding this comment.
What if a shell/notebook user adds some parameter to the config during a session?
They'll need to call removeKmsClientFactory and then setKmsClientFactory each time?
Are there usecases where config changes are hard to trace?
Maybe there is an alternative approach? (eg using something similar to the kms instance, say "parquet.encryption.kms.factory.instance")
There was a problem hiding this comment.
Mutating the same Configuration does not require deregistration. Registrations use object identity, so changes to its properties do not affect lookup.
I added a test that mutates the configuration after registration and verifies that the factory remains registered and receives the updated configuration. If a client has already been cached, changing a property will not recreate it; that is also the behavior of the existing reflective path. Calling setKmsClientFactory again replaces the registration and clears its caches, so a separate remove call is not needed.
A parquet.encryption.kms.factory.instance string would require a static ID-to-object registry because Configuration cannot contain the live factory itself. That would recreate the side channel this API is intended to remove and would have ambiguous behavior when a copied configuration is serialized to another JVM. Configuration copies therefore still need their own explicit registration.
There was a problem hiding this comment.
Sounds good, thanks. Indeed, the Configuration.equals implementation ignores the content.
There was a problem hiding this comment.
One more question on this. What if Spark/Flink/etc copies a Configuration into another object (today or in the future), some time during a session?
There was a problem hiding this comment.
Good point. That would be a copied Configuration, so the identity-based registration would not be found.
One option is to store an opaque registration ID in the configuration. That would support copies within the same JVM, but it could not carry the live factory into another JVM. Transparent cross-JVM support would require Spark/Flink-side integration to register or create the factory on each worker.
Would same-JVM copy support, with an explicit error when no local factory is registered, be the right scope here? Or would you suggest a different approach?
There was a problem hiding this comment.
Nope, this option sounds good to me. As for an error, can you check that an exception is indeed thrown? To make sure a dataframe is not written unencrypted silently, when no factory or kms class parameter are found.
| * | ||
| * @return a new or pre-built KMS client | ||
| */ | ||
| KmsClient createKmsClient(); |
There was a problem hiding this comment.
should we pass the current Configuration object here? It might have a useful input for creation of custom KMS clients.
There was a problem hiding this comment.
Done. I went slightly further and now pass the full creation context: Configuration, KMS instance ID, KMS instance URL, and access token. That lets an immutable or constructor-injected implementation select dependencies using the resolved values before initialize is called.
I also documented that each factory invocation must return a distinct, uninitialized client, since the factory is invoked separately for uncached token/KMS-instance combinations. Tests now verify the complete context and separate initialized clients for two access tokens.
There was a problem hiding this comment.
Cool, it certainly makes sense to add these three.
What changes?
KmsClientFactoryandKeyToolkit.setKmsClientFactory(Configuration, KmsClientFactory).KeyToolkitinitializes and caches.Configurationcopies using an opaque registration ID.removeKmsClientFactory(Configuration)and document the API and lifecycle in the README.Why?
Some KMS clients need live, constructor-injected dependencies such as SDK clients or dependency-injection state. Those clients cannot be created through a public no-argument constructor and should not need a static registry as a side channel.
Testing
KeyToolkitTestandTestKmsUrlReadsuites pass: 21 tests, 0 failures.Closes #3683