ursa_engine/
stores.rs

1//! Object stores, registered the first time a query names a URL.
2//!
3//! DataFusion only knows the local filesystem until a store is registered for a
4//! URL's `scheme://authority`. As datafusion-cli does, the engine registers one
5//! lazily: when a query first reads (or `COPY`s to) `s3://bucket/…`, a store for
6//! `s3://bucket` is built from the standard environment and kept for the
7//! session.
8//!
9//! Credentials come from each backend's environment variables and default
10//! credential chain (`object_store`'s `from_env`): `AWS_*` (plus web identity,
11//! container and instance credentials), `GOOGLE_*` (plus application default
12//! credentials), and `AZURE_*`. A public S3 bucket can be read with
13//! `AWS_SKIP_SIGNATURE=true`.
14//!
15//! With the `aws-profile` feature (on by default in the `ursa` binary), S3
16//! also reads the AWS profile files, `~/.aws/config` and `~/.aws/credentials`
17//! (or `AWS_CONFIG_FILE` and `AWS_SHARED_CREDENTIALS_FILE`), for the profile
18//! `AWS_PROFILE` names, or `default`. It comes after the environment's keys
19//! and before web identity, container and instance credentials, the order the
20//! AWS CLI uses. A profile can hold static keys, an SSO login, a role to
21//! assume or a `credential_process`; `aws-config` resolves each as the AWS
22//! SDKs do. See [`aws_profile`].
23
24use std::sync::Arc;
25
26use datafusion::error::{DataFusionError, Result};
27use datafusion::execution::runtime_env::RuntimeEnv;
28#[cfg(not(target_arch = "wasm32"))]
29use object_store::aws::{resolve_bucket_region, AmazonS3Builder};
30#[cfg(not(target_arch = "wasm32"))]
31use object_store::azure::MicrosoftAzureBuilder;
32#[cfg(not(target_arch = "wasm32"))]
33use object_store::gcp::GoogleCloudStorageBuilder;
34use object_store::http::HttpBuilder;
35use object_store::{ClientOptions, ObjectStore};
36use url::{Position, Url};
37
38/// Make sure `runtime` has an object store for `url`, building and registering
39/// one if this is the first time the session sees its `scheme://authority`.
40/// Local paths need none.
41pub async fn ensure_object_store(runtime: &RuntimeEnv, url: &Url) -> Result<()> {
42    if url.scheme() == "file" || runtime.object_store_registry.get_store(url).is_ok() {
43        return Ok(());
44    }
45    let store = build_object_store(url).await?;
46    runtime.register_object_store(&store_base(url)?, store);
47    Ok(())
48}
49
50/// [`ensure_object_store`] for a location written in SQL (a `COPY … TO` target
51/// or an external table's `LOCATION`). A local path needs no store.
52pub async fn ensure_object_store_for(runtime: &RuntimeEnv, location: &str) -> Result<()> {
53    match Url::parse(location) {
54        // A one-letter scheme is a Windows drive letter.
55        Ok(url) if url.scheme().len() > 1 => ensure_object_store(runtime, &url).await,
56        _ => Ok(()),
57    }
58}
59
60/// The `scheme://authority[:port]` a store for `url` is registered under.
61fn store_base(url: &Url) -> Result<Url> {
62    Url::parse(&url[..Position::BeforePath])
63        .map_err(|e| DataFusionError::Plan(format!("'{url}' is not a valid URL: {e}")))
64}
65
66/// A store for `url`'s scheme and bucket, configured from the environment.
67pub async fn build_object_store(url: &Url) -> Result<Arc<dyn ObjectStore>> {
68    let base = &url[..Position::BeforePath];
69    let failed = |e: object_store::Error| {
70        DataFusionError::Execution(format!("cannot connect to {base}: {e}"))
71    };
72    let store: Arc<dyn ObjectStore> = match url.scheme() {
73        #[cfg(not(target_arch = "wasm32"))]
74        "s3" | "s3a" => {
75            let mut builder = AmazonS3Builder::from_env().with_url(url.as_str());
76            // The region a profile names, used when the environment names
77            // none and S3 cannot say where the bucket is.
78            #[allow(unused_mut)]
79            let mut profile_region: Option<String> = None;
80            #[cfg(feature = "aws-profile")]
81            if let Some(profile) = aws_profile::load()
82                .await
83                .map_err(|e| DataFusionError::Execution(format!("cannot connect to {base}: {e}")))?
84            {
85                if let Some(credentials) = profile.credentials {
86                    builder = builder.with_credentials(credentials);
87                }
88                profile_region = profile.region;
89            }
90            if !env_set(&["AWS_REGION", "AWS_DEFAULT_REGION"]) {
91                // Without a region the client assumes us-east-1, and any other
92                // bucket answers with a redirect. Ask S3 where the bucket is,
93                // unless the environment names an endpoint (a local or
94                // S3-compatible store), then fall back to the profile's region.
95                let mut region = None;
96                if !env_set(&["AWS_ENDPOINT", "AWS_ENDPOINT_URL"]) {
97                    if let Some(bucket) = url.host_str() {
98                        region = resolve_bucket_region(bucket, &ClientOptions::new())
99                            .await
100                            .ok();
101                    }
102                }
103                if let Some(region) = region.or(profile_region) {
104                    builder = builder.with_region(region);
105                }
106            }
107            Arc::new(builder.build().map_err(failed)?)
108        }
109        #[cfg(not(target_arch = "wasm32"))]
110        "gs" | "gcs" => Arc::new(
111            GoogleCloudStorageBuilder::from_env()
112                .with_url(url.as_str())
113                .build()
114                .map_err(failed)?,
115        ),
116        #[cfg(not(target_arch = "wasm32"))]
117        "az" | "azure" | "abfs" | "abfss" | "adl" => Arc::new(
118            MicrosoftAzureBuilder::from_env()
119                .with_url(url.as_str())
120                .build()
121                .map_err(failed)?,
122        ),
123        // A plain HTTP(S) URL reads single files with ranged GETs. `http://` is
124        // allowed because the user chose the scheme.
125        "http" | "https" => Arc::new(
126            HttpBuilder::new()
127                .with_url(base)
128                .with_client_options(ClientOptions::new().with_allow_http(true))
129                .build()
130                .map_err(failed)?,
131        ),
132        #[cfg(target_arch = "wasm32")]
133        other @ ("s3" | "s3a" | "gs" | "gcs" | "az" | "azure" | "abfs" | "abfss" | "adl") => {
134            return Err(DataFusionError::Plan(format!(
135                "'{other}://' URLs are not available in the browser build; use an http:// or \
136                 https:// URL"
137            )))
138        }
139        other => {
140            return Err(DataFusionError::Plan(format!(
141                "Ursa cannot read '{other}://' URLs; use a local path or an s3://, gs://, az://, \
142                 http:// or https:// URL"
143            )))
144        }
145    };
146    Ok(store)
147}
148
149/// Whether any of the environment variables `keys` is set and not empty.
150#[cfg(not(target_arch = "wasm32"))]
151fn env_set(keys: &[&str]) -> bool {
152    keys.iter()
153        .any(|k| std::env::var_os(k).is_some_and(|v| !v.is_empty()))
154}
155
156/// S3 credentials and region from the AWS profile files, as the AWS SDKs read
157/// them (`aws-config`).
158///
159/// The profile is `AWS_PROFILE`'s, or `default`. It is read only when the
160/// environment holds no keys (`AWS_ACCESS_KEY_ID`) and does not ask to skip
161/// signing (`AWS_SKIP_SIGNATURE`), so the environment wins, as it does for the
162/// AWS CLI. A profile with no credentials (none of static keys, SSO, a role or
163/// a `credential_process`) leaves the rest of `object_store`'s chain to find
164/// them: web identity, container and instance credentials. A profile that
165/// `AWS_PROFILE` names and that cannot give credentials is an error, rather
166/// than a quiet fall back to other credentials.
167#[cfg(feature = "aws-profile")]
168pub mod aws_profile {
169    use std::sync::Arc;
170    use std::time::{Duration, SystemTime};
171
172    use async_trait::async_trait;
173    use aws_config::meta::region::ProvideRegion;
174    use aws_config::profile::{ProfileFileCredentialsProvider, ProfileFileRegionProvider};
175    use aws_config::provider_config::ProviderConfig;
176    use aws_credential_types::provider::error::CredentialsError;
177    use aws_credential_types::provider::ProvideCredentials;
178    use aws_credential_types::Credentials;
179    use aws_smithy_http_client::tls::{self, rustls_provider::CryptoMode};
180    use object_store::aws::AwsCredential;
181    use object_store::CredentialProvider;
182    use tokio::sync::Mutex;
183
184    /// What a profile gives S3.
185    pub struct Profile {
186        /// Its credentials, or `None` when it has none.
187        pub credentials: Option<Arc<ProfileCredentials>>,
188        pub region: Option<String>,
189    }
190
191    /// The profile's credentials and region, or `None` when the environment
192    /// supplies credentials itself; an error says what is wrong with the
193    /// profile.
194    pub async fn load() -> Result<Option<Profile>, String> {
195        let skip_signature = std::env::var("AWS_SKIP_SIGNATURE")
196            .is_ok_and(|v| matches!(v.to_ascii_lowercase().as_str(), "true" | "1"));
197        if super::env_set(&["AWS_ACCESS_KEY_ID"]) || skip_signature {
198            return Ok(None);
199        }
200        let named = std::env::var("AWS_PROFILE").ok().filter(|p| !p.is_empty());
201        let name = named.clone().unwrap_or_else(|| "default".to_string());
202        // SSO and assume-role reach AWS over HTTPS; ring is the TLS crypto the
203        // rest of the tree already uses.
204        let http = aws_smithy_http_client::Builder::new()
205            .tls_provider(tls::Provider::Rustls(CryptoMode::Ring))
206            .build_https();
207        let config = ProviderConfig::without_region().with_http_client(http);
208        let region = ProfileFileRegionProvider::builder()
209            .profile_name(&name)
210            .configure(&config)
211            .build()
212            .region()
213            .await;
214        let config = config.with_region(region.clone());
215        let provider = ProfileFileCredentialsProvider::builder()
216            .profile_name(&name)
217            .configure(&config)
218            .build();
219        let region = region.map(|r| r.to_string());
220        match provider.provide_credentials().await {
221            Ok(first) => Ok(Some(Profile {
222                credentials: Some(Arc::new(ProfileCredentials {
223                    provider,
224                    current: Mutex::new(first),
225                })),
226                region,
227            })),
228            Err(CredentialsError::CredentialsNotLoaded(_)) if named.is_none() => {
229                Ok(Some(Profile {
230                    credentials: None,
231                    region,
232                }))
233            }
234            Err(e) => Err(format!(
235                "AWS profile '{name}' gives no credentials: {}",
236                error_chain(&e)
237            )),
238        }
239    }
240
241    /// `e` and its sources, joined, as `aws-config` words them.
242    fn error_chain(e: &(dyn std::error::Error + 'static)) -> String {
243        let mut parts = vec![e.to_string()];
244        let mut source = e.source();
245        while let Some(s) = source {
246            let text = s.to_string();
247            if !parts.iter().any(|p| p.contains(&text)) {
248                parts.push(text);
249            }
250            source = s.source();
251        }
252        parts.join(": ")
253    }
254
255    /// A profile's credentials for `object_store`, refreshed shortly before
256    /// they expire (an SSO login or an assumed role's are temporary).
257    #[derive(Debug)]
258    pub struct ProfileCredentials {
259        provider: ProfileFileCredentialsProvider,
260        current: Mutex<Credentials>,
261    }
262
263    /// How long before expiry credentials are refreshed.
264    const REFRESH_BEFORE: Duration = Duration::from_secs(5 * 60);
265
266    #[async_trait]
267    impl CredentialProvider for ProfileCredentials {
268        type Credential = AwsCredential;
269
270        async fn get_credential(&self) -> object_store::Result<Arc<AwsCredential>> {
271            let mut current = self.current.lock().await;
272            let stale = current
273                .expiry()
274                .is_some_and(|at| at <= SystemTime::now() + REFRESH_BEFORE);
275            if stale {
276                *current = self.provider.provide_credentials().await.map_err(|e| {
277                    object_store::Error::Generic {
278                        store: "S3",
279                        source: error_chain(&e).into(),
280                    }
281                })?;
282            }
283            Ok(Arc::new(AwsCredential {
284                key_id: current.access_key_id().to_string(),
285                secret_key: current.secret_access_key().to_string(),
286                token: current.session_token().map(str::to_string),
287            }))
288        }
289    }
290}
ursa_engine/stores.rs

ursa_engine/
stores.rs
