andygrove commented on code in PR #6031:
URL: https://github.com/apache/datafusion-comet/pull/6031#discussion_r4098072028
##########
native/core/src/parquet/objectstore/s3.rs:
##########
@@ -97,48 +101,153 @@ pub fn create_store(
store: "S3",
source: format!("CometS3CredentialBridge init failed for
{bucket}: {e}").into(),
})?;
- builder.with_credentials(Arc::new(bridge))
+ let locations =
+ bridge
+ .policy_locations()
+ .map_err(|e| object_store::Error::Generic {
+ store: "S3",
+ source: format!("Failed to get policy locations for
{bucket}: {e}").into(),
+ })?;
+ if let Some(locations) = locations {
+ let template = S3StoreTemplate::new(url, configs, bucket)?;
+ let store =
+ location_scoped_store(template, provider_class, bucket,
bridge, locations)?;
+ return Ok((Box::new(store), path));
+ }
+ S3Credentials::Provider(Arc::new(bridge))
}
None => {
match get_runtime().block_on(build_credential_provider(configs,
bucket, min_ttl))? {
- Some(provider) => builder.with_credentials(Arc::new(provider)),
- None => builder.with_skip_signature(true),
+ Some(provider) => S3Credentials::Provider(Arc::new(provider)),
+ None => S3Credentials::SkipSignature,
}
}
};
- let s3_configs = extract_s3_config_options(configs, bucket);
- debug!("S3 configs for bucket {bucket}: {s3_configs:?}");
+ let object_store = S3StoreTemplate::new(url, configs,
bucket)?.build(credentials)?;
- // When using the default AWS S3 endpoint (no custom endpoint configured),
a valid region
- // is required. If no region is explicitly configured, attempt to
auto-resolve it by
- // making a HeadBucket request to determine the bucket's region.
- if !s3_configs.contains_key(&AmazonS3ConfigKey::Endpoint)
- && !s3_configs.contains_key(&AmazonS3ConfigKey::Region)
- {
- let region = get_runtime()
- .block_on(resolve_bucket_region(bucket))
- .map_err(|e| object_store::Error::Generic {
- store: "S3",
- source: format!(
- "Failed to resolve region: {e}. If '{bucket}' is on a
non-AWS S3-compatible \
- service, set fs.s3a.endpoint (and optionally
fs.s3a.endpoint.region, \
- fs.s3a.path.style.access) or the per-bucket variants \
- fs.s3a.bucket.{bucket}.endpoint[.region] so Comet skips
the AWS HEAD probe."
- )
- .into(),
- })?;
- debug!("resolved region: {region:?}");
- builder = builder.with_config(AmazonS3ConfigKey::Region,
region.to_string());
+ Ok((Box::new(object_store), path))
+}
+
+/// How a store built from an [`S3StoreTemplate`] signs its requests.
+enum S3Credentials {
+ Provider(AwsCredentialProvider),
+ SkipSignature,
+}
+
+/// Builder settings shared by every store for one bucket. Creating a template
may block on a
+/// region lookup; building a store from it does not, so location-scoped
stores can be built from
+/// async code on a Tokio worker.
+struct S3StoreTemplate {
+ url: String,
+ region: Option<String>,
+ s3_configs: HashMap<AmazonS3ConfigKey, String>,
+}
+
+impl S3StoreTemplate {
+ fn new(
+ url: &Url,
+ configs: &HashMap<String, String>,
+ bucket: &str,
+ ) -> Result<Self, object_store::Error> {
+ let s3_configs = extract_s3_config_options(configs, bucket);
+ debug!("S3 configs for bucket {bucket}: {s3_configs:?}");
+
+ // When using the default AWS S3 endpoint (no custom endpoint
configured), a valid region
+ // is required. If no region is explicitly configured, attempt to
auto-resolve it by
+ // making a HeadBucket request to determine the bucket's region.
+ let region = if !s3_configs.contains_key(&AmazonS3ConfigKey::Endpoint)
+ && !s3_configs.contains_key(&AmazonS3ConfigKey::Region)
+ {
+ let region = get_runtime()
+ .block_on(resolve_bucket_region(bucket))
+ .map_err(|e| object_store::Error::Generic {
+ store: "S3",
+ source: format!(
+ "Failed to resolve region: {e}. If '{bucket}' is on a
non-AWS S3-compatible \
+ service, set fs.s3a.endpoint (and optionally
fs.s3a.endpoint.region, \
+ fs.s3a.path.style.access) or the per-bucket variants \
+ fs.s3a.bucket.{bucket}.endpoint[.region] so Comet
skips the AWS HEAD probe."
+ )
+ .into(),
+ })?;
+ debug!("resolved region: {region:?}");
+ Some(region)
+ } else {
+ None
+ };
+
+ Ok(Self {
+ url: url.to_string(),
+ region,
+ s3_configs,
+ })
}
- for (key, value) in s3_configs {
- builder = builder.with_config(key, value);
+ fn build(&self, credentials: S3Credentials) -> Result<AmazonS3,
object_store::Error> {
+ let builder = AmazonS3Builder::new()
+ .with_url(self.url.clone())
+ .with_allow_http(true);
+ let mut builder = match credentials {
+ S3Credentials::Provider(provider) =>
builder.with_credentials(provider),
+ S3Credentials::SkipSignature => builder.with_skip_signature(true),
+ };
+ if let Some(region) = &self.region {
+ builder = builder.with_config(AmazonS3ConfigKey::Region,
region.clone());
+ }
+ for (key, value) in &self.s3_configs {
+ builder = builder.with_config(*key, value.clone());
+ }
+ builder.build()
}
+}
- let object_store = builder.build()?;
+/// Builds the store for a `CometS3LocationScopedCredentialProvider`. `bridge`
was created on this
+/// thread, which registered the provider; it is kept to fetch the locations
again after a 403.
+/// Each location's bridge is created on first use, often on a Tokio worker,
and reuses that
+/// registration.
+fn location_scoped_store(
+ template: S3StoreTemplate,
+ provider_class: &str,
+ bucket: &str,
+ bridge: CometS3CredentialBridge,
+ locations: Vec<String>,
+) -> Result<LocationScopedObjectStore, object_store::Error> {
+ let source_bucket = bucket.to_string();
+ let source: LocationSource = Arc::new(move || {
+ let locations = bridge
+ .policy_locations()
+ .map_err(|e| object_store::Error::Generic {
+ store: "S3",
+ source: format!("Failed to get policy locations for
{source_bucket}: {e}").into(),
+ })?;
+ locations.ok_or_else(|| object_store::Error::Generic {
+ store: "S3",
+ source: format!("The provider for {source_bucket} stopped
returning policy locations")
+ .into(),
+ })
+ });
+
+ let provider_class = provider_class.to_string();
+ let factory_bucket = bucket.to_string();
+ let factory: LocationStoreFactory = Arc::new(move |credential_path: &str| {
+ let bridge = CometS3CredentialBridge::new(
+ provider_class.as_str(),
+ factory_bucket.as_str(),
+ factory_bucket.as_str(),
+ credential_path,
+ AccessMode::Read,
+ &HashMap::new(),
Review Comment:
Each location's bridge goes through `CometS3CredentialBridge::new` again, so
it only lands on the bucket's registration because this `&HashMap::new()`
matches what `create_store` passes. #6023 changes `create_store` to forward the
`fs.s3a.*` map. In a trial merge of the two, `s3.rs` only conflicts in
`create_store`, so this line survives unchanged. After that,
`ensureInitialized` would miss the existing key and create a second provider
with an empty map. It would do that on whatever thread first reads the
location, which is often a Tokio worker with no context class loader, so a
vendor jar that comes in through `--jars` can't be loaded when Comet is on
`extraClassPath`.
Could the location bridges be derived from `bridge` instead, reusing its
handle and bucket string and only creating a new path string? They would then
share the registration by construction, whatever `create_store` passes, and
we'd skip an `ensureInitialized` round trip per location.
This doesn't need to hold up this PR. We can file a follow-on issue for it,
as long as the fix lands before or together with #6023.
##########
native/core/src/parquet/objectstore/location_scoped.rs:
##########
@@ -0,0 +1,879 @@
+// 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.
+
+//! The object store for a `CometS3LocationScopedCredentialProvider`.
+//!
+//! `object_store::CredentialProvider::get_credential` receives no request
path, so one S3 store
+//! presents one credential. A bucket whose policies differ by location needs
a store per location
+//! and something that picks the right one for each request.
[`LocationScopedObjectStore`] is
+//! registered once per bucket in place of a plain S3 store. It keeps the
provider's policy
+//! locations and serves each request with the store of the longest location
that covers the
+//! request's path, compared one path segment at a time. The bucket root is an
implicit location
+//! that covers every other path. A location's store is built the first time a
request needs it and
+//! kept for the life of this store.
+//!
+//! The locations are a snapshot, so a 403 can mean a location was added or
removed after it was
+//! taken. A read (`get_opts` or `get_ranges`) that gets a 403 fetches the
locations again, unless
+//! another read already tried since this one was routed, and retries once if
its path now routes to
+//! a different location; otherwise the 403 is returned. A failed fetch fails
every read that shared
+//! it. Other operations route by path without retrying, because Comet only
reads through this store.
+
+use std::collections::HashMap;
+use std::fmt;
+use std::ops::Range;
+use std::sync::{Arc, PoisonError, RwLock};
+
+use async_trait::async_trait;
+use bytes::Bytes;
+use futures::stream::{self, BoxStream, StreamExt, TryStreamExt};
+use object_store::path::Path;
+use object_store::{
+ CopyOptions, Error, GetOptions, GetResult, ListResult, MultipartUpload,
ObjectMeta,
+ ObjectStore, ObjectStoreExt, PutMultipartOptions, PutOptions, PutPayload,
PutResult,
+ RenameOptions, Result,
+};
+use tokio::sync::Mutex;
+
+const STORE: &str = "LocationScopedS3";
+
+/// The credential path used for paths that no returned location covers.
+const ROOT_CREDENTIAL_PATH: &str = "/";
+
+/// Fetches the provider's current policy locations for the bucket.
+pub(crate) type LocationSource = Arc<dyn Fn() -> Result<Vec<String>> + Send +
Sync>;
+
+/// Builds the store for one location from the path passed to
`getCredentialsForPath`.
+pub(crate) type LocationStoreFactory =
+ Arc<dyn Fn(&str) -> Result<Arc<dyn ObjectStore>> + Send + Sync>;
+
+/// One snapshot of the provider's locations.
+struct LocationIndex {
+ /// Counts refresh attempts, including failed ones.
+ generation: u64,
+ /// Canonical location to credential path: the location as the provider
returned it, with a
+ /// leading slash.
+ locations: Arc<HashMap<Path, String>>,
+ /// Why the attempt that produced this snapshot failed, when it kept the
previous locations.
+ failure: Option<Arc<str>>,
+}
+
+impl LocationIndex {
+ fn new(generation: u64, locations: Vec<String>) -> Result<Self> {
+ let mut index = HashMap::with_capacity(locations.len());
+ for location in locations {
+ // Request paths are percent-decoded the same way, so both sides
compare as raw keys.
+ let canonical = Path::from_url_path(&location).map_err(|e|
Error::Generic {
+ store: STORE,
+ source: format!("Invalid policy location {location:?}:
{e}").into(),
+ })?;
+ // A duplicate keeps the first spelling, which is the path the
provider is given.
+ index
+ .entry(canonical)
+ .or_insert_with(|| credential_path(&location));
+ }
+ Ok(Self {
+ generation,
+ locations: Arc::new(index),
+ failure: None,
+ })
+ }
+
+ /// The snapshot after a failed refresh: the same locations under the next
generation, with the
+ /// failure recorded for the requests routed before it.
+ fn after_failure(&self, failure: &Error) -> Self {
+ Self {
+ generation: self.generation + 1,
+ locations: Arc::clone(&self.locations),
+ failure: Some(failure.to_string().into()),
+ }
+ }
+
+ /// Returns the credential path of the longest location that covers `path`.
+ fn route(&self, path: &Path) -> &str {
+ let mut longest = self
+ .locations
+ .get(&Path::default())
+ .map_or(ROOT_CREDENTIAL_PATH, String::as_str);
+ let mut prefix = Path::default();
+ for part in path.parts() {
+ prefix = prefix.join(part);
+ if let Some(credential_path) = self.locations.get(&prefix) {
+ longest = credential_path;
+ }
+ }
+ longest
+ }
+}
+
+fn credential_path(location: &str) -> String {
+ if location.starts_with('/') {
+ location.to_string()
+ } else {
+ format!("/{location}")
+ }
+}
+
+fn is_forbidden(err: &Error) -> bool {
Review Comment:
The refresh only fires on a 403 from S3. A provider with no policy for the
path it's asked about will usually throw instead, and `get_credential` turns
that into a `Generic` error, which skips the refresh. I tried two cases with a
fake location store that returns that error, and both fetched the locations
zero times across three reads. In the first, the provider has no bucket-wide
policy and throws for `/`, then adds `warehouse/finance` mid-job. Reads under
it route to `/` and keep failing, so the new location is never picked up. In
the second, the provider folds `warehouse/finance` into `warehouse` and throws
for the old location. Later reads under it keep failing, while a freshly built
store serves the same path from `/warehouse`. The store lives as long as the
executor, so both stay broken until a restart.
The user guide says a location added while a job runs is picked up, and
suggests keying the provider's cache on the location, which leads naturally to
that throw. Should a credential failure from a location's store trigger the
same refresh as a 403? Or should the contract require `getCredentialsForPath`
to answer for `/` and for locations the provider has dropped, and the guide say
so?
This one can also be a follow-on issue rather than blocking the PR.
--
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]
