query_runtime/runtime/
object_store.rs

1//! Object-store control-plane handles and shared runtime registry.
2
3use std::{
4    collections,
5    collections::HashMap,
6    fmt,
7    sync::{Arc, Mutex, OnceLock},
8    time,
9    time::{Duration, Instant},
10};
11
12use bytes::Bytes;
13use object_store::aws::{AmazonS3, AmazonS3Builder, AwsCredential, AwsCredentialProvider};
14use query_planning::error::{QueryError, QueryErrorCode, Result};
15use tokio::runtime::{Builder, Runtime as TokioRuntime};
16
17/// Storage profile used by native remote reads and cache maintenance.
18///
19/// The adapter translates its S3 storage settings into this shape before
20/// execution. Secrets are held only as object-store builder inputs and are
21/// redacted from [`Debug`].
22#[derive(Clone, PartialEq, Eq)]
23pub struct SourceStorageProfile {
24    id: String,
25    s3: S3SourceStorageProfile,
26}
27
28impl SourceStorageProfile {
29    /// Builds an S3 storage profile with a stable non-secret identifier.
30    pub fn s3(id: impl Into<String>, s3: S3SourceStorageProfile) -> Self {
31        Self { id: id.into(), s3 }
32    }
33
34    /// Stable profile identifier referenced by resolved source contracts.
35    pub fn id(&self) -> &str {
36        &self.id
37    }
38
39    /// S3-specific source storage settings.
40    pub fn s3_profile(&self) -> &S3SourceStorageProfile {
41        &self.s3
42    }
43
44    /// Builds an S3 object store with the same credential, endpoint, region,
45    /// and addressing policy used by native source execution.
46    ///
47    /// Adapter-side metadata continuations use this to read Parquet footers
48    /// without maintaining a second storage-profile interpretation.
49    pub fn build_s3_object_store(
50        &self,
51        bucket: impl Into<String>,
52        fallback_region: impl Into<String>,
53    ) -> Result<Arc<dyn object_store::ObjectStore>> {
54        let key = S3ObjectStoreKey {
55            bucket: bucket.into(),
56            region: fallback_region.into(),
57            storage_profile_id: Some(self.id.clone()),
58        };
59        let store: Arc<dyn object_store::ObjectStore> = Arc::new(build_s3_store(&key, Some(self))?);
60        Ok(store)
61    }
62
63    /// Whether this profile needs per-profile settings that libcudf/KvikIO
64    /// cannot currently receive through the typed remote Parquet contract.
65    pub fn requires_scoped_remote_direct_profile(&self) -> bool {
66        self.s3.requires_scoped_remote_direct_profile()
67    }
68}
69
70impl fmt::Debug for SourceStorageProfile {
71    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
72        f.debug_struct("SourceStorageProfile").field("id", &self.id).field("s3", &self.s3).finish()
73    }
74}
75
76/// S3-compatible storage settings for source object-store fetches.
77#[derive(Clone, PartialEq, Eq)]
78pub struct S3SourceStorageProfile {
79    /// Region passed to the S3 client. When absent, the parsed source/fallback
80    /// region is used.
81    pub region: Option<String>,
82    /// Optional S3-compatible endpoint such as `http://localhost:9000`.
83    pub endpoint: Option<String>,
84    /// Optional virtual-hosted addressing flag. Set to `Some(false)` for
85    /// path-style S3-compatible stores such as RustFS/MinIO.
86    pub virtual_hosted_style_request: Option<bool>,
87    /// Optional HTTP allowance. If unset, `http://` endpoints imply `true`.
88    pub(crate) allow_http: Option<bool>,
89    /// Credential source used by this profile.
90    pub credential: S3SourceStorageCredential,
91}
92
93impl S3SourceStorageProfile {
94    /// Builds an S3 profile that uses the object-store crate's default
95    /// credential/config chain unless fields are overridden.
96    pub fn default_chain() -> Self {
97        Self {
98            region: None,
99            endpoint: None,
100            virtual_hosted_style_request: None,
101            allow_http: None,
102            credential: S3SourceStorageCredential::DefaultChain,
103        }
104    }
105
106    /// Overrides the S3 region.
107    pub fn with_region(mut self, region: impl Into<String>) -> Self {
108        self.region = Some(region.into());
109        self
110    }
111
112    /// Overrides the S3 endpoint.
113    pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
114        self.endpoint = Some(endpoint.into());
115        self
116    }
117
118    /// Selects path-style or virtual-hosted-style request addressing.
119    pub fn with_path_style_access(mut self, path_style: bool) -> Self {
120        self.virtual_hosted_style_request = Some(!path_style);
121        self
122    }
123
124    /// Uses unsigned S3 requests without consulting ambient AWS credentials.
125    pub fn with_unsigned_requests(mut self) -> Self {
126        self.credential = S3SourceStorageCredential::Unsigned;
127        self
128    }
129
130    /// Uses explicit static credentials for this profile.
131    pub fn with_static_credentials(
132        mut self,
133        access_key_id: impl Into<String>,
134        secret_access_key: impl Into<String>,
135        session_token: Option<String>,
136    ) -> Self {
137        self.credential = S3SourceStorageCredential::Static {
138            access_key_id: access_key_id.into(),
139            secret_access_key: secret_access_key.into(),
140            session_token,
141        };
142        self
143    }
144
145    /// Whether remote-direct cuDF reads need settings scoped to this profile.
146    /// Process-global KvikIO config does not provide them.
147    pub fn requires_scoped_remote_direct_profile(&self) -> bool {
148        self.region.as_ref().is_some_and(|region| !region.trim().is_empty())
149            || !matches!(self.credential, S3SourceStorageCredential::DefaultChain)
150            || self.endpoint.as_ref().is_some_and(|endpoint| !endpoint.trim().is_empty())
151            || self.virtual_hosted_style_request.is_some()
152            || self.allow_http.is_some()
153    }
154}
155
156impl fmt::Debug for S3SourceStorageProfile {
157    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
158        f.debug_struct("S3SourceStorageProfile")
159            .field("region", &self.region)
160            .field("endpoint", &self.endpoint.as_ref().map(|_| "<configured>"))
161            .field("virtual_hosted_style_request", &self.virtual_hosted_style_request)
162            .field("allow_http", &self.allow_http)
163            .field("credential", &self.credential)
164            .finish()
165    }
166}
167
168/// Credential source for an S3 source storage profile.
169#[derive(Clone, PartialEq, Eq)]
170pub enum S3SourceStorageCredential {
171    /// Use the object-store crate's default AWS credential/config chain.
172    DefaultChain,
173    /// Use unsigned S3 requests and do not consult ambient AWS credentials.
174    Unsigned,
175    /// Use explicit static credentials supplied by the adapter.
176    Static {
177        /// Access key id. Redacted from [`Debug`] to keep logs uniformly secret-free.
178        access_key_id: String,
179        /// Secret access key. Never render in diagnostics.
180        secret_access_key: String,
181        /// Optional session token. Never render in diagnostics.
182        session_token: Option<String>,
183    },
184}
185
186impl fmt::Debug for S3SourceStorageCredential {
187    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
188        match self {
189            Self::DefaultChain => f.write_str("DefaultChain"),
190            Self::Unsigned => f.write_str("Unsigned"),
191            Self::Static { session_token, .. } => f
192                .debug_struct("Static")
193                .field("access_key_id", &"<redacted>")
194                .field("secret_access_key", &"<redacted>")
195                .field("session_token", &session_token.as_ref().map(|_| "<redacted>"))
196                .finish(),
197        }
198    }
199}
200
201#[derive(Clone, Debug, PartialEq, Eq, Hash)]
202pub struct S3ObjectStoreKey {
203    #[doc(hidden)]
204    pub bucket: String,
205    #[doc(hidden)]
206    pub region: String,
207    #[doc(hidden)]
208    pub storage_profile_id: Option<String>,
209}
210
211pub struct ObjectStoreFetchResult {
212    #[doc(hidden)]
213    pub bytes: Result<Bytes>,
214}
215
216pub struct ObjectStoreCustomFetchResult<T> {
217    #[doc(hidden)]
218    pub value: Result<T>,
219    #[doc(hidden)]
220    pub elapsed: Duration,
221}
222
223/// Maximum age for reusing a resolved default-chain credential.
224///
225/// The provider owns refresh; this only avoids repeated cross-runtime lookups
226/// and must stay well below the shortest credential lifetime.
227pub(super) const RESOLVED_SOURCE_CREDENTIAL_REUSE_WINDOW: Duration = Duration::from_secs(60);
228
229struct ResolvedSourceCredential {
230    resolved_at: Instant,
231    credential: Arc<AwsCredential>,
232}
233
234// Per-context S3 store and credential cache keyed by source storage profile.
235pub(super) struct SourceObjectStoreRegistry {
236    profiles: HashMap<String, SourceStorageProfile>,
237    stores: Mutex<HashMap<S3ObjectStoreKey, Arc<AmazonS3>>>,
238    resolved_credentials: Mutex<HashMap<S3ObjectStoreKey, ResolvedSourceCredential>>,
239}
240
241impl SourceObjectStoreRegistry {
242    #[cfg(test)]
243    pub(super) fn new() -> Self {
244        Self {
245            profiles: HashMap::new(),
246            stores: Mutex::new(HashMap::new()),
247            resolved_credentials: Mutex::new(HashMap::new()),
248        }
249    }
250
251    pub(super) fn with_profiles(profiles: impl IntoIterator<Item = SourceStorageProfile>) -> Result<Self> {
252        let profiles = profiles.into_iter().collect::<Vec<_>>();
253        Self::validate_profiles(&profiles)?;
254        let by_id = profiles.into_iter().map(|profile| (profile.id().to_string(), profile)).collect();
255        Ok(Self {
256            profiles: by_id,
257            stores: Mutex::new(HashMap::new()),
258            resolved_credentials: Mutex::new(HashMap::new()),
259        })
260    }
261
262    pub(super) fn validate_profiles(profiles: &[SourceStorageProfile]) -> Result<()> {
263        let mut profile_ids = collections::HashSet::with_capacity(profiles.len());
264        for profile in profiles {
265            if profile.id().trim().is_empty() {
266                return Err(QueryError::new(
267                    QueryErrorCode::InvalidConfiguration,
268                    "query_engine.source_object_store.profile",
269                    "source storage profile id must not be empty",
270                )
271                .with_fact("reason", "empty_profile_id"));
272            }
273            if !profile_ids.insert(profile.id()) {
274                return Err(QueryError::new(
275                    QueryErrorCode::InvalidConfiguration,
276                    "query_engine.source_object_store.profile",
277                    "source storage profile ids must be unique",
278                )
279                .with_fact("storage_profile_id", profile.id())
280                .with_fact("reason", "duplicate_profile_id"));
281            }
282        }
283        Ok(())
284    }
285
286    pub(super) fn s3_store(&self, key: S3ObjectStoreKey) -> Result<Arc<dyn object_store::ObjectStore>> {
287        let store: Arc<dyn object_store::ObjectStore> = self.s3_typed_store(key)?;
288        Ok(store)
289    }
290
291    /// Returns the S3 credential provider shared with remote KvikIO reads.
292    pub(super) fn s3_credential_provider(&self, key: S3ObjectStoreKey) -> Result<AwsCredentialProvider> {
293        Ok(Arc::clone(self.s3_typed_store(key)?.credentials()))
294    }
295
296    /// Returns a cached resolved credential still within the reuse window.
297    pub(super) fn recently_resolved_s3_credential(&self, key: &S3ObjectStoreKey) -> Option<Arc<AwsCredential>> {
298        let resolved_credentials = self.resolved_credentials.lock().ok()?;
299        let resolved = resolved_credentials.get(key)?;
300        if resolved.resolved_at.elapsed() < RESOLVED_SOURCE_CREDENTIAL_REUSE_WINDOW {
301            Some(Arc::clone(&resolved.credential))
302        } else {
303            None
304        }
305    }
306
307    /// Records a freshly resolved credential for reuse.
308    pub(super) fn record_resolved_s3_credential(&self, key: S3ObjectStoreKey, credential: Arc<AwsCredential>) {
309        if let Ok(mut resolved_credentials) = self.resolved_credentials.lock() {
310            resolved_credentials
311                .insert(key, ResolvedSourceCredential { resolved_at: time::Instant::now(), credential });
312        }
313    }
314
315    fn s3_typed_store(&self, key: S3ObjectStoreKey) -> Result<Arc<AmazonS3>> {
316        let mut stores = self.stores.lock().map_err(|_error| {
317            QueryError::new(
318                QueryErrorCode::Internal,
319                "query_engine.source_object_store.registry",
320                "source object-store registry mutex was poisoned",
321            )
322        })?;
323        if let Some(store) = stores.get(&key) {
324            return Ok(Arc::clone(store));
325        }
326        let profile = key
327            .storage_profile_id
328            .as_ref()
329            .map(|id| {
330                self.profiles.get(id).ok_or_else(|| {
331                    QueryError::new(
332                        QueryErrorCode::InvalidPlan,
333                        "query_engine.source_object_store.profile",
334                        "source storage profile id was not registered",
335                    )
336                    .with_fact("storage_profile_id", id.as_str())
337                })
338            })
339            .transpose()?;
340        let store = Arc::new(build_s3_store(&key, profile)?);
341        stores.insert(key, Arc::clone(&store));
342        Ok(store)
343    }
344
345    pub(super) fn profile(&self, id: &str) -> Result<&SourceStorageProfile> {
346        self.profiles.get(id).ok_or_else(|| {
347            QueryError::new(
348                QueryErrorCode::InvalidPlan,
349                "query_engine.source_object_store.profile",
350                "source storage profile id was not registered",
351            )
352            .with_fact("storage_profile_id", id)
353        })
354    }
355}
356
357fn build_s3_store(key: &S3ObjectStoreKey, profile: Option<&SourceStorageProfile>) -> Result<AmazonS3> {
358    let s3_profile = profile.map(SourceStorageProfile::s3_profile);
359    let region = s3_profile
360        .and_then(|profile| profile.region.as_deref())
361        .filter(|region| !region.is_empty())
362        .unwrap_or(key.region.as_str());
363    let mut builder = match s3_profile.map(|profile| &profile.credential) {
364        Some(S3SourceStorageCredential::Static { access_key_id, secret_access_key, session_token }) => {
365            let mut builder =
366                AmazonS3Builder::new().with_access_key_id(access_key_id).with_secret_access_key(secret_access_key);
367            if let Some(token) = session_token.as_ref().filter(|token| !token.is_empty()) {
368                builder = builder.with_token(token);
369            }
370            builder
371        },
372        Some(S3SourceStorageCredential::Unsigned) => AmazonS3Builder::new().with_skip_signature(true),
373        _ => AmazonS3Builder::from_env(),
374    }
375    .with_bucket_name(key.bucket.as_str())
376    .with_region(region);
377    if let Some(profile) = s3_profile {
378        if let Some(endpoint) = profile.endpoint.as_ref().filter(|endpoint| !endpoint.trim().is_empty()) {
379            if endpoint_authority_contains_userinfo(endpoint) {
380                return Err(
381                    s3_store_build_error(key, region, s3_profile).with_fact("reason", "endpoint_userinfo_unsupported")
382                );
383            }
384            builder = builder.with_endpoint(endpoint);
385            if profile.allow_http.unwrap_or_else(|| endpoint.trim_start().starts_with("http://")) {
386                builder = builder.with_allow_http(true);
387            }
388        } else if let Some(allow_http) = profile.allow_http {
389            builder = builder.with_allow_http(allow_http);
390        }
391        if let Some(virtual_hosted_style_request) = profile.virtual_hosted_style_request {
392            builder = builder.with_virtual_hosted_style_request(virtual_hosted_style_request);
393        }
394    }
395    builder.build().map_err(|_error| s3_store_build_error(key, region, s3_profile))
396}
397
398fn s3_store_build_error(
399    key: &S3ObjectStoreKey,
400    region: &str,
401    s3_profile: Option<&S3SourceStorageProfile>,
402) -> QueryError {
403    let mut query_error = QueryError::new(
404        QueryErrorCode::CudfDependencyFailure,
405        "query_engine.source_object_store.s3",
406        "failed to build S3 object-store client for native control-plane work",
407    )
408    .with_fact("bucket", key.bucket.as_str())
409    .with_fact("region", region)
410    .with_fact("endpoint_configured", s3_profile.and_then(|profile| profile.endpoint.as_ref()).is_some());
411    if let Some(profile_id) = key.storage_profile_id.as_deref() {
412        query_error = query_error.with_fact("storage_profile_id", profile_id);
413    }
414    query_error
415}
416
417fn endpoint_authority_contains_userinfo(endpoint: &str) -> bool {
418    let trimmed = endpoint.trim();
419    let after_scheme = trimmed.split_once("://").map_or(trimmed, |(_, after_scheme)| after_scheme);
420    let authority = after_scheme.split(['/', '?', '#']).next().unwrap_or(after_scheme);
421    authority.contains('@')
422}
423
424impl fmt::Debug for SourceObjectStoreRegistry {
425    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
426        let cached_store_count = self.stores.lock().map_or(0, |stores| stores.len());
427        f.debug_struct("SourceObjectStoreRegistry")
428            .field("profile_count", &self.profiles.len())
429            .field("cached_store_count", &cached_store_count)
430            .finish()
431    }
432}
433
434pub(super) fn default_object_store_control_runtime() -> Result<Arc<TokioRuntime>> {
435    static RUNTIME: OnceLock<Arc<TokioRuntime>> = OnceLock::new();
436    static INIT_LOCK: Mutex<()> = Mutex::new(());
437    if let Some(runtime) = RUNTIME.get() {
438        return Ok(Arc::clone(runtime));
439    }
440    let _guard = INIT_LOCK.lock().map_err(|_| {
441        QueryError::new(
442            QueryErrorCode::CudfDependencyFailure,
443            "query_engine.runtime.object_store_control_runtime",
444            "object-store control runtime initialization lock was poisoned",
445        )
446    })?;
447    if let Some(runtime) = RUNTIME.get() {
448        return Ok(Arc::clone(runtime));
449    }
450    let runtime =
451        Arc::new(Builder::new_multi_thread().thread_name("algeon-object-store-control").enable_all().build().map_err(
452            |error| {
453                QueryError::new(
454                    QueryErrorCode::CudfDependencyFailure,
455                    "query_engine.runtime.object_store_control_runtime",
456                    "failed to initialize object-store control runtime",
457                )
458                .with_source(error)
459            },
460        )?);
461    assert!(
462        RUNTIME.set(Arc::clone(&runtime)).is_ok(),
463        "object-store control runtime initialized while holding init lock"
464    );
465    Ok(runtime)
466}
467
468#[cfg(test)]
469mod tests {
470    use super::*;
471
472    #[test]
473    fn source_storage_profile_debug_redacts_endpoint_and_credentials() {
474        let profile = SourceStorageProfile::s3(
475            "rest-s3",
476            S3SourceStorageProfile::default_chain()
477                .with_region("us-east-1")
478                .with_endpoint("http://user:password@localhost:9000")
479                .with_path_style_access(true)
480                .with_static_credentials(
481                    "access-key-value",
482                    "secret-key-value",
483                    Some("session-token-value".to_string()),
484                ),
485        );
486
487        let rendered = format!("{profile:?}");
488
489        assert!(rendered.contains("rest-s3"));
490        assert!(rendered.contains("<configured>"));
491        assert!(rendered.contains("<redacted>"));
492        assert!(!rendered.contains("user:password"));
493        assert!(!rendered.contains("access-key-value"));
494        assert!(!rendered.contains("secret-key-value"));
495        assert!(!rendered.contains("session-token-value"));
496    }
497
498    #[test]
499    fn source_object_store_registry_rejects_duplicate_profile_ids() {
500        let profile = SourceStorageProfile::s3("same", S3SourceStorageProfile::default_chain());
501        let error = SourceObjectStoreRegistry::with_profiles([profile.clone(), profile]).unwrap_err();
502
503        assert_eq!(error.code(), QueryErrorCode::InvalidConfiguration);
504        assert_eq!(error.fact_value("reason"), Some("duplicate_profile_id"));
505        assert_eq!(error.fact_value("storage_profile_id"), Some("same"));
506    }
507
508    #[test]
509    fn source_object_store_registry_rejects_empty_profile_ids() {
510        let profile = SourceStorageProfile::s3("  ", S3SourceStorageProfile::default_chain());
511        let error = SourceObjectStoreRegistry::with_profiles([profile]).unwrap_err();
512
513        assert_eq!(error.code(), QueryErrorCode::InvalidConfiguration);
514        assert_eq!(error.fact_value("reason"), Some("empty_profile_id"));
515    }
516
517    #[test]
518    fn source_object_store_registry_requires_registered_profile_id() {
519        let error = SourceObjectStoreRegistry::new()
520            .s3_store(S3ObjectStoreKey {
521                bucket: "bucket".to_string(),
522                region: "us-east-1".to_string(),
523                storage_profile_id: Some("missing".to_string()),
524            })
525            .unwrap_err();
526
527        assert_eq!(error.code(), QueryErrorCode::InvalidPlan);
528        assert_eq!(error.fact_value("storage_profile_id"), Some("missing"));
529    }
530
531    #[test]
532    fn source_storage_profile_reports_scoped_remote_direct_requirement_without_secrets() {
533        let default_chain = SourceStorageProfile::s3("aws-default", S3SourceStorageProfile::default_chain());
534        assert!(!default_chain.requires_scoped_remote_direct_profile());
535
536        let region_override =
537            SourceStorageProfile::s3("region", S3SourceStorageProfile::default_chain().with_region("us-west-2"));
538        assert!(region_override.requires_scoped_remote_direct_profile());
539
540        let static_credentials = SourceStorageProfile::s3(
541            "static",
542            S3SourceStorageProfile::default_chain().with_static_credentials("access", "secret", None),
543        );
544        assert!(static_credentials.requires_scoped_remote_direct_profile());
545
546        let custom_endpoint = SourceStorageProfile::s3(
547            "rustfs",
548            S3SourceStorageProfile::default_chain().with_endpoint("http://localhost:9000").with_path_style_access(true),
549        );
550        assert!(custom_endpoint.requires_scoped_remote_direct_profile());
551    }
552
553    #[test]
554    fn source_storage_profile_builds_object_store_for_adapter_metadata_reads() {
555        let profile =
556            SourceStorageProfile::s3("s3-profile", S3SourceStorageProfile::default_chain().with_unsigned_requests());
557
558        let store = profile.build_s3_object_store("lakehouse", "us-east-1");
559
560        if let Err(error) = store {
561            panic!("native storage profile should build an S3 object store: {error}");
562        }
563    }
564
565    #[test]
566    fn s3_store_build_errors_do_not_expose_endpoint_userinfo() {
567        let profile = SourceStorageProfile::s3(
568            "bad-endpoint",
569            S3SourceStorageProfile::default_chain()
570                .with_endpoint("http://user:password@[::1")
571                .with_path_style_access(true),
572        );
573        let key = S3ObjectStoreKey {
574            bucket: "bucket".to_string(),
575            region: "us-east-1".to_string(),
576            storage_profile_id: Some("bad-endpoint".to_string()),
577        };
578
579        let error = build_s3_store(&key, Some(&profile)).unwrap_err();
580        let rendered = format!("{error:?}\n{error}");
581
582        assert_eq!(error.code(), QueryErrorCode::CudfDependencyFailure);
583        assert_eq!(error.fact_value("reason"), Some("endpoint_userinfo_unsupported"));
584        assert_eq!(error.fact_value("endpoint_configured"), Some("true"));
585        assert!(error.source_error().is_none());
586        assert!(!rendered.contains("user:password"));
587        assert!(!rendered.contains("password"));
588    }
589
590    #[test]
591    fn s3_credential_provider_is_shared_with_control_plane_store_and_resolves_static_credentials() {
592        let registry = SourceObjectStoreRegistry::with_profiles([SourceStorageProfile::s3(
593            "static",
594            S3SourceStorageProfile::default_chain().with_region("us-east-1").with_static_credentials(
595                "access-key-value",
596                "secret-key-value",
597                Some("session-token-value".to_string()),
598            ),
599        )])
600        .expect("registry builds");
601        let key = S3ObjectStoreKey {
602            bucket: "bucket".to_string(),
603            region: "us-east-1".to_string(),
604            storage_profile_id: Some("static".to_string()),
605        };
606
607        let _control_plane_store = registry.s3_store(key.clone()).expect("control-plane store builds");
608        let provider = registry.s3_credential_provider(key).expect("credential provider resolves");
609        let credential = default_object_store_control_runtime()
610            .unwrap()
611            .block_on(provider.get_credential())
612            .expect("credential resolves");
613
614        assert_eq!(credential.key_id, "access-key-value");
615        assert_eq!(credential.secret_key, "secret-key-value");
616        assert_eq!(credential.token.as_deref(), Some("session-token-value"));
617    }
618
619    #[test]
620    fn resolved_s3_credentials_are_reused_per_key_within_the_reuse_window() {
621        let registry = SourceObjectStoreRegistry::new();
622        let key = S3ObjectStoreKey {
623            bucket: "bucket".to_string(),
624            region: "us-east-1".to_string(),
625            storage_profile_id: Some("glue".to_string()),
626        };
627        let other_key = S3ObjectStoreKey { bucket: "other-bucket".to_string(), ..key.clone() };
628        assert!(registry.recently_resolved_s3_credential(&key).is_none());
629
630        let credential = Arc::new(AwsCredential {
631            key_id: "resolved-access-key".to_string(),
632            secret_key: "resolved-secret-key".to_string(),
633            token: None,
634        });
635        registry.record_resolved_s3_credential(key.clone(), Arc::clone(&credential));
636
637        let reused = registry.recently_resolved_s3_credential(&key).expect("fresh credential is reused");
638        assert!(Arc::ptr_eq(&reused, &credential));
639        assert!(registry.recently_resolved_s3_credential(&other_key).is_none());
640    }
641
642    #[test]
643    fn s3_store_rejects_endpoint_userinfo_without_scheme() {
644        let profile = SourceStorageProfile::s3(
645            "bad-endpoint",
646            S3SourceStorageProfile::default_chain()
647                .with_endpoint("user:password@localhost:9000")
648                .with_path_style_access(true),
649        );
650        let key = S3ObjectStoreKey {
651            bucket: "bucket".to_string(),
652            region: "us-east-1".to_string(),
653            storage_profile_id: Some("bad-endpoint".to_string()),
654        };
655
656        let error = build_s3_store(&key, Some(&profile)).unwrap_err();
657
658        assert_eq!(error.code(), QueryErrorCode::CudfDependencyFailure);
659        assert_eq!(error.fact_value("reason"), Some("endpoint_userinfo_unsupported"));
660        assert!(error.source_error().is_none());
661    }
662}
query_runtime/runtime/object_store.rs

query_runtime/runtime/
object_store.rs
