Skip to content

GH-3683: Support programmatic KMS client factories - #3785

Open
stevenwarejones wants to merge 6 commits into
apache:masterfrom
stevenwarejones:stevenwarejones_kms_client_factory
Open

GH-3683: Support programmatic KMS client factories#3785
stevenwarejones wants to merge 6 commits into
apache:masterfrom
stevenwarejones:stevenwarejones_kms_client_factory

Conversation

@stevenwarejones

@stevenwarejones stevenwarejones commented Sep 9, 2026

Copy link
Copy Markdown

What changes?

  • Add a public KmsClientFactory and KeyToolkit.setKmsClientFactory(Configuration, KmsClientFactory).
  • Pass the current configuration and resolved KMS instance ID, URL, and access token to the factory. Each invocation returns a distinct client that KeyToolkit initializes and caches.
  • Prefer the registered factory while keeping class-name reflection as the fallback.
  • Preserve registrations across same-JVM Configuration copies using an opaque registration ID.
  • Isolate KMS-client and KEK caches by factory registration, and isolate write-side KEKs by KMS instance within a registration.
  • Add 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

  • Focused KeyToolkitTest and TestKmsUrlRead suites pass: 21 tests, 0 failures.
  • Regression coverage verifies factory registration and lifecycle, copied configurations, cache isolation between registrations, and write-side KEK isolation between KMS instances.
  • The new KMS-instance cache test fails on the prior PR head because the second KMS client is not called, and passes with this change.

Closes #3683

@stevenwarejones

stevenwarejones commented Sep 9, 2026

Copy link
Copy Markdown
Author

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.

@ggershinsky

Copy link
Copy Markdown
Contributor

Sure, thanks

Objects.requireNonNull(configuration, "configuration");
Objects.requireNonNull(kmsClientFactory, "kmsClientFactory");
KmsClientCacheContext previous =
KMS_CLIENT_FACTORY_REGISTRATIONS.put(configuration, new KmsClientCacheContext(kmsClientFactory));

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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")

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds good, thanks. Indeed, the Configuration.equals implementation ignores the content.

@ggershinsky ggershinsky Sep 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.

*
* @return a new or pre-built KMS client
*/
KmsClient createKmsClient();

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should we pass the current Configuration object here? It might have a useful input for creation of custom KMS clients.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cool, it certainly makes sense to add these three.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Allow supplying a KmsClient instance/supplier instead of only a reflectively-instantiated class name

2 participants