sunchao commented on code in PR #6031:
URL: https://github.com/apache/datafusion-comet/pull/6031#discussion_r4051461775
##########
native/core/src/parquet/parquet_support.rs:
##########
@@ -735,11 +752,70 @@ type ObjectStoreCache =
RwLock<HashMap<ObjectStoreCacheKey, Arc<dyn ObjectStore>
/// (e.g. `fs.s3a.access.key` / `fs.s3a.secret.key`) produce a different
`config_hash` when
/// those values change, which causes a new store to be created and inserted
under the new
/// key; the old entry is harmlessly superseded.
+///
+/// ## Scope-aware entries (S3 SPI path)
+///
+/// When the S3 `CometS3CredentialProvider` implements the
+/// `CometS3ScopedCredentialProvider` sub-interface, its
`getPolicyLocationsFor` report is
+/// carried alongside the store as a `ScopeEntry`. Multiple entries can
coexist under one key
+/// so that distinct policy scopes on the same bucket each get their own
credentials without
+/// evicting each other. Base (non-scoped) providers land as a single entry
with empty
+/// `prefixes`, preserving the historical single-entry-per-bucket behavior.
The advisory
+/// nature of the scope hint is enforced by wrapping each SPI store in
+/// [`RetryOn403ObjectStore`], whose rebuild closure re-fires the SPI with the
failing path
+/// in context and *appends* a new scope entry to this cache — leaving any
pre-existing
+/// entries intact so disjoint scoped stores on the same bucket continue to
coexist. See
+/// `docs/source/contributor-guide/s3-credential-provider-design.md` for the
full contract.
fn object_store_cache() -> &'static ObjectStoreCache {
static CACHE: OnceLock<ObjectStoreCache> = OnceLock::new();
CACHE.get_or_init(|| RwLock::new(HashMap::new()))
}
+/// True when `path` is covered by `prefixes`. Empty `prefixes` is a catchall.
+fn path_covered(path: &str, prefixes: &[String]) -> bool {
+ if prefixes.is_empty() {
+ return true;
+ }
+ prefixes.iter().any(|p| path.starts_with(p.as_str()))
+}
+
+/// Length of the longest prefix in `prefixes` that covers `path`. Empty
`prefixes` is a
+/// catchall and returns 0 (so any covering scoped entry wins the tie-break);
returns `None`
+/// when nothing matches.
+fn longest_covering_prefix_len(path: &str, prefixes: &[String]) ->
Option<usize> {
+ if prefixes.is_empty() {
+ return Some(0);
+ }
+ prefixes
+ .iter()
+ .filter(|p| path.starts_with(p.as_str()))
+ .map(|p| p.len())
Review Comment:
### Correctness
[P2] Could both sides be canonicalized to the bucket-relative representation
required by `CometS3ScopedCredentialProvider` before this comparison?
`requested_path` is taken directly from `url.path()`, so it starts with `/`,
while the Java contract explicitly requires prefixes without a leading slash. A
conforming hint such as `warehouse/table` therefore never covers
`/warehouse/table/file`. With only that scoped entry present, the exact lookup
helper produced zero hits in 1,000 requests, causing a new store and HTTP pool
to be constructed on each preparation. The absolute S3 URI used by the new
Minio test also never matches. The comparison should enforce the documented
path-component boundary too, so `table` cannot match `table_extra`. Please
exercise the Java contract's prefix format through the native cache.
##########
native/core/src/parquet/objectstore/retry.rs:
##########
@@ -0,0 +1,628 @@
+// 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.
+
+//! Correctness safety net for the scope-hint SPI.
+//!
+//! # Why
+//!
+//! `CometS3ScopedCredentialProvider::getPolicyLocationsFor` is *advisory*:
vendors are
+//! encouraged to report narrower scopes than the policy grants, so a request
outside the
+//! reported scope can still legitimately fail with 403 at S3. The scope hint
is a latency
+//! optimization that lets Comet share a bridge across paths inside one scope,
saving one
+//! JVM round-trip per get. S3 itself remains authoritative.
+//!
+//! This wrapper is the mechanism that keeps that split honest: on a 403 from
a cached
+//! (scope-bound) store, we invoke a caller-provided rebuild closure exactly
once, passing
+//! it the `Path` of the request that failed. The rebuild constructs a fresh
bridge bound
+//! to that path, re-fires the SPI against it (so the vendor answers
`getPolicyLocationsFor`
+//! *for the actual failing request*), and appends the resulting scope entry
to the outer
+//! `object_store` registry cache — alongside any pre-existing entries, not
replacing them.
+//! We then retry the same operation against the rebuilt store; a second 403
propagates
+//! unchanged. Because the new entry is inserted with the vendor's fresh
(narrower) scope
+//! rather than an all-covering catchall, disjoint scoped stores on the same
bucket can
+//! continue to coexist after a 403 recovery.
+//!
+//! # Non-goals
+//!
+//! - Not a generic retry policy. Only `Error::PermissionDenied` (403) is
intercepted; every
+//! other error (including 401/`Unauthenticated`) passes through as-is.
+//! - Not an infinite retry. Exactly one rebuild per wrapper instance, exactly
one retry per
+//! operation.
+//! - Not a stream-level retry. `list`, `list_with_offset`, and
`delete_stream` return
+//! `BoxStream`s whose per-item errors are surfaced as-is; wrapping them
would require
+//! materializing the stream. Comet's parquet path first hits 403 at
`get_opts`/`get_ranges`
+//! which are covered; the rebuild there populates the shared cache and
subsequent stream
+//! requests use the newly-inserted scope entry.
+//! - Not applied to `put_multipart_opts`: a partial multipart upload cannot
be transparently
+//! retried, so a 403 mid-upload propagates for the caller to handle.
+
+use std::fmt;
+use std::ops::Range;
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use bytes::Bytes;
+use futures::stream::BoxStream;
+use log::debug;
+use object_store::path::Path;
+use object_store::{
+ CopyOptions, Error, GetOptions, GetResult, ListResult, MultipartUpload,
ObjectMeta,
+ ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult,
RenameOptions, Result,
+};
+use once_cell::sync::OnceCell;
+
+/// Rebuild function contract: construct a fresh backing `ObjectStore` scoped
for the failing
+/// request and register it in the outer `object_store` registry cache. Called
at most once
+/// per wrapper.
+///
+/// The `Option<&Path>` argument is the location that triggered the 403
(source path for
+/// copy/rename). The rebuild closure passes this into the fresh bridge so the
SPI's
+/// `getPolicyLocationsFor` reflects the actual failing request rather than
the path baked in
+/// at construction time. `None` is reserved for callers that intercept a 403
without a path
+/// (none of the current retry sites) and lets the closure fall back to its
pre-baked path.
+pub type RebuildFn =
+ Arc<dyn Fn(Option<&Path>) -> Result<Arc<dyn ObjectStore>> + Send + Sync +
'static>;
+
+/// Wraps an `Arc<dyn ObjectStore>` so a single 403 rebuilds the store once
and retries.
+///
+/// See the module-level doc-comment for rationale and non-goals.
+pub struct RetryOn403ObjectStore {
+ inner: Arc<dyn ObjectStore>,
+ rebuild: RebuildFn,
+ /// Populated on the first 403 we successfully recover from. Cached so a
subsequent
+ /// request against this same wrapper skips straight to the rebuilt store,
and a 403 there
+ /// is treated as authoritative.
+ rebuilt: OnceCell<Arc<dyn ObjectStore>>,
+}
+
+impl RetryOn403ObjectStore {
+ pub fn new(inner: Arc<dyn ObjectStore>, rebuild: RebuildFn) -> Self {
+ Self {
+ inner,
+ rebuild,
+ rebuilt: OnceCell::new(),
+ }
+ }
+
+ /// Return the store that should service *this* request: the post-rebuild
store when we
+ /// have one, else the original.
+ fn current(&self) -> Arc<dyn ObjectStore> {
+ self.rebuilt
+ .get()
+ .cloned()
+ .unwrap_or_else(|| Arc::clone(&self.inner))
Review Comment:
### Correctness
[P2] Could the rebuilt store be selected by request path without permanently
replacing the backing store for every operation on this wrapper? The planner
constructs one store from the first file and uses it for all files in a Parquet
partition. With an A-only initial store and a B-only rebuilt store, reading A1,
B1, then A2 returns success, success, then 403. The original store still
authorizes A2, but `current()` now routes it to B and `already_rebuilt()`
suppresses recovery. Appending a new scope entry does not repair the original
cache entry or scans already holding this wrapper. I reproduced that sequence
using these exact methods and synthetic scoped stores. Please retain
path-specific routing and bound retries per operation, with A/B/A and
three-disjoint-scope coverage.
##########
native/core/src/parquet/objectstore/retry.rs:
##########
@@ -0,0 +1,628 @@
+// 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.
+
+//! Correctness safety net for the scope-hint SPI.
+//!
+//! # Why
+//!
+//! `CometS3ScopedCredentialProvider::getPolicyLocationsFor` is *advisory*:
vendors are
+//! encouraged to report narrower scopes than the policy grants, so a request
outside the
+//! reported scope can still legitimately fail with 403 at S3. The scope hint
is a latency
+//! optimization that lets Comet share a bridge across paths inside one scope,
saving one
+//! JVM round-trip per get. S3 itself remains authoritative.
+//!
+//! This wrapper is the mechanism that keeps that split honest: on a 403 from
a cached
+//! (scope-bound) store, we invoke a caller-provided rebuild closure exactly
once, passing
+//! it the `Path` of the request that failed. The rebuild constructs a fresh
bridge bound
+//! to that path, re-fires the SPI against it (so the vendor answers
`getPolicyLocationsFor`
+//! *for the actual failing request*), and appends the resulting scope entry
to the outer
+//! `object_store` registry cache — alongside any pre-existing entries, not
replacing them.
+//! We then retry the same operation against the rebuilt store; a second 403
propagates
+//! unchanged. Because the new entry is inserted with the vendor's fresh
(narrower) scope
+//! rather than an all-covering catchall, disjoint scoped stores on the same
bucket can
+//! continue to coexist after a 403 recovery.
+//!
+//! # Non-goals
+//!
+//! - Not a generic retry policy. Only `Error::PermissionDenied` (403) is
intercepted; every
+//! other error (including 401/`Unauthenticated`) passes through as-is.
+//! - Not an infinite retry. Exactly one rebuild per wrapper instance, exactly
one retry per
+//! operation.
+//! - Not a stream-level retry. `list`, `list_with_offset`, and
`delete_stream` return
+//! `BoxStream`s whose per-item errors are surfaced as-is; wrapping them
would require
+//! materializing the stream. Comet's parquet path first hits 403 at
`get_opts`/`get_ranges`
+//! which are covered; the rebuild there populates the shared cache and
subsequent stream
+//! requests use the newly-inserted scope entry.
+//! - Not applied to `put_multipart_opts`: a partial multipart upload cannot
be transparently
+//! retried, so a 403 mid-upload propagates for the caller to handle.
+
+use std::fmt;
+use std::ops::Range;
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use bytes::Bytes;
+use futures::stream::BoxStream;
+use log::debug;
+use object_store::path::Path;
+use object_store::{
+ CopyOptions, Error, GetOptions, GetResult, ListResult, MultipartUpload,
ObjectMeta,
+ ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult,
RenameOptions, Result,
+};
+use once_cell::sync::OnceCell;
+
+/// Rebuild function contract: construct a fresh backing `ObjectStore` scoped
for the failing
+/// request and register it in the outer `object_store` registry cache. Called
at most once
+/// per wrapper.
+///
+/// The `Option<&Path>` argument is the location that triggered the 403
(source path for
+/// copy/rename). The rebuild closure passes this into the fresh bridge so the
SPI's
+/// `getPolicyLocationsFor` reflects the actual failing request rather than
the path baked in
+/// at construction time. `None` is reserved for callers that intercept a 403
without a path
+/// (none of the current retry sites) and lets the closure fall back to its
pre-baked path.
+pub type RebuildFn =
+ Arc<dyn Fn(Option<&Path>) -> Result<Arc<dyn ObjectStore>> + Send + Sync +
'static>;
+
+/// Wraps an `Arc<dyn ObjectStore>` so a single 403 rebuilds the store once
and retries.
+///
+/// See the module-level doc-comment for rationale and non-goals.
+pub struct RetryOn403ObjectStore {
+ inner: Arc<dyn ObjectStore>,
+ rebuild: RebuildFn,
+ /// Populated on the first 403 we successfully recover from. Cached so a
subsequent
+ /// request against this same wrapper skips straight to the rebuilt store,
and a 403 there
+ /// is treated as authoritative.
+ rebuilt: OnceCell<Arc<dyn ObjectStore>>,
+}
+
+impl RetryOn403ObjectStore {
+ pub fn new(inner: Arc<dyn ObjectStore>, rebuild: RebuildFn) -> Self {
+ Self {
+ inner,
+ rebuild,
+ rebuilt: OnceCell::new(),
+ }
+ }
+
+ /// Return the store that should service *this* request: the post-rebuild
store when we
+ /// have one, else the original.
+ fn current(&self) -> Arc<dyn ObjectStore> {
+ self.rebuilt
+ .get()
+ .cloned()
+ .unwrap_or_else(|| Arc::clone(&self.inner))
+ }
+
+ /// Whether we have already spent our single rebuild allowance.
+ fn already_rebuilt(&self) -> bool {
+ self.rebuilt.get().is_some()
+ }
+
+ /// Attempt to install a rebuilt store, calling `rebuild` at most once.
Concurrent 403s
+ /// race harmlessly here — `once_cell::sync::OnceCell` guarantees only one
initializer
+ /// runs; the rest observe the same result. The `path` of the request that
first triggered
+ /// the rebuild is threaded into the closure so it can request a scope for
the actual
+ /// failing location.
+ fn rebuild_once(&self, path: Option<&Path>) -> Result<Arc<dyn
ObjectStore>> {
+ self.rebuilt
+ .get_or_try_init(|| (self.rebuild)(path))
+ .map(Arc::clone)
+ }
+}
+
+/// Only `PermissionDenied` (S3 403) triggers the retry. `Unauthenticated`
(401) is treated as
+/// a permanent credential-config error and not retried.
+fn is_forbidden(err: &Error) -> bool {
+ matches!(err, Error::PermissionDenied { .. })
+}
+
+impl fmt::Debug for RetryOn403ObjectStore {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ f.debug_struct("RetryOn403ObjectStore")
+ .field("inner", &self.inner)
+ .field("rebuilt_cached", &self.rebuilt.get().is_some())
+ .finish()
+ }
+}
+
+impl fmt::Display for RetryOn403ObjectStore {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ write!(f, "RetryOn403({})", self.inner)
+ }
+}
+
+#[async_trait]
+impl ObjectStore for RetryOn403ObjectStore {
+ async fn put_opts(
+ &self,
+ location: &Path,
+ payload: PutPayload,
+ opts: PutOptions,
+ ) -> Result<PutResult> {
+ let store = self.current();
+ match store.put_opts(location, payload.clone(), opts.clone()).await {
+ Err(e) if is_forbidden(&e) && !self.already_rebuilt() => {
+ debug!("RetryOn403: 403 on put({location}); rebuilding store");
+ let rebuilt = self.rebuild_once(Some(location))?;
+ rebuilt.put_opts(location, payload, opts).await
+ }
+ other => other,
+ }
+ }
+
+ async fn put_multipart_opts(
+ &self,
+ location: &Path,
+ opts: PutMultipartOptions,
+ ) -> Result<Box<dyn MultipartUpload>> {
+ // Multipart uploads can't be transparently retried mid-flight;
propagate 403 as-is.
+ self.current().put_multipart_opts(location, opts).await
+ }
+
+ async fn get_opts(&self, location: &Path, options: GetOptions) ->
Result<GetResult> {
+ let store = self.current();
+ match store.get_opts(location, options.clone()).await {
+ Err(e) if is_forbidden(&e) && !self.already_rebuilt() => {
+ debug!("RetryOn403: 403 on get({location}); rebuilding store");
+ let rebuilt = self.rebuild_once(Some(location))?;
Review Comment:
### Correctness
[P2] Could this decision distinguish a 403 from the original store from a
403 returned by the rebuilt store? Two requests can capture the original store
before either completes. If the first finishes, rebuilds successfully, and
retries, the second then observes `already_rebuilt() == true` and returns its
original 403 without trying the fresh store. A deterministic interleaving of
these exact methods reproduced one successful request and one failed request,
even though a later call for the failed path succeeds on the rebuilt store.
`OnceCell` prevents duplicate construction but does not make this guard safe.
Please let each operation retry once against the published replacement and
cover the same race in `get_ranges`.
##########
native/core/src/parquet/parquet_support.rs:
##########
@@ -832,18 +939,96 @@ pub(crate) fn prepare_object_store_with_configs(
let (store, path): (Box<dyn ObjectStore>, Path) = if
is_hdfs_scheme {
create_hdfs_object_store(&url)
} else if scheme == "s3" {
- objectstore::s3::create_store(&url, object_store_configs,
Duration::from_secs(300))
+ objectstore::s3::create_store_with_bridge(
+ &url,
+ object_store_configs,
+ bridge_opt.clone(),
+ Duration::from_secs(300),
+ )
} else if is_azure_scheme(scheme) {
objectstore::azure::create_store(&url, object_store_configs)
} else {
parse_url(&url)
}
.map_err(|e| ExecutionError::GeneralError(e.to_string()))?;
- let store: Arc<dyn ObjectStore> = Arc::from(store);
- // Insert into cache
+ let raw_store: Arc<dyn ObjectStore> = Arc::from(store);
+
+ // Wrap SPI-backed stores in the 403-retry safety net. Non-SPI
paths keep their
+ // plain store so we don't add overhead for the base credential
chain (which is
+ // already correct without the wrapper).
+ let store: Arc<dyn ObjectStore> = if bridge_opt.is_some() {
+ let cache_key_for_rebuild = cache_key.clone();
+ let url_for_rebuild = url.clone();
+ let configs_for_rebuild: HashMap<String, String> =
object_store_configs.clone();
+ let rebuild: RebuildFn = Arc::new(move |failing_path:
Option<&Path>| {
+ // Fresh bridge → fresh SPI dispatch → fresh credentials
scoped for the
+ // actual failing request. We rebind the bridge's baked-in
path to the
+ // 403'd location (falling back to the URL's path when the
retry site did
+ // not supply one) so `fetch_policy_locations` returns the
vendor's scope
+ // for *this* request rather than whatever the initial
construction picked.
+ let failing_path_str = failing_path.map(|p|
format!("/{p}"));
+ let rebuilt_bridge =
objectstore::s3::try_construct_bridge_with_path(
+ &url_for_rebuild,
+ &configs_for_rebuild,
+ failing_path_str.as_deref(),
+ )
+ .map_err(|e| object_store::Error::Generic {
+ store: "S3",
+ source: format!("rebuild bridge failed: {e}").into(),
+ })?;
+ // Query the vendor for the fresh session's scope hint.
Errors normalize
+ // to an empty prefix list (catchall for the new entry
only) — matching
+ // the fallback the initial builder uses on the same call.
+ let fresh_prefixes = if let Some(bridge) =
rebuilt_bridge.as_ref() {
+ bridge.fetch_policy_locations().unwrap_or_else(|e| {
+ debug!("fetch_policy_locations on rebuild failed:
{e}");
+ Vec::new()
+ })
+ } else {
+ Vec::new()
+ };
+ let (rebuilt_raw, _path) =
objectstore::s3::create_store_with_bridge(
+ &url_for_rebuild,
+ &configs_for_rebuild,
+ rebuilt_bridge,
+ Duration::from_secs(300),
+ )?;
Review Comment:
### Correctness
[P2] Could recovery reuse the resolved region or perform the rebuild without
synchronously entering Tokio? This closure is invoked inside the async
`get_opts`/`get_ranges` operation. For an ordinary AWS bucket with neither
`fs.s3a.endpoint` nor `fs.s3a.endpoint.region` configured,
`create_store_with_bridge` calls
`get_runtime().block_on(resolve_bucket_region(bucket))`. Tokio rejects that
nested runtime entry with a panic, including when the region cache would return
immediately. I verified the latter with the locked Tokio 1.53.1 and a ready
cached-region future, without making an AWS request. Thus the first recoverable
403 can abort the scan instead of retrying. Please cover recovery without an
explicit endpoint or region.
--
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]
