andygrove opened a new issue, #2532:
URL: https://github.com/apache/datafusion-ballista/issues/2532
**Is your feature request related to a problem or challenge? Please describe
what you are trying to do.**
Comet recently added a pluggable S3 credential provider SPI
(apache/datafusion-comet#4309), followed by built-in adapters
(apache/datafusion-comet#6023), per-location credentials
(apache/datafusion-comet#6031) and expiry-bounded credential reuse
(apache/datafusion-comet#6509). Ballista has the same gap, in a different
shape. This issue proposes a high-level design for discussion.
Today Ballista gets S3 credentials in two ways, both in
`ballista_core::object_store::CustomObjectStoreRegistry`:
1. **Ambient.** `AmazonS3Builder::from_env()` covers static env keys, web
identity (IRSA), ECS/EKS container credentials and IMDS. That is one identity
per process, shared by every session on the cluster.
2. **Static session keys.** `SET s3.access_key_id / s3.secret_access_key /
s3.session_token`, which travel in the session config to the scheduler and
executors.
That leaves no good answer for:
- **Multi-tenant clusters** where each session should read as a different
principal with credentials that refresh. Static session keys are a snapshot, so
a session token that expires mid-job fails the job. From reading the code (not
yet reproduced), issuing a new `SET` doesn't help either:
`DefaultSessionRuntimeCache` keeps the registry that was built from the
session's first task, so an executor that already cached the session's runtime
keeps the old keys.
- **Vended credentials**, where a catalog (Iceberg REST, Unity, Polaris,
Lake Formation) or an in-house broker issues short-lived and often path-scoped
credentials.
- **Credential sources `object_store` doesn't cover**, such as profiles,
SSO, `credential_process` and plain `AssumeRole`.
The only extension point is to replace the whole `ObjectStoreRegistry`
through `override_runtime_producer` and `override_session_builder`. That means
reimplementing store construction and shipping custom scheduler and executor
binaries, just to change where three strings come from.
Two existing problems in the registry are worth fixing as part of this work:
- `get_store` builds a new `AmazonS3` on every call, each with its own
credential cache, so nothing is shared between scans. Comet hit STS throttling
from the same pattern (apache/datafusion-comet#6025).
- Static secrets ride in the session config and task props, and both
`get_store` and the executor's task launch print them at `trace` level.
**Describe the solution you'd like**
### What carries over from Comet, and what doesn't
| | Comet | Ballista |
| --- | --- | --- |
| Where the provider runs | JVM, called from native code over JNI | Rust, in
process |
| How a provider is loaded | Class name in config, loaded reflectively | No
runtime class loading in Rust, so providers are compiled in and selected by
name |
| Who touches S3 | Executors only (the driver stays on Hadoop S3A) | Client
(schema inference), scheduler (listing, footer stats) and executors (reads and
writes) |
| Tenancy | Mostly one application per JVM | Long-lived shared executors
serving many sessions |
| Property bag | `fs.s3a.*` / catalog properties | Session config, already
propagated to scheduler and executors |
The shape of the SPI, the "returns or throws" contract, config-driven
activation, expiry-bounded reuse and the location-scoped extension all carry
over. The JNI dispatcher, handles and Hadoop/AWS SDK adapters don't.
### 1. A narrow provider trait in `ballista-core`
A sketch, not a final API:
```rust
#[async_trait]
pub trait S3CredentialProvider: Debug + Send + Sync {
/// Returns credentials for the request, or an error. There is no
/// "I don't know" value and no fall-through to another source.
async fn credentials(&self, ctx: &S3CredentialContext) ->
Result<S3Credentials>;
}
#[non_exhaustive]
pub struct S3CredentialContext {
pub bucket: String,
pub path: String, // bucket root unless location scoping is in
use
pub mode: S3AccessMode, // Read | Write
}
pub struct S3Credentials {
pub access_key_id: String,
pub secret_access_key: String,
pub session_token: Option<String>,
pub expiry: Option<SystemTime>,
}
pub trait S3CredentialProviderFactory: Send + Sync {
/// `options` is the session's `s3.credential.option.*` map.
fn create(&self, options: &HashMap<String, String>)
-> Result<Arc<dyn S3CredentialProvider>>;
}
```
Factories are registered by name when a process is built, next to the
existing `override_*` hooks on `SchedulerConfig` and `ExecutorProcessConfig`.
Someone with a custom provider still builds their own binaries, but writes one
small trait impl instead of a registry.
### 2. Activation through session config
```sql
SET s3.credential.provider = 'acme';
SET s3.credential.option.tenant = 'finance';
```
- `s3.credential.provider` names a registered factory. Unset keeps today's
behavior exactly.
- `s3.credential.option.<key>` is a free-form bag handed to the factory. It
rides the existing `S3Options` extension, so it reaches the scheduler and every
executor with no new plumbing.
- Only the provider name and its options travel. The credentials themselves
are produced where they are used and never enter session config, plans or task
props.
- Naming an unregistered provider is an error on first use. It must not fall
back to the ambient chain, since that would silently read as the executor's own
role.
### 3. A bridge that owns caching
One internal type implements `object_store::CredentialProvider` on top of
the trait and takes the lessons from apache/datafusion-comet#6509:
- Reuse a credential until 5 minutes before its reported expiry, then ask
again. A credential with no expiry is not kept.
- Single-flight: one provider call at a time per bridge, with waiters
sharing the outcome (credential or error).
- A process-wide store cache keyed by `(bucket, store config, provider name,
options)`. Sessions with identical config share one store and one credential.
Sessions with different options are isolated. This also fixes the "new store
per `get_store`" problem for the ambient path.
- The executor's session runtime cache needs the `s3.*` config in its key
(or the registry needs to stop capturing the first task's `S3Options`), so a
session that changes provider options is not served from a stale runtime.
### 4. A built-in `process` provider for stock binaries
Most users run the published binaries and images, so there has to be a path
that needs no Rust. The proposal is one built-in provider that runs an external
command and parses the standard AWS
[`credential_process`](https://docs.aws.amazon.com/sdkref/latest/guide/feature-process-credentials.html)
JSON (`AccessKeyId`, `SecretAccessKey`, `SessionToken`, `Expiration`).
- The command comes from **process-level** config on the scheduler and
executor (a CLI flag or env var), never from session config. A client must not
be able to choose what executors execute.
- The request context (bucket, path, mode, and the session's options) is
passed to the command as environment variables.
- Existing `credential_process` helpers work unchanged, since they ignore
the extra variables. That also gives stock binaries profile, SSO and broker
support by way of whatever helper the operator installs.
This plays the role of Comet's built-in adapters: a config-only path for the
common cases, built on the same SPI.
### 5. The same path in all three processes
The provider resolves through the same registry code wherever S3 is touched:
`state_with_s3_support()` on a Rust client, `session_state_with_s3_support` on
the scheduler, and `runtime_env_with_s3_support` on the executor. A session
that selects a provider gets the same identity for schema inference, planning
and execution.
### 6. Location-scoped credentials (later phase)
`get_credential()` receives no request path and `get_store` sees only
`s3://bucket`, so one store means one credential per bucket. That is the same
limit Comet addressed in apache/datafusion-comet#6031. The same answer should
work here: an opt-in extension trait that returns the bucket's policy
locations, and a wrapper `ObjectStore` that routes each request to the store
for the longest location covering its path, refreshing the list once on a 403.
This is additive and doesn't change the base trait.
### Security notes
- Provider options are client-asserted. A provider that vends per-tenant
credentials has to authenticate what it is given (for example an option
carrying a token the broker validates), not trust a bare `tenant=finance`.
- Option values may be secrets, so they should be redacted in `Debug` output
and in the `trace` logs that print session config. `SHOW ALL` needs more care,
because the same `entries()` call that feeds it also drives config propagation
in `to_key_value_pairs`.
- A provider failure fails the task with the provider's error. No downgrade
to another identity.
### Suggested phases
1. Trait, named registration, bridge with expiry-bounded reuse, store cache,
activation config, docs, and a MinIO test (the `examples/tests/object_store.rs`
harness already runs MinIO) where the executor's ambient credentials are wrong,
so a successful read proves the provider was used.
2. Built-in `process` provider.
3. Location-scoped credentials.
### Open questions
- **Authenticated identity.** The Flight SQL frontend already resolves an
`Identity` per session. Should the scheduler inject it into the credential
context as a field clients can't set, so providers can vend per-user
credentials without trusting options?
- **Shared statistics cache.** `share_file_statistics_cache` shares file
statistics across sessions on the scheduler. With per-session credentials, a
cache hit could serve one tenant metadata for files it can't read. Should the
cache be partitioned by credential config when a provider is active?
- **Per-bucket selection.** Is `s3.bucket.<name>.credential.provider` needed
up front, or can a single provider dispatch on `ctx.bucket`? The existing
`s3.region` and `s3.endpoint` are also global today.
- **Python client.** `register_object_store` with datafusion-python's
`AmazonS3` is a separate path for client-side schema inference. Out of scope
here, but the mismatch should at least be documented.
- **Sharing code with Comet.** The expiry-bounded bridge and the
location-scoped store are engine-neutral. Is it worth pulling them into a
shared crate instead of porting them?
**Describe alternatives you've considered**
- **Keep `override_runtime_producer` as the answer.** It works today, but
every user reimplements store construction, caching and expiry handling, and
there is no path for stock binaries.
- **Client-resolved credentials through `SET`.** Simple and already
supported, but they are a snapshot that can't refresh during a long job, and
the secrets travel in config, plans and logs.
- **Dynamically loaded plugins (`dlopen`).** The closest analogue to Comet's
class-name loading. Rust has no stable ABI, so this needs a C ABI around an
async interface plus version pinning between plugin and host. A lot of surface
for what `process` covers.
- **An HTTP credential endpoint as the built-in.** Reasonable, and could be
added later as a second built-in. `process` comes first because the contract is
an existing AWS standard with existing helpers, and it adds no listener or SSRF
surface.
- **Python callbacks.** Would only help the client, since the scheduler and
executors are Rust processes.
**Additional context**
Non-goals for this proposal: Azure and GCS (the same shape could follow),
Iceberg or other table-format integration, and changing the default ambient
behavior.
Comet references: the [design
notes](https://github.com/apache/datafusion-comet/blob/main/docs/source/contributor-guide/s3-credential-provider-design.md)
and [user
guide](https://github.com/apache/datafusion-comet/blob/main/docs/source/user-guide/latest/s3-credential-providers.md).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]