andygrove commented on code in PR #6031:
URL: https://github.com/apache/datafusion-comet/pull/6031#discussion_r4104635544
##########
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:
Thanks, deriving from the bucket's bridge closes this. The location bridges
now share the handle by construction. I also checked that the bucket bridge is
always `Read` in `create_store`, so the inherited mode matches what the factory
passed before.
--
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]
