Skip to main content

unitycatalog_object_store/
lib.rs

1//! `object_store` integration for Unity Catalog.
2//!
3//! This crate adapts Unity Catalog's credential-vending APIs to the
4//! [`object_store`](https://docs.rs/object_store) trait, so any framework
5//! that accepts an `Arc<dyn ObjectStore>` (DataFusion, `delta_kernel`,
6//! `parquet`, …) can read and write data governed by Unity Catalog
7//! volumes, tables, or external locations with no extra glue.
8//!
9//! # Quickstart
10//!
11//! ```no_run
12//! use unitycatalog_object_store::{Operation, UnityObjectStoreFactory};
13//!
14//! # async fn run() -> Result<(), Box<dyn std::error::Error>> {
15//! let factory = UnityObjectStoreFactory::builder()
16//!     .with_uri("https://my-workspace.cloud.databricks.com/api/2.1/unity-catalog/")
17//!     .with_token(std::env::var("DATABRICKS_TOKEN").unwrap())
18//!     .build()
19//!     .await?;
20//!
21//! // Address a UC securable directly with a `uc://` URL …
22//! let store = factory
23//!     .for_url("uc:///Volumes/main/default/landing/raw/", Operation::Read)
24//!     .await?;
25//!
26//! // … and use it like any other `object_store`.
27//! let listing = futures::StreamExt::collect::<Vec<_>>(store.as_dyn().list(None)).await;
28//! # Ok(()) }
29//! ```
30//!
31//! # URL scheme
32//!
33//! See [`UCReference`] for the full grammar. In short:
34//!
35//! - `uc:///Volumes/<catalog>/<schema>/<volume>[/<path>]`
36//! - `uc:///Tables/<catalog>/<schema>/<table>`
37//! - `s3://`, `gs://`, `abfss://`, `r2://`, … — raw cloud URLs, vended via
38//!   `temporary-path-credentials`.
39//!
40//! The kind segment is **case-insensitive** (so `/Volumes/`, `/volumes/`,
41//! and `/VOLUMES/` all work); the capitalised form is canonical because it
42//! mirrors the Databricks workspace POSIX path convention. The Databricks
43//! `vol+dbfs:/Volumes/...` alias is also accepted.
44
45use std::sync::Arc;
46
47use object_store::azure::MicrosoftAzureBuilder;
48use object_store::path::Path;
49use object_store::prefix::PrefixStore;
50use object_store::{ObjectStore, Result};
51use unitycatalog_client::{TemporaryCredentialClient, UnityCatalogClient};
52
53// Native-only object_store backends and I/O handle. On `wasm32` the AWS/GCP
54// stores are unsupported (Azure-first), `LocalFileSystem` is absent from
55// object_store, and there is no tokio runtime.
56#[cfg(not(target_arch = "wasm32"))]
57use object_store::aws::AmazonS3Builder;
58#[cfg(not(target_arch = "wasm32"))]
59use object_store::client::SpawnedReqwestConnector;
60#[cfg(not(target_arch = "wasm32"))]
61use object_store::gcp::GoogleCloudStorageBuilder;
62#[cfg(not(target_arch = "wasm32"))]
63use object_store::local::LocalFileSystem;
64#[cfg(not(target_arch = "wasm32"))]
65use tokio::runtime::Handle;
66
67/// HTTP transport for credential vending, selected per target — mirrors the
68/// alias in `unitycatalog-client` (`crates/client/src/lib.rs`). Native builds
69/// use `olai_http::CloudClient`; `wasm32` builds use the browser Fetch
70/// `olai_http_wasm::WasmClient`. The per-service clients
71/// (`TemporaryCredentialClient`, `UnityCatalogClient`) hold this same type.
72#[cfg(not(target_arch = "wasm32"))]
73use olai_http::CloudClient as Transport;
74#[cfg(not(target_arch = "wasm32"))]
75use olai_http::service::{HttpService, ReqwestService};
76#[cfg(target_arch = "wasm32")]
77use olai_http_wasm::WasmClient as Transport;
78use unitycatalog_common::tables::v1::GetTableRequest;
79use unitycatalog_common::temporary_credentials::v1::TemporaryCredential;
80use unitycatalog_common::volumes::v1::GetVolumeRequest;
81use url::Url;
82
83use crate::credential::{SecurableRef, as_aws, as_azure, as_gcp, new_azure};
84// AWS/GCP provider constructors + the access-point helper are native-only: on
85// wasm the store is Azure-first and these are never constructed.
86#[cfg(not(target_arch = "wasm32"))]
87use crate::credential::{aws_access_point, new_aws, new_gcp};
88pub use crate::error::Error;
89pub use unitycatalog_common::UCReference;
90// Re-export the reference / operation enums so consumers do not need a direct
91// dependency on `unitycatalog-client` for the common case.
92pub use unitycatalog_client::{
93    PathOperation, TableOperation, TableReference, VolumeOperation, VolumeReference,
94};
95
96mod credential;
97mod error;
98mod proxy;
99
100/// An [`HttpService`] decorator that injects a fixed request header on every
101/// outbound request before delegating to an inner service.
102///
103/// Used by [`UnityObjectStoreFactory::with_forwarded_user`] to forward a trusted
104/// reverse-proxy identity header (e.g. `x-forwarded-user`) onto the upstream
105/// credential-vending + metadata calls, so the upstream Unity Catalog attributes
106/// the vend to the end user rather than the proxy's own service principal. The
107/// header is `insert`ed (replacing any existing value of the same name), not
108/// appended.
109///
110/// Native-only: the browser Fetch transport (`wasm32`) forwards identity through
111/// its own same-origin session, so this decorator is unused there.
112#[cfg(not(target_arch = "wasm32"))]
113#[derive(Debug)]
114struct ForwardedHeaderService {
115    inner: Arc<dyn HttpService>,
116    name: reqwest::header::HeaderName,
117    value: reqwest::header::HeaderValue,
118}
119
120#[cfg(not(target_arch = "wasm32"))]
121impl HttpService for ForwardedHeaderService {
122    fn call(
123        &self,
124        mut request: reqwest::Request,
125    ) -> std::pin::Pin<
126        Box<dyn std::future::Future<Output = olai_http::Result<reqwest::Response>> + Send + '_>,
127    > {
128        request
129            .headers_mut()
130            .insert(self.name.clone(), self.value.clone());
131        self.inner.call(request)
132    }
133}
134
135/// Builder for [`UnityObjectStoreFactory`].
136#[derive(Debug, Clone, Default)]
137pub struct UnityObjectStoreFactoryBuilder {
138    /// Base URL of the Unity Catalog REST API
139    /// (e.g. `https://<workspace>.cloud.databricks.com/api/2.1/unity-catalog/`).
140    uri: Option<String>,
141    /// Bearer token used for authentication.
142    token: Option<String>,
143    /// Permit construction without a token. Useful for local development
144    /// against an unauthenticated OSS server; do not use in production.
145    allow_unauthenticated: bool,
146    /// Optional AWS region hint. Required when the data lives in a region
147    /// other than `us-east-1` and the server does not return region info
148    /// alongside the vended credential.
149    aws_region: Option<String>,
150    /// Optional dedicated tokio runtime for HTTP I/O. When set, all
151    /// object-store and credential-vending requests are spawned on this
152    /// runtime instead of the ambient one. See [`with_io_runtime`].
153    ///
154    /// [`with_io_runtime`]: UnityObjectStoreFactoryBuilder::with_io_runtime
155    #[cfg(not(target_arch = "wasm32"))]
156    io_handle: Option<Handle>,
157}
158
159impl UnityObjectStoreFactoryBuilder {
160    pub fn new() -> Self {
161        Self::default()
162    }
163
164    /// Set the URI of the Unity Catalog API
165    /// (e.g. `https://<workspace>/api/2.1/unity-catalog/`).
166    pub fn with_uri(mut self, uri: impl Into<String>) -> Self {
167        self.uri = Some(uri.into());
168        self
169    }
170
171    /// Set the [access token] used for bearer authentication.
172    ///
173    /// Accepts both `String` and `Option<String>` — pass `None` to clear
174    /// a previously-set token (e.g. when reusing the builder).
175    ///
176    /// [access token]: https://docs.databricks.com/aws/en/dev-tools/auth/pat
177    pub fn with_token(mut self, token: impl Into<Option<String>>) -> Self {
178        self.token = token.into();
179        self
180    }
181
182    /// Allow construction without any authentication credentials.
183    ///
184    /// Only intended for local development against an unauthenticated OSS
185    /// Unity Catalog server — there should not be any unauthenticated UC
186    /// servers in production deployments.
187    pub fn with_allow_unauthenticated(mut self, allow_unauthenticated: bool) -> Self {
188        self.allow_unauthenticated = allow_unauthenticated;
189        self
190    }
191
192    /// Override the AWS region used for vended AWS credentials.
193    ///
194    /// When unset the factory falls back to (in order):
195    /// 1. The `AWS_REGION` environment variable.
196    /// 2. The `object_store` default region (`us-east-1`).
197    ///
198    /// This is a stop-gap until the server reliably returns region info
199    /// alongside the credential.
200    pub fn with_aws_region(mut self, aws_region: impl Into<Option<String>>) -> Self {
201        self.aws_region = aws_region.into();
202        self
203    }
204
205    /// Route all HTTP I/O onto a dedicated tokio runtime.
206    ///
207    /// In production DataFusion deployments it is common to segregate network
208    /// I/O onto a separate runtime so that CPU-bound query work on the main
209    /// runtime cannot starve object-store requests (and vice versa). When a
210    /// handle is supplied here:
211    ///
212    /// - every cloud object store ([`AmazonS3Builder`], [`MicrosoftAzureBuilder`],
213    ///   [`GoogleCloudStorageBuilder`]) is built with a
214    ///   [`SpawnedReqwestConnector`] that spawns its requests on this runtime; and
215    /// - the credential-vending [`CloudClient`] is configured with the same
216    ///   runtime via [`CloudClient::with_runtime`].
217    ///
218    /// When unset (the default), I/O runs on the ambient runtime — current
219    /// behaviour, fully backwards compatible.
220    ///
221    /// Pass `None` to clear a previously set handle (e.g. when reusing the
222    /// builder).
223    ///
224    /// Not available on `wasm32`: the browser has a single-threaded executor
225    /// and no tokio runtime to route I/O onto.
226    #[cfg(not(target_arch = "wasm32"))]
227    pub fn with_io_runtime(mut self, handle: impl Into<Option<Handle>>) -> Self {
228        self.io_handle = handle.into();
229        self
230    }
231
232    pub async fn build(self) -> Result<UnityObjectStoreFactory> {
233        let url = if let Some(uri) = self.uri.as_ref() {
234            url::Url::parse(uri).map_err(Error::from)?
235        } else {
236            return Err(Error::invalid_config("missing `uri` for Unity Catalog endpoint").into());
237        };
238
239        let transport = self.build_transport()?;
240
241        // On wasm, discover the server's storage-access posture once, up front,
242        // so every `for_*` call routes consistently. A discovery failure (older
243        // server, no proxy) resolves to `Direct` — the historical behavior.
244        #[cfg(target_arch = "wasm32")]
245        let storage_access = proxy::discover_storage_access(&url, self.token.as_deref()).await;
246
247        let creds = TemporaryCredentialClient::new_with_url(transport.clone(), url.clone());
248        let uc = UnityCatalogClient::new(transport, url.clone());
249        Ok(UnityObjectStoreFactory {
250            creds,
251            uc,
252            aws_region: self.aws_region,
253            #[cfg(not(target_arch = "wasm32"))]
254            io_handle: self.io_handle,
255            #[cfg(not(target_arch = "wasm32"))]
256            base_url: url,
257            #[cfg(not(target_arch = "wasm32"))]
258            allow_unauthenticated: self.allow_unauthenticated,
259            #[cfg(target_arch = "wasm32")]
260            storage_access,
261            #[cfg(target_arch = "wasm32")]
262            token: self.token,
263        })
264    }
265
266    /// Build the credential-vending transport for the current target.
267    ///
268    /// Native: the `olai-http` cloud client, optionally routed onto the
269    /// dedicated I/O runtime. `wasm32`: the browser Fetch client, carrying a
270    /// bearer token via `with_auth` when one was supplied (otherwise it relies
271    /// on the ambient browser session — cookies / forwarded auth).
272    #[cfg(not(target_arch = "wasm32"))]
273    fn build_transport(&self) -> Result<Transport> {
274        let cloud_client = if let Some(token) = &self.token {
275            Transport::new_with_token(token)
276        } else if self.allow_unauthenticated {
277            Transport::new_unauthenticated()
278        } else {
279            return Err(Error::invalid_config(
280                "no token and `allow_unauthenticated` not set: cannot build credential client",
281            )
282            .into());
283        };
284
285        // Route credential-vending HTTP onto the dedicated I/O runtime when one
286        // was supplied; otherwise leave it on the ambient runtime.
287        Ok(match &self.io_handle {
288            Some(handle) => cloud_client.with_runtime(handle.clone()),
289            None => cloud_client,
290        })
291    }
292
293    #[cfg(target_arch = "wasm32")]
294    fn build_transport(&self) -> Result<Transport> {
295        use olai_http_wasm::CredentialsMode;
296        use reqwest::header::{AUTHORIZATION, HeaderMap, HeaderValue};
297
298        // No `allow_unauthenticated` gate on wasm: a bare browser session is a
299        // valid unauthenticated mode (cookies / same-origin forwarded auth).
300        let client = Transport::new();
301        Ok(match &self.token {
302            Some(token) => {
303                let token = token.clone();
304                // Bearer auth: attach the header per request and stop the
305                // browser from shadowing it with a stale cookie.
306                client
307                    .with_credentials(CredentialsMode::Omit)
308                    .with_auth(move || {
309                        let mut headers = HeaderMap::new();
310                        if let Ok(value) = HeaderValue::from_str(&format!("Bearer {token}")) {
311                            headers.insert(AUTHORIZATION, value);
312                        }
313                        headers
314                    })
315            }
316            None => client,
317        })
318    }
319}
320
321/// A configured Unity Catalog `ObjectStore` ready for use.
322///
323/// The default [`Self::as_dyn`] returns a store that is automatically
324/// prefixed to the credential-scoped sub-path (e.g. just the volume's
325/// storage root); paths passed to `list`/`get`/`put` are interpreted
326/// relative to that prefix. The unprefixed [`Self::root`] is an escape
327/// hatch for callers that need to work at the bucket level.
328#[derive(Clone)]
329pub struct UCStore {
330    /// Bucket-rooted store (credentials may be scoped to a sub-path).
331    root: Arc<dyn ObjectStore>,
332    /// The full cloud URL of the credential-scoped root.
333    url: Url,
334    /// Path within `root` the credential is scoped to.
335    path: Path,
336}
337
338impl UCStore {
339    /// Returns the credential-scoped store (prefixed at [`Self::prefix`]).
340    ///
341    /// This is the common case: callers list / read / write paths inside
342    /// the volume or table the credential was vended for.
343    pub fn as_dyn(&self) -> Arc<dyn ObjectStore> {
344        if self.path.as_ref().is_empty() {
345            self.root.clone()
346        } else {
347            Arc::new(PrefixStore::new(self.root.clone(), self.path.clone()))
348        }
349    }
350
351    /// Returns the bucket-rooted store.
352    ///
353    /// The vended credential may not authorise access to siblings of
354    /// [`Self::prefix`]; callers using `root()` are responsible for not
355    /// accessing paths outside the scoped region.
356    pub fn root(&self) -> Arc<dyn ObjectStore> {
357        self.root.clone()
358    }
359
360    /// The full cloud URL of the credential-scoped root
361    /// (e.g. `s3://bucket/prefix/inside/volume/`).
362    pub fn url(&self) -> &Url {
363        &self.url
364    }
365
366    /// The prefix inside [`Self::root`] the credential is scoped to.
367    pub fn prefix(&self) -> &Path {
368        &self.path
369    }
370}
371
372/// Factory that mints `object_store` instances backed by Unity Catalog
373/// credential vending.
374#[derive(Clone)]
375pub struct UnityObjectStoreFactory {
376    creds: TemporaryCredentialClient,
377    uc: UnityCatalogClient,
378    // Read only by the native AWS store branch; the wasm build is Azure-first and
379    // never consults it, but keeping the field (and `with_aws_region`) uniform
380    // across targets avoids splitting the builder API.
381    #[cfg_attr(target_arch = "wasm32", allow(dead_code))]
382    aws_region: Option<String>,
383    /// Dedicated runtime for object-store HTTP I/O, if configured via
384    /// [`UnityObjectStoreFactoryBuilder::with_io_runtime`]. Absent on `wasm32`.
385    #[cfg(not(target_arch = "wasm32"))]
386    io_handle: Option<Handle>,
387    /// The UC REST base URL, retained so [`with_forwarded_user`] can rebuild the
388    /// credential + metadata clients over a derived transport.
389    ///
390    /// [`with_forwarded_user`]: UnityObjectStoreFactory::with_forwarded_user
391    #[cfg(not(target_arch = "wasm32"))]
392    base_url: Url,
393    /// Whether the factory was built without a token, retained for parity when
394    /// [`with_forwarded_user`] derives a transport (the forwarded-identity path
395    /// is unauthenticated by design; this field documents that the base factory
396    /// also tolerated no token).
397    ///
398    /// [`with_forwarded_user`]: UnityObjectStoreFactory::with_forwarded_user
399    #[cfg(not(target_arch = "wasm32"))]
400    #[allow(dead_code)]
401    allow_unauthenticated: bool,
402    /// Storage-access posture discovered from the server's `/capabilities`
403    /// endpoint at build time. `Proxy` routes every byte through the server's
404    /// same-origin storage byte-proxy instead of vending + direct cloud IO.
405    /// Only meaningful on `wasm32` (native always reads storage directly).
406    #[cfg(target_arch = "wasm32")]
407    storage_access: proxy::StorageAccess,
408    /// Bearer token, retained on `wasm32` so the proxy store can relay it to the
409    /// same-origin `/storage-proxy` endpoint (the endpoint may require auth).
410    #[cfg(target_arch = "wasm32")]
411    token: Option<String>,
412}
413
414impl UnityObjectStoreFactory {
415    pub fn builder() -> UnityObjectStoreFactoryBuilder {
416        UnityObjectStoreFactoryBuilder::default()
417    }
418
419    /// Derive a factory that forwards a trusted reverse-proxy identity header on
420    /// every upstream Unity Catalog request (both the metadata lookups and the
421    /// credential vend).
422    ///
423    /// Intended for a service — e.g. the standalone storage byte-proxy — that
424    /// sits behind the same reverse proxy as the upstream UC and has already
425    /// validated the caller's identity. Forwarding the header verbatim lets UC's
426    /// own reverse-proxy authenticator attribute the vend to the real end user.
427    ///
428    /// Semantics:
429    /// - `user == None` (anonymous request) → returns `self` cloned unchanged, so
430    ///   the upstream calls use whatever auth the base factory was built with
431    ///   (its static token, or unauthenticated). Zero added cost.
432    /// - `user == Some(name)` → returns a factory whose upstream transport is
433    ///   **unauthenticated** and injects `header: name` on every request. The
434    ///   base factory's static token is intentionally dropped for these calls:
435    ///   the forwarded identity is the auth, not the proxy's own principal.
436    ///
437    /// `header` is the outgoing header name (typically the same
438    /// `x-forwarded-user` the proxy read the identity from). Returns an error if
439    /// `header` or `name` is not a valid HTTP header name / value.
440    ///
441    /// Native-only: on `wasm32` identity is forwarded by the browser session, so
442    /// this method is absent.
443    #[cfg(not(target_arch = "wasm32"))]
444    pub fn with_forwarded_user(&self, header: &str, user: Option<&str>) -> Result<Self> {
445        use reqwest::header::{HeaderName, HeaderValue};
446
447        let Some(name) = user else {
448            // Anonymous: keep the base factory (and its auth) untouched.
449            return Ok(self.clone());
450        };
451
452        let header_name = HeaderName::from_bytes(header.as_bytes()).map_err(|e| {
453            Error::invalid_config(format!(
454                "invalid forwarded-user header name `{header}`: {e}"
455            ))
456        })?;
457        let header_value = HeaderValue::from_str(name).map_err(|e| {
458            Error::invalid_config(format!("invalid forwarded-user header value: {e}"))
459        })?;
460
461        // A fresh unauthenticated transport (no bearer signer): the forwarded
462        // header is the auth for these calls. Route I/O onto the dedicated
463        // runtime when one is configured, mirroring `build_transport`.
464        let reqwest_client = reqwest::Client::new();
465        let base_service: Arc<dyn HttpService> =
466            Arc::new(ReqwestService::new(reqwest_client.clone()));
467        let transport =
468            Transport::new_unauthenticated().with_http_service(Arc::new(ForwardedHeaderService {
469                inner: base_service,
470                name: header_name,
471                value: header_value,
472            }));
473        let transport = match &self.io_handle {
474            Some(handle) => transport.with_runtime(handle.clone()),
475            None => transport,
476        };
477
478        // Rebuilding both clients is cheap — each just holds a `Transport` clone
479        // (Arc bumps) plus a `Url`.
480        let creds =
481            TemporaryCredentialClient::new_with_url(transport.clone(), self.base_url.clone());
482        let uc = UnityCatalogClient::new(transport, self.base_url.clone());
483        Ok(UnityObjectStoreFactory {
484            creds,
485            uc,
486            aws_region: self.aws_region.clone(),
487            io_handle: self.io_handle.clone(),
488            base_url: self.base_url.clone(),
489            allow_unauthenticated: self.allow_unauthenticated,
490        })
491    }
492
493    /// Borrow the underlying [`UnityCatalogClient`] for catalog metadata
494    /// operations (listing volumes, resolving table names, …).
495    pub fn unity_client(&self) -> &UnityCatalogClient {
496        &self.uc
497    }
498
499    /// Borrow the underlying credential-vending client. Most users want
500    /// [`for_url`](Self::for_url) / [`for_volume`](Self::for_volume) /
501    /// [`for_table`](Self::for_table) / [`for_path`](Self::for_path) instead.
502    pub fn credentials_client(&self) -> &TemporaryCredentialClient {
503        &self.creds
504    }
505
506    /// Build an [`UCStore`] for any supported URL.
507    ///
508    /// See [`UCReference`] for the supported URL grammar. Raw cloud URLs
509    /// (`s3://`, `gs://`, `abfss://`, …) are routed to
510    /// [`for_path`](Self::for_path).
511    pub async fn for_url(&self, url: &str, op: Operation) -> Result<UCStore> {
512        let reference = UCReference::parse(url)
513            .map_err(crate::error::Error::from)
514            .map_err(object_store::Error::from)?;
515        match reference {
516            UCReference::Volume {
517                catalog,
518                schema,
519                volume,
520                path,
521            } => {
522                let name = format!("{catalog}.{schema}.{volume}");
523                let store = self.for_volume(name, op.into_volume()).await?;
524                if path.is_empty() {
525                    Ok(store)
526                } else {
527                    Ok(extend_prefix(store, &path))
528                }
529            }
530            UCReference::Table {
531                catalog,
532                schema,
533                table,
534            } => {
535                let name = format!("{catalog}.{schema}.{table}");
536                self.for_table(name, op.into_table()).await
537            }
538            UCReference::Path(url) => self.for_path(&url, op.into_path()).await,
539        }
540    }
541
542    /// Vend credentials for a table and return a prefixed store rooted at
543    /// the table's storage location.
544    ///
545    /// The `table` argument accepts a `Uuid`, a [`String`] / `&str`
546    /// containing a three-level `<catalog>.<schema>.<table>` name, or any
547    /// [`TableReference`].
548    pub async fn for_table(
549        &self,
550        table: impl Into<TableReference>,
551        operation: TableOperation,
552    ) -> Result<UCStore> {
553        let table = table.into();
554        // For name references, resolve the storage location up front (a lookup
555        // name-based vending makes anyway). Two branches consume it before
556        // vending:
557        //   - a `file://` location has no cloud credential to vend → local store;
558        //   - on wasm under a proxy posture → route through the byte-proxy.
559        // A UUID-only reference skips both (a local-fs or proxied table by UUID
560        // is an unsupported edge case — use the three-level name).
561        if let TableReference::Name(name) = &table
562            && let Some(location) = self.table_storage_location(name).await?
563            && let Ok(url) = Url::parse(&location)
564        {
565            if url.scheme() == "file" {
566                let path_op = match operation {
567                    TableOperation::Read => PathOperation::Read,
568                    TableOperation::ReadWrite => PathOperation::ReadWrite,
569                };
570                return local_store(&url, path_op);
571            }
572            #[cfg(target_arch = "wasm32")]
573            if matches!(self.storage_access, proxy::StorageAccess::Proxy { .. }) {
574                let seg = format!("table:{name}");
575                return self.proxy_store(&seg, &url);
576            }
577        }
578        let (credential, table_id) = self
579            .creds
580            .temporary_table_credential(table, operation)
581            .await
582            .map_err(Error::from)?;
583        let securable = SecurableRef::Table(table_id, operation);
584        self.build_store(credential, securable).await
585    }
586
587    /// Look up a table's `storage_location` by its three-level name.
588    ///
589    /// Returns `None` when the table has no storage location set. Used by
590    /// [`for_table`](Self::for_table) to detect `file://`-backed tables before
591    /// vending.
592    async fn table_storage_location(&self, full_name: &str) -> Result<Option<String>> {
593        let table = self
594            .uc
595            .tables_client()
596            .get_table(&GetTableRequest {
597                full_name: full_name.to_string(),
598                include_browse: Some(false),
599                include_delta_metadata: Some(false),
600                include_manifest_capabilities: Some(false),
601                ..Default::default()
602            })
603            .await
604            .map_err(Error::from)?;
605        Ok(table.storage_location.filter(|s| !s.is_empty()))
606    }
607
608    /// Vend credentials for a volume and return a prefixed store rooted at
609    /// the volume's storage location.
610    ///
611    /// A volume whose storage location is a `file://` path is served by a local
612    /// [`LocalFileSystem`] store and never hits the credential-vending API —
613    /// local storage has no cloud credential to vend.
614    pub async fn for_volume(
615        &self,
616        volume: impl Into<VolumeReference>,
617        operation: VolumeOperation,
618    ) -> Result<UCStore> {
619        let volume = volume.into();
620        // As with `for_table`: a volume backed by local filesystem storage has
621        // no cloud credential to vend. Resolve its storage location up front
622        // (the same `GetVolume` lookup name-based vending makes) and, when it is
623        // `file://`, build a local store directly — skipping vending.
624        //
625        // Only possible for name references; a caller holding only the volume
626        // UUID still vends (a local-fs volume addressed by UUID is unsupported —
627        // use the three-level name).
628        if let VolumeReference::Name(name) = &volume
629            && let Some(location) = self.volume_storage_location(name).await?
630            && let Ok(url) = Url::parse(&location)
631        {
632            if url.scheme() == "file" {
633                let path_op = match operation {
634                    VolumeOperation::Read => PathOperation::Read,
635                    VolumeOperation::ReadWrite => PathOperation::ReadWrite,
636                };
637                return local_store(&url, path_op);
638            }
639            // Proxy posture (wasm): route volume bytes (read + write) through
640            // the server's storage byte-proxy.
641            #[cfg(target_arch = "wasm32")]
642            if matches!(self.storage_access, proxy::StorageAccess::Proxy { .. }) {
643                let seg = format!("vol:{name}");
644                return self.proxy_store(&seg, &url);
645            }
646        }
647        let (credential, volume_id) = self
648            .creds
649            .temporary_volume_credential(volume, operation)
650            .await
651            .map_err(Error::from)?;
652        let securable = SecurableRef::Volume(volume_id, operation);
653        self.build_store(credential, securable).await
654    }
655
656    /// Look up a volume's `storage_location` by its three-level name.
657    ///
658    /// Returns `None` when the volume has no storage location set. Used by
659    /// [`for_volume`](Self::for_volume) to detect `file://`-backed volumes
660    /// before vending.
661    async fn volume_storage_location(&self, name: &str) -> Result<Option<String>> {
662        let volume = self
663            .uc
664            .volumes_client()
665            .get_volume(&GetVolumeRequest {
666                name: name.to_string(),
667                include_browse: Some(false),
668                ..Default::default()
669            })
670            .await
671            .map_err(Error::from)?;
672        Ok(Some(volume.storage_location).filter(|s| !s.is_empty()))
673    }
674
675    /// Vend credentials for a raw cloud URL (`s3://`, `gs://`, `abfss://`,
676    /// …). Uses `temporary-path-credentials` under the hood.
677    ///
678    /// `file://` URLs are served by a local [`LocalFileSystem`] store and
679    /// never hit the credential-vending API — local storage has no cloud
680    /// credential to vend.
681    pub async fn for_path(&self, path: &Url, operation: PathOperation) -> Result<UCStore> {
682        if path.scheme() == "file" {
683            return local_store(path, operation);
684        }
685        // Proxy posture (wasm): route bytes for the raw cloud URL through the
686        // server's storage byte-proxy (`path:<pct-encoded url>` securable). The
687        // URL itself is the securable root, so the routing prefix is its path.
688        #[cfg(target_arch = "wasm32")]
689        if matches!(self.storage_access, proxy::StorageAccess::Proxy { .. }) {
690            let seg = proxy::path_securable(path);
691            return self.proxy_store(&seg, path);
692        }
693        let (credential, _resolved) = self
694            .creds
695            .temporary_path_credential(path.clone(), operation, false)
696            .await
697            .map_err(Error::from)?;
698        let securable = SecurableRef::Path(path.clone(), operation, Some(false));
699        self.build_store(credential, securable).await
700    }
701
702    /// Vend credentials for a raw cloud URL with `dry_run` set to true.
703    /// The server validates that credentials *could* be issued but the
704    /// returned token is not usable for IO; useful for permission probes.
705    ///
706    /// For `file://` URLs there is nothing to probe — a local store is
707    /// returned directly, identical to [`for_path`](Self::for_path).
708    pub async fn dry_run_path(&self, path: &Url, operation: PathOperation) -> Result<UCStore> {
709        if path.scheme() == "file" {
710            return local_store(path, operation);
711        }
712        let (credential, _resolved) = self
713            .creds
714            .temporary_path_credential(path.clone(), operation, true)
715            .await
716            .map_err(Error::from)?;
717        let securable = SecurableRef::Path(path.clone(), operation, Some(true));
718        self.build_store(credential, securable).await
719    }
720
721    /// Build a proxy-backed [`UCStore`] for `securable_seg` (a typed
722    /// `{securable}` segment) whose data lives at `location` (the real cloud
723    /// URL of the securable root).
724    ///
725    /// Used on wasm when the server announces `storageAccess: "proxy"`: no
726    /// credential is vended here (the server vends internally). `url` stays the
727    /// real cloud URL so the read path's routing/log-discovery math is
728    /// unchanged; `root` is a proxy `HttpStore` wrapped to strip the bucket
729    /// prefix so forwarded bucket-rooted keys become securable-relative.
730    #[cfg(target_arch = "wasm32")]
731    fn proxy_store(&self, securable_seg: &str, location: &Url) -> Result<UCStore> {
732        let proxy::StorageAccess::Proxy { base } = &self.storage_access else {
733            // Only called from the proxy branches, which are gated on `Proxy`.
734            return Err(Error::invalid_config(
735                "proxy_store called without a proxy posture".to_string(),
736            )
737            .into());
738        };
739        // The bucket-relative prefix of the securable root — the same value
740        // `build_store` records as `UCStore.path`, and exactly what the routing
741        // store forwards ahead of the object key.
742        let path = match parse_azurite(location) {
743            Some(loc) => Path::from(loc.prefix),
744            None => Path::from_url_path(location.path())?,
745        };
746        let root = proxy::build_proxy_store(base, securable_seg, path.clone(), self.token())?;
747        Ok(UCStore {
748            root,
749            url: location.clone(),
750            path,
751        })
752    }
753
754    /// The bearer token this factory authenticates with, if any. On wasm the
755    /// proxy store relays it to the same-origin proxy endpoint.
756    #[cfg(target_arch = "wasm32")]
757    fn token(&self) -> Option<&str> {
758        self.token.as_deref()
759    }
760
761    async fn build_store(
762        &self,
763        credential: TemporaryCredential,
764        securable: SecurableRef,
765    ) -> Result<UCStore> {
766        let url = Url::parse(&credential.url).map_err(Error::from)?;
767        // The emulator store is rooted at the container, so the prefix is the
768        // blob path *within* the container — not the full URL path (which for
769        // path-style Azurite also carries `/<account>/<container>`).
770        let path = match parse_azurite(&url) {
771            Some(loc) => Path::from(loc.prefix),
772            None => Path::from_url_path(url.path())?,
773        };
774        let store = self.to_store(credential, securable).await?;
775        Ok(UCStore {
776            root: store,
777            url,
778            path,
779        })
780    }
781
782    async fn to_store(
783        &self,
784        credential: TemporaryCredential,
785        securable: SecurableRef,
786    ) -> Result<Arc<dyn ObjectStore>> {
787        if as_azure(&credential).is_ok() {
788            let url = Url::parse(&credential.url).map_err(Error::from)?;
789
790            // Azurite (local Blob emulator): the URL is path-style and
791            // `with_url` does not understand it, and in emulator mode the
792            // builder ignores a `with_credentials` provider — it only honours
793            // an explicit access key / bearer / SAS. So set account + container
794            // explicitly and pass the vended SAS via `SasKey`.
795            if let Some(loc) = parse_azurite(&url) {
796                let sas = azure_sas_token(&credential).ok_or_else(|| {
797                    Error::invalid_config(
798                        "Azurite store requires a SAS-token credential".to_string(),
799                    )
800                })?;
801                #[allow(unused_mut)]
802                let mut builder = MicrosoftAzureBuilder::new()
803                    .with_use_emulator(true)
804                    .with_account(loc.account)
805                    .with_container_name(loc.container)
806                    .with_config(object_store::azure::AzureConfigKey::SasKey, sas);
807                // Native: route object-store HTTP onto the dedicated I/O runtime
808                // when one was supplied. On wasm, use object_store's default
809                // browser-Fetch connector (no runtime to route onto).
810                #[cfg(not(target_arch = "wasm32"))]
811                if let Some(handle) = &self.io_handle {
812                    builder =
813                        builder.with_http_connector(SpawnedReqwestConnector::new(handle.clone()));
814                }
815                // Belt-and-braces: object_store's retry path panics on wasm
816                // without a timer, so disable retries on the browser store.
817                #[cfg(target_arch = "wasm32")]
818                {
819                    builder = builder.with_retry(object_store::RetryConfig {
820                        max_retries: 0,
821                        ..Default::default()
822                    });
823                }
824                let store = builder.build()?;
825                return Ok(Arc::new(store));
826            }
827
828            let provider = new_azure(self.creds.clone(), &credential, securable).await?;
829            #[allow(unused_mut)]
830            let mut builder = MicrosoftAzureBuilder::new()
831                .with_url(url.to_string())
832                .with_credentials(Arc::new(provider));
833            #[cfg(not(target_arch = "wasm32"))]
834            if let Some(handle) = &self.io_handle {
835                builder = builder.with_http_connector(SpawnedReqwestConnector::new(handle.clone()));
836            }
837            #[cfg(target_arch = "wasm32")]
838            {
839                builder = builder.with_retry(object_store::RetryConfig {
840                    max_retries: 0,
841                    ..Default::default()
842                });
843            }
844            let store = builder.build()?;
845            return Ok(Arc::new(store));
846        }
847
848        #[cfg(not(target_arch = "wasm32"))]
849        if as_aws(&credential).is_ok() {
850            let access_point = aws_access_point(&credential);
851            let provider = new_aws(self.creds.clone(), &credential, securable).await?;
852            let url = Url::parse(&credential.url).map_err(Error::from)?;
853            let mut builder = AmazonS3Builder::new()
854                .with_url(url.to_string())
855                .with_credentials(Arc::new(provider));
856            // Prefer an explicit override; otherwise honour `AWS_REGION`
857            // before falling back to the object_store default.
858            if let Some(region) = self
859                .aws_region
860                .clone()
861                .or_else(|| std::env::var("AWS_REGION").ok())
862            {
863                builder = builder.with_region(region);
864            }
865            // Where the server returns an S3 access-point ARN, use it as the
866            // bucket so SigV4 signatures match what STS authorised.
867            if let Some(ap) = access_point {
868                builder = builder.with_bucket_name(ap);
869            }
870            if let Some(handle) = &self.io_handle {
871                builder = builder.with_http_connector(SpawnedReqwestConnector::new(handle.clone()));
872            }
873            let store = builder.build()?;
874            return Ok(Arc::new(store));
875        }
876
877        #[cfg(not(target_arch = "wasm32"))]
878        if as_gcp(&credential).is_ok() {
879            let provider = new_gcp(self.creds.clone(), &credential, securable).await?;
880            let url = Url::parse(&credential.url).map_err(Error::from)?;
881            let mut builder = GoogleCloudStorageBuilder::new()
882                .with_url(url.to_string())
883                .with_credentials(Arc::new(provider));
884            if let Some(handle) = &self.io_handle {
885                builder = builder.with_http_connector(SpawnedReqwestConnector::new(handle.clone()));
886            }
887            let store = builder.build()?;
888            return Ok(Arc::new(store));
889        }
890
891        // wasm is Azure-first: AWS/GCP stores are gated out above, so a
892        // non-Azure credential here is unsupported rather than unmatched.
893        #[cfg(target_arch = "wasm32")]
894        if as_aws(&credential).is_ok() || as_gcp(&credential).is_ok() {
895            return Err(Error::invalid_config(
896                "AWS/GCP object stores are not supported on wasm (Azure-first)",
897            )
898            .into());
899        }
900
901        Err(
902            Error::InvalidCredential("Failed to match credential with storage type".to_string())
903                .into(),
904        )
905    }
906}
907
908/// Unified read/write operation used by [`UnityObjectStoreFactory::for_url`].
909///
910/// The factory translates this to the operation enum expected by each
911/// vending endpoint:
912///
913/// | `Operation`  | Volume        | Table       | Path             |
914/// |--------------|---------------|-------------|------------------|
915/// | `Read`       | `READ_VOLUME` | `READ`      | `PATH_READ`      |
916/// | `ReadWrite`  | `WRITE_VOLUME`| `READ_WRITE`| `PATH_READ_WRITE`|
917#[derive(Debug, Clone, Copy, PartialEq, Eq)]
918pub enum Operation {
919    Read,
920    ReadWrite,
921}
922
923impl Operation {
924    fn into_volume(self) -> VolumeOperation {
925        match self {
926            Operation::Read => VolumeOperation::Read,
927            Operation::ReadWrite => VolumeOperation::ReadWrite,
928        }
929    }
930
931    fn into_table(self) -> TableOperation {
932        match self {
933            Operation::Read => TableOperation::Read,
934            Operation::ReadWrite => TableOperation::ReadWrite,
935        }
936    }
937
938    fn into_path(self) -> PathOperation {
939        match self {
940            Operation::Read => PathOperation::Read,
941            Operation::ReadWrite => PathOperation::ReadWrite,
942        }
943    }
944}
945
946impl From<Operation> for TableOperation {
947    fn from(op: Operation) -> Self {
948        op.into_table()
949    }
950}
951
952impl From<Operation> for VolumeOperation {
953    fn from(op: Operation) -> Self {
954        op.into_volume()
955    }
956}
957
958impl From<Operation> for PathOperation {
959    fn from(op: Operation) -> Self {
960        op.into_path()
961    }
962}
963
964/// Build a [`UCStore`] backed by the local filesystem for a `file://` URL,
965/// bypassing credential vending entirely.
966///
967/// Mirrors the bucket-root + prefix model of vended cloud stores: `root` is an
968/// unrooted [`LocalFileSystem`] (paths are absolute, relative to the filesystem
969/// root) and `path` is the URL's directory. This keeps both consumption modes
970/// correct:
971///
972/// - [`UCStore::as_dyn`] wraps `root` in a `PrefixStore` at the directory, so
973///   callers address paths *relative* to the `file://` directory; and
974/// - [`UCStore::root`] is the unrooted store, so the DataFusion routing store
975///   — which forwards the full request path unchanged — resolves absolute
976///   table paths.
977///
978/// For [`PathOperation::ReadWrite`] / [`PathOperation::CreateTable`] the target
979/// directory is created if missing, so writes to a fresh local root succeed.
980///
981/// # Platform support
982///
983/// Local `file://` storage is **POSIX-only** for now. On Windows it returns an
984/// error: `object_store::path::Path` cannot represent a Windows absolute path
985/// (the drive-letter colon is percent-encoded), so an unrooted
986/// [`LocalFileSystem`] addressed by full path does not resolve back to the real
987/// on-disk location. Tracked for a follow-up; cloud schemes are unaffected.
988///
989/// Not available on `wasm32`: object_store has no `LocalFileSystem` there and a
990/// browser has no local filesystem — a `file://` reference returns an error.
991#[cfg(not(target_arch = "wasm32"))]
992fn local_store(url: &Url, operation: PathOperation) -> Result<UCStore> {
993    if cfg!(windows) {
994        return Err(Error::invalid_config(format!(
995            "local (file://) storage is not supported on Windows: {url}"
996        ))
997        .into());
998    }
999
1000    let dir = url
1001        .to_file_path()
1002        .map_err(|_| Error::invalid_url(format!("not a valid local file path: {url}")))?;
1003
1004    if matches!(
1005        operation,
1006        PathOperation::ReadWrite | PathOperation::CreateTable
1007    ) {
1008        std::fs::create_dir_all(&dir).map_err(|e| {
1009            Error::invalid_config(format!("failed to create local directory {dir:?}: {e}"))
1010        })?;
1011    }
1012
1013    // Unrooted: object_store `Path`s are relative to the filesystem root, so a
1014    // full path like `tmp/data/part-0` resolves to `/tmp/data/part-0`. The
1015    // directory is carried as the `UCStore` prefix instead (see `build_store`).
1016    let store = LocalFileSystem::new();
1017    let path = Path::from_url_path(url.path())?;
1018    Ok(UCStore {
1019        root: Arc::new(store),
1020        url: url.clone(),
1021        path,
1022    })
1023}
1024
1025/// A [`CredentialProvider`] that always hands back a fixed, already-materialized
1026/// credential. Unlike `object_store::StaticCredentialProvider` (which takes the
1027/// credential by value), this holds an `Arc<T>` directly — the shape
1028/// [`as_azure`]/[`as_aws`]/[`as_gcp`] already produce — so no `Clone` on the
1029/// (non-`Clone`) cloud credential types is required.
1030#[cfg(not(target_arch = "wasm32"))]
1031#[derive(Debug)]
1032struct StaticArcProvider<T>(Arc<T>);
1033
1034#[cfg(not(target_arch = "wasm32"))]
1035#[async_trait::async_trait]
1036impl<T: std::fmt::Debug + Send + Sync> object_store::CredentialProvider for StaticArcProvider<T> {
1037    type Credential = T;
1038    async fn get_credential(&self) -> Result<Arc<T>> {
1039        Ok(self.0.clone())
1040    }
1041}
1042
1043/// Build a [`UCStore`] from a credential that has **already been vended**
1044/// out-of-band, using a *static* credential provider (no Unity Catalog client,
1045/// no auto-refresh).
1046///
1047/// This is the escape hatch for an in-process caller — e.g. a server-side storage
1048/// byte-proxy — that authorizes and vends the credential itself and then only
1049/// needs it turned into an [`ObjectStore`]. Unlike [`UnityObjectStoreFactory`]'s
1050/// `for_*` entry points, this takes no factory and constructs no
1051/// [`UnityCatalogClient`]: the store is bound to the exact `credential` passed,
1052/// valid for that credential's lifetime. A single short-lived proxy request never
1053/// outlives the vended credential, so the lack of refresh is fine; a caller that
1054/// needs long-lived refresh should use the factory `for_*` methods instead.
1055///
1056/// The returned [`UCStore`] is scoped exactly as [`UnityObjectStoreFactory::for_path`]
1057/// would produce it; call [`UCStore::as_dyn`] for a store prefixed at the
1058/// credential-scoped root. `io_handle` (when `Some`) routes object-store HTTP onto
1059/// a dedicated I/O runtime; `aws_region` overrides the S3 region (else `AWS_REGION`
1060/// / the object_store default).
1061///
1062/// A `file://` URL yields a [`LocalFileSystem`] store and needs no credential.
1063#[cfg(not(target_arch = "wasm32"))]
1064pub fn store_from_vended_credential(
1065    credential: &TemporaryCredential,
1066    io_handle: Option<tokio::runtime::Handle>,
1067    aws_region: Option<String>,
1068) -> Result<UCStore> {
1069    let url = Url::parse(&credential.url).map_err(Error::from)?;
1070
1071    // Local storage: no cloud credential to apply.
1072    if url.scheme() == "file" {
1073        let store = LocalFileSystem::new();
1074        let path = Path::from_url_path(url.path())?;
1075        return Ok(UCStore {
1076            root: Arc::new(store),
1077            url,
1078            path,
1079        });
1080    }
1081
1082    // The credential-scoped prefix within the bucket/container root. For
1083    // path-style Azurite the store is rooted at the container, so the prefix is
1084    // the blob path only (not the full `/<account>/<container>/...` URL path).
1085    let path = match parse_azurite(&url) {
1086        Some(loc) => Path::from(loc.prefix),
1087        None => Path::from_url_path(url.path())?,
1088    };
1089
1090    let root: Arc<dyn ObjectStore> = if as_azure(credential).is_ok() {
1091        // Azurite: the emulator ignores a credentials provider — set account +
1092        // container explicitly and pass the vended SAS via `SasKey` (identical to
1093        // the factory path).
1094        if let Some(loc) = parse_azurite(&url) {
1095            let sas = azure_sas_token(credential).ok_or_else(|| {
1096                Error::invalid_config("Azurite store requires a SAS-token credential".to_string())
1097            })?;
1098            let mut builder = MicrosoftAzureBuilder::new()
1099                .with_use_emulator(true)
1100                .with_account(loc.account)
1101                .with_container_name(loc.container)
1102                .with_config(object_store::azure::AzureConfigKey::SasKey, sas);
1103            if let Some(handle) = &io_handle {
1104                builder = builder.with_http_connector(SpawnedReqwestConnector::new(handle.clone()));
1105            }
1106            Arc::new(builder.build()?)
1107        } else {
1108            let token = as_azure(credential)?.token;
1109            let mut builder = MicrosoftAzureBuilder::new()
1110                .with_url(url.to_string())
1111                .with_credentials(Arc::new(StaticArcProvider(token)));
1112            if let Some(handle) = &io_handle {
1113                builder = builder.with_http_connector(SpawnedReqwestConnector::new(handle.clone()));
1114            }
1115            Arc::new(builder.build()?)
1116        }
1117    } else if as_aws(credential).is_ok() {
1118        let access_point = aws_access_point(credential);
1119        let token = as_aws(credential)?.token;
1120        let mut builder = AmazonS3Builder::new()
1121            .with_url(url.to_string())
1122            .with_credentials(Arc::new(StaticArcProvider(token)));
1123        if let Some(region) = aws_region.or_else(|| std::env::var("AWS_REGION").ok()) {
1124            builder = builder.with_region(region);
1125        }
1126        if let Some(ap) = access_point {
1127            builder = builder.with_bucket_name(ap);
1128        }
1129        if let Some(handle) = &io_handle {
1130            builder = builder.with_http_connector(SpawnedReqwestConnector::new(handle.clone()));
1131        }
1132        Arc::new(builder.build()?)
1133    } else if as_gcp(credential).is_ok() {
1134        let token = as_gcp(credential)?.token;
1135        let mut builder = GoogleCloudStorageBuilder::new()
1136            .with_url(url.to_string())
1137            .with_credentials(Arc::new(StaticArcProvider(token)));
1138        if let Some(handle) = &io_handle {
1139            builder = builder.with_http_connector(SpawnedReqwestConnector::new(handle.clone()));
1140        }
1141        Arc::new(builder.build()?)
1142    } else {
1143        return Err(Error::InvalidCredential(
1144            "Failed to match credential with storage type".to_string(),
1145        )
1146        .into());
1147    };
1148
1149    Ok(UCStore { root, url, path })
1150}
1151
1152/// wasm has no local filesystem; a `file://` reference is unsupported.
1153#[cfg(target_arch = "wasm32")]
1154fn local_store(url: &Url, _operation: PathOperation) -> Result<UCStore> {
1155    Err(Error::invalid_config(format!(
1156        "local (file://) storage is not supported on wasm: {url}"
1157    ))
1158    .into())
1159}
1160
1161/// Parsed components of an Azurite (Azure Blob emulator) storage URL.
1162///
1163/// Azurite is path-style — the account, container, and blob prefix are all
1164/// segments of the path rather than the host (real Azure encodes the account in
1165/// the host: `https://<account>.blob.core.windows.net/<container>/<prefix>`).
1166struct AzuriteLocation {
1167    account: String,
1168    container: String,
1169    prefix: String,
1170}
1171
1172/// Detect and parse an Azurite endpoint from a storage URL.
1173///
1174/// Recognised forms (mirroring the server's `is_azurite` / `StorageLocationUrl`):
1175/// - `http://127.0.0.1:10000/<account>/<container>/<prefix>` (default Azurite
1176///   blob endpoint on localhost / 127.0.0.1, port 10000), and
1177/// - `azurite://<container>/<prefix>` (custom scheme; the account is implicit,
1178///   defaulting to the well-known emulator account).
1179///
1180/// Returns `None` for any non-Azurite URL so the caller falls through to the
1181/// real-Azure path. `object_store`'s `MicrosoftAzureBuilder::with_url` does not
1182/// understand either form, so the Azurite branch must set account/container
1183/// explicitly rather than going through `with_url`.
1184fn parse_azurite(url: &Url) -> Option<AzuriteLocation> {
1185    const EMULATOR_ACCOUNT: &str = "devstoreaccount1";
1186
1187    if url.scheme() == "azurite" {
1188        // azurite://<container>/<prefix>
1189        let container = url.host_str().filter(|s| !s.is_empty())?.to_owned();
1190        let prefix = url.path().trim_start_matches('/').to_owned();
1191        return Some(AzuriteLocation {
1192            account: EMULATOR_ACCOUNT.to_owned(),
1193            container,
1194            prefix,
1195        });
1196    }
1197
1198    let is_localhost = matches!(url.host_str(), Some("localhost") | Some("127.0.0.1"));
1199    if url.scheme() == "http" && is_localhost && url.port() == Some(10000) {
1200        // http://127.0.0.1:10000/<account>/<container>/<prefix>
1201        let mut segments = url.path_segments()?;
1202        let account = segments.next().filter(|s| !s.is_empty())?.to_owned();
1203        let container = segments.next().filter(|s| !s.is_empty())?.to_owned();
1204        let prefix = segments.collect::<Vec<_>>().join("/");
1205        return Some(AzuriteLocation {
1206            account,
1207            container,
1208            prefix,
1209        });
1210    }
1211
1212    None
1213}
1214
1215/// Extract the raw SAS query string (`k1=v1&k2=v2&…`) from a vended credential,
1216/// if it carries an Azure User Delegation / service SAS.
1217fn azure_sas_token(credential: &TemporaryCredential) -> Option<String> {
1218    use unitycatalog_common::temporary_credentials::v1::temporary_credential::Credentials;
1219    match credential.credentials.as_ref()? {
1220        Credentials::AzureUserDelegationSas(sas) => Some(sas.sas_token.clone()),
1221        _ => None,
1222    }
1223}
1224
1225/// Returns a [`UCStore`] whose prefix is `store.prefix() + extra`, leaving
1226/// the underlying bucket-rooted store untouched.
1227fn extend_prefix(store: UCStore, extra: &str) -> UCStore {
1228    let mut url = store.url.clone();
1229    // Append the extra path component(s), keeping the trailing slash.
1230    {
1231        let mut segs = url.path_segments_mut().expect("cloud URL has a path");
1232        segs.pop_if_empty();
1233        for part in extra.split('/').filter(|p| !p.is_empty()) {
1234            segs.push(part);
1235        }
1236    }
1237    let new_path = if store.path.as_ref().is_empty() {
1238        Path::from(extra)
1239    } else {
1240        let base = store.path.as_ref().trim_end_matches('/');
1241        let extra = extra.trim_start_matches('/');
1242        Path::from(format!("{base}/{extra}"))
1243    };
1244    UCStore {
1245        root: store.root,
1246        url,
1247        path: new_path,
1248    }
1249}
1250
1251// Tests use tokio/tempfile/mockito and the native transport; there is no wasm
1252// test runner in this crate's CI (the wasm build is proven in the query-wasm
1253// workspace).
1254#[cfg(all(test, not(target_arch = "wasm32")))]
1255mod tests {
1256    use futures::TryStreamExt;
1257    // `put`/`PutPayload` are only used by the POSIX-only local-storage tests.
1258    #[cfg(not(windows))]
1259    use object_store::{ObjectStoreExt, PutPayload};
1260
1261    use super::*;
1262
1263    /// A factory whose UC endpoint points nowhere reachable. Any code path that
1264    /// tries to vend a credential will fail to connect — so a successful local
1265    /// operation proves vending was skipped entirely.
1266    async fn offline_factory() -> UnityObjectStoreFactory {
1267        UnityObjectStoreFactory::builder()
1268            // An unroutable, non-listening endpoint: a vend attempt cannot succeed.
1269            .with_uri("http://127.0.0.1:0/api/2.1/unity-catalog/")
1270            .with_allow_unauthenticated(true)
1271            .build()
1272            .await
1273            .unwrap()
1274    }
1275
1276    /// `file://` URLs are served by a local store and round-trip read/write/list
1277    /// without any credential-vending call (the factory points at a dead endpoint).
1278    // Local file:// storage is POSIX-only for now (see `local_store`).
1279    #[cfg(not(windows))]
1280    #[tokio::test]
1281    async fn for_url_file_roundtrips_without_vending() {
1282        let dir = tempfile::tempdir().unwrap();
1283        let url = Url::from_directory_path(dir.path()).unwrap();
1284
1285        let factory = offline_factory().await;
1286        let store = factory
1287            .for_url(url.as_str(), Operation::ReadWrite)
1288            .await
1289            .unwrap();
1290
1291        let dyn_store = store.as_dyn();
1292        let path = object_store::path::Path::from("hello.txt");
1293        dyn_store
1294            .put(&path, PutPayload::from_static(b"world"))
1295            .await
1296            .unwrap();
1297
1298        let listing: Vec<_> = dyn_store.list(None).try_collect().await.unwrap();
1299        assert_eq!(listing.len(), 1, "expected exactly one object");
1300        assert_eq!(listing[0].location, path);
1301
1302        let got = dyn_store.get(&path).await.unwrap().bytes().await.unwrap();
1303        assert_eq!(&got[..], b"world");
1304    }
1305
1306    /// `for_path` with a `file://` URL returns a usable store with no network call.
1307    #[cfg(not(windows))]
1308    #[tokio::test]
1309    async fn for_path_file_skips_vending() {
1310        let dir = tempfile::tempdir().unwrap();
1311        let url = Url::from_directory_path(dir.path()).unwrap();
1312
1313        let factory = offline_factory().await;
1314        // Read-only: the directory already exists, nothing is created.
1315        let store = factory.for_path(&url, PathOperation::Read).await.unwrap();
1316        // Empty directory lists to nothing — and crucially, no vend was attempted.
1317        let listing: Vec<_> = store.as_dyn().list(None).try_collect().await.unwrap();
1318        assert!(listing.is_empty());
1319        assert_eq!(store.url(), &url);
1320    }
1321
1322    /// A read-write local store auto-creates a missing target directory.
1323    #[cfg(not(windows))]
1324    #[tokio::test]
1325    async fn local_store_read_write_creates_dir() {
1326        let dir = tempfile::tempdir().unwrap();
1327        let missing = dir.path().join("not-yet-here");
1328        let url = Url::from_directory_path(&missing).unwrap();
1329
1330        let store = local_store(&url, PathOperation::ReadWrite).unwrap();
1331        assert!(missing.exists(), "ReadWrite must create the root directory");
1332
1333        let path = object_store::path::Path::from("a.bin");
1334        store
1335            .as_dyn()
1336            .put(&path, PutPayload::from_static(b"x"))
1337            .await
1338            .unwrap();
1339        assert!(missing.join("a.bin").exists());
1340    }
1341
1342    /// The unrooted `root()` store resolves the **full** path (the DataFusion
1343    /// routing-store contract): a write addressed by the full absolute path
1344    /// round-trips through that same path, and lands on disk under the table
1345    /// directory.
1346    ///
1347    /// On-disk placement is checked with a real `read_dir` of the table
1348    /// directory, not a native `PathBuf::join` against an object_store `Path`.
1349    /// POSIX-only: local file:// is gated off on Windows (see `local_store`),
1350    /// where an unrooted store cannot round-trip a drive-letter absolute path.
1351    #[cfg(not(windows))]
1352    #[tokio::test]
1353    async fn local_store_root_resolves_full_path() {
1354        let dir = tempfile::tempdir().unwrap();
1355        let table_dir = dir.path().join("mytable");
1356        let url = Url::from_directory_path(&table_dir).unwrap();
1357
1358        let store = local_store(&url, PathOperation::ReadWrite).unwrap();
1359
1360        // `prefix()` is the table directory; the routing store registers and
1361        // forwards the full path beneath it unchanged.
1362        let full = Path::from(format!("{}/part-0.parquet", store.prefix()));
1363        store
1364            .root()
1365            .put(&full, PutPayload::from_static(b"data"))
1366            .await
1367            .unwrap();
1368
1369        // Reading back the same full path proves the unrooted store resolves it
1370        // consistently — exactly what the DataFusion routing store relies on.
1371        let got = store
1372            .root()
1373            .get(&full)
1374            .await
1375            .unwrap()
1376            .bytes()
1377            .await
1378            .unwrap();
1379        assert_eq!(&got[..], b"data");
1380
1381        // And it landed on disk under the table directory. Check the real
1382        // filesystem directly (not via an object_store `Path`) so the assertion
1383        // is OS-agnostic.
1384        let entries: Vec<_> = std::fs::read_dir(&table_dir)
1385            .unwrap()
1386            .map(|e| e.unwrap().file_name())
1387            .collect();
1388        assert_eq!(entries, vec![std::ffi::OsString::from("part-0.parquet")]);
1389    }
1390
1391    /// A non-`file` URL handed to the local helper is a clean error, not a panic.
1392    /// POSIX-only: on Windows `local_store` rejects every call up front, so the
1393    /// specific "not a valid local file path" message does not apply.
1394    #[cfg(not(windows))]
1395    #[test]
1396    fn local_store_rejects_non_file_url() {
1397        let url = Url::parse("s3://bucket/prefix/").unwrap();
1398        match local_store(&url, PathOperation::Read) {
1399            Ok(_) => panic!("expected an error for a non-file URL"),
1400            Err(e) => assert!(
1401                e.to_string().contains("not a valid local file path"),
1402                "unexpected error: {e}"
1403            ),
1404        }
1405    }
1406
1407    /// Building a factory with a dedicated I/O runtime handle succeeds and the
1408    /// handle is carried through to the factory. The `None` path (no handle) is
1409    /// exercised everywhere else and must remain the default.
1410    ///
1411    /// A plain `#[test]` (not `#[tokio::test]`) so the dedicated I/O `Runtime`
1412    /// can be dropped here — dropping a `Runtime` inside an async context panics.
1413    #[test]
1414    fn build_with_io_runtime_carries_handle() {
1415        // A dedicated I/O runtime — the canonical segregation pattern (mirrors
1416        // object_store's own spawn-connector test).
1417        let io_runtime = tokio::runtime::Builder::new_current_thread()
1418            .enable_all()
1419            .build()
1420            .unwrap();
1421        let handle = io_runtime.handle().clone();
1422
1423        // Drive `build()` on a separate, lightweight runtime so this test owns
1424        // (and can safely drop) `io_runtime` itself.
1425        let driver = tokio::runtime::Builder::new_current_thread()
1426            .enable_all()
1427            .build()
1428            .unwrap();
1429
1430        driver.block_on(async {
1431            let factory = UnityObjectStoreFactory::builder()
1432                .with_uri("http://localhost:8080/api/2.1/unity-catalog/")
1433                .with_allow_unauthenticated(true)
1434                .with_io_runtime(handle.clone())
1435                .build()
1436                .await
1437                .unwrap();
1438            assert!(factory.io_handle.is_some());
1439
1440            // Clearing via `None` is honoured.
1441            let factory = UnityObjectStoreFactory::builder()
1442                .with_uri("http://localhost:8080/api/2.1/unity-catalog/")
1443                .with_allow_unauthenticated(true)
1444                .with_io_runtime(handle.clone())
1445                .with_io_runtime(None)
1446                .build()
1447                .await
1448                .unwrap();
1449            assert!(factory.io_handle.is_none());
1450        });
1451    }
1452
1453    /// Live test against a Databricks workspace. Requires
1454    /// `DATABRICKS_HOST` + `DATABRICKS_TOKEN` to be set. Marked `#[ignore]`
1455    /// because CI shouldn't hit a real workspace.
1456    ///
1457    /// The calling runtime is a current-thread runtime with **I/O disabled**
1458    /// (no `enable_all`); a successful `list` therefore proves every request
1459    /// was spawned onto the separate, I/O-enabled runtime via the connector —
1460    /// exactly the assertion object_store's own spawn-connector test makes.
1461    #[test]
1462    #[ignore]
1463    fn list_store_via_temp_credential_on_io_runtime() {
1464        let databricks_host = std::env::var("DATABRICKS_HOST").unwrap();
1465        let databricks_token = std::env::var("DATABRICKS_TOKEN").unwrap();
1466
1467        // Dedicated, I/O-enabled runtime on its own thread.
1468        let io_runtime = std::thread::spawn(|| {
1469            tokio::runtime::Builder::new_multi_thread()
1470                .enable_all()
1471                .build()
1472                .unwrap()
1473        })
1474        .join()
1475        .unwrap();
1476        let io_handle = io_runtime.handle().clone();
1477
1478        // Calling runtime: current-thread, no I/O driver enabled. Any request
1479        // not spawned onto `io_handle` would panic ("no reactor running").
1480        let main_runtime = tokio::runtime::Builder::new_current_thread()
1481            .build()
1482            .unwrap();
1483
1484        main_runtime.block_on(async move {
1485            let factory = UnityObjectStoreFactory::builder()
1486                .with_uri(format!("{databricks_host}/api/2.1/unity-catalog/"))
1487                .with_token(databricks_token)
1488                .with_aws_region("eu-north-1".to_string())
1489                .with_io_runtime(io_handle)
1490                .build()
1491                .await
1492                .unwrap();
1493
1494            let volume_path = url::Url::parse("s3://open-lakehouse-dev/volumes/").unwrap();
1495            let store = factory
1496                .for_path(&volume_path, PathOperation::Read)
1497                .await
1498                .unwrap();
1499            let files: Vec<_> = store.as_dyn().list(None).try_collect().await.unwrap();
1500            println!("files: {files:?}");
1501        });
1502    }
1503
1504    #[test]
1505    fn parse_azurite_path_style() {
1506        let url =
1507            Url::parse("http://127.0.0.1:10000/devstoreaccount1/mycontainer/some/prefix").unwrap();
1508        let loc = parse_azurite(&url).expect("should parse path-style azurite URL");
1509        assert_eq!(loc.account, "devstoreaccount1");
1510        assert_eq!(loc.container, "mycontainer");
1511        assert_eq!(loc.prefix, "some/prefix");
1512    }
1513
1514    #[test]
1515    fn parse_azurite_custom_scheme() {
1516        let url = Url::parse("azurite://mycontainer/some/prefix").unwrap();
1517        let loc = parse_azurite(&url).expect("should parse azurite:// URL");
1518        assert_eq!(loc.account, "devstoreaccount1");
1519        assert_eq!(loc.container, "mycontainer");
1520        assert_eq!(loc.prefix, "some/prefix");
1521    }
1522
1523    #[test]
1524    fn parse_azurite_localhost_alias() {
1525        let url = Url::parse("http://localhost:10000/acct/cont/p").unwrap();
1526        let loc = parse_azurite(&url).expect("localhost is also an azurite host");
1527        assert_eq!(loc.account, "acct");
1528        assert_eq!(loc.container, "cont");
1529        assert_eq!(loc.prefix, "p");
1530    }
1531
1532    #[test]
1533    fn parse_azurite_rejects_real_cloud_urls() {
1534        // Real Azure, S3, and a non-Azurite localhost port must all be `None`.
1535        for raw in [
1536            "https://acct.blob.core.windows.net/container/path",
1537            "s3://bucket/prefix",
1538            "http://127.0.0.1:9000/acct/cont/p", // not the Azurite port
1539        ] {
1540            let url = Url::parse(raw).unwrap();
1541            assert!(parse_azurite(&url).is_none(), "should not match: {raw}");
1542        }
1543    }
1544
1545    /// Building a store from an Azurite SAS credential targets the emulator and
1546    /// roots the prefix at the blob path within the container (not the full URL
1547    /// path, which also carries `/<account>/<container>`).
1548    #[tokio::test]
1549    async fn azurite_store_targets_emulator_and_roots_prefix() {
1550        use unitycatalog_common::temporary_credentials::v1::AzureUserDelegationSas;
1551
1552        let url = "http://127.0.0.1:10000/devstoreaccount1/mycontainer/tbl/data";
1553        let credential = TemporaryCredential {
1554            expiration_time: now_epoch_millis() + 3_600_000,
1555            url: url.to_string(),
1556            credentials: Some(
1557                AzureUserDelegationSas {
1558                    // A minimal, well-formed SAS query string. `split_sas` parses it
1559                    // into query pairs; the emulator never sees it in this test.
1560                    sas_token: "sv=2021-08-06&ss=b&srt=co&sp=rl&se=2999-01-01T00:00:00Z&sig=AAAA"
1561                        .to_string(),
1562                    ..Default::default()
1563                }
1564                .into(),
1565            ),
1566            ..Default::default()
1567        };
1568
1569        let factory = offline_factory().await;
1570        let securable =
1571            SecurableRef::Path(Url::parse(url).unwrap(), PathOperation::Read, Some(false));
1572        // `build_store` does no network I/O — it only constructs the emulator
1573        // store and computes the prefix — so the dead-endpoint factory is fine.
1574        let store = factory.build_store(credential, securable).await.unwrap();
1575
1576        // The container-relative blob prefix, with `/<account>/<container>` stripped.
1577        assert_eq!(store.prefix(), &Path::from("tbl/data"));
1578    }
1579
1580    /// `now_epoch_millis` is defined in the credential-vending crate, not here;
1581    /// the store crate just needs a future timestamp for the test credential.
1582    fn now_epoch_millis() -> i64 {
1583        use std::time::{SystemTime, UNIX_EPOCH};
1584        SystemTime::now()
1585            .duration_since(UNIX_EPOCH)
1586            .unwrap()
1587            .as_millis() as i64
1588    }
1589
1590    // ---- forwarded-identity header forwarding ------------------------------
1591
1592    /// A JSON `temporary-path-credentials` response carrying a minimal, well-formed
1593    /// Azurite SAS credential — enough for `for_path` to build an emulator store
1594    /// with no real cloud I/O. `url` is a path-style Azurite endpoint.
1595    fn azurite_credential_body() -> String {
1596        format!(
1597            r#"{{"expiration_time":{},"url":"http://127.0.0.1:10000/devstoreaccount1/mycontainer/prefix/","azure_user_delegation_sas":{{"sas_token":"sv=2021-08-06&ss=b&srt=co&sp=rl&se=2999-01-01T00:00:00Z&sig=AAAA"}}}}"#,
1598            now_epoch_millis() + 3_600_000
1599        )
1600    }
1601
1602    /// A factory pointed at the given (mock) UC base URL, authenticated with a
1603    /// static token — the token the forwarded-identity path must *drop*.
1604    async fn factory_at(base_url: &str) -> UnityObjectStoreFactory {
1605        UnityObjectStoreFactory::builder()
1606            .with_uri(format!("{base_url}/api/2.1/unity-catalog/"))
1607            .with_token("proxy-service-token".to_string())
1608            .build()
1609            .await
1610            .unwrap()
1611    }
1612
1613    /// With a forwarded user, the upstream vend carries `x-forwarded-user` and
1614    /// drops the proxy's own bearer token (header-replaces-token).
1615    #[tokio::test]
1616    async fn forwarded_user_injects_header_and_drops_token() {
1617        let mut server = mockito::Server::new_async().await;
1618        let mock = server
1619            .mock("POST", "/api/2.1/unity-catalog/temporary-path-credentials")
1620            .match_header("x-forwarded-user", "alice")
1621            // The static bearer token must NOT ride along on a forwarded call.
1622            .match_header("authorization", mockito::Matcher::Missing)
1623            .with_status(200)
1624            .with_header("content-type", "application/json")
1625            .with_body(azurite_credential_body())
1626            .expect(1)
1627            .create_async()
1628            .await;
1629
1630        let factory = factory_at(&server.url()).await;
1631        let scoped = factory
1632            .with_forwarded_user("x-forwarded-user", Some("alice"))
1633            .unwrap();
1634        let url = Url::parse("abfss://mycontainer@devstoreaccount1/prefix/").unwrap();
1635        scoped
1636            .for_path(&url, PathOperation::Read)
1637            .await
1638            .expect("vend + store build should succeed against the mock");
1639
1640        mock.assert_async().await;
1641    }
1642
1643    /// An anonymous request (no forwarded user) sends NO `x-forwarded-user` header
1644    /// and still authenticates with the proxy's static bearer token — unchanged
1645    /// from today.
1646    #[tokio::test]
1647    async fn anonymous_sends_no_forwarded_header_and_keeps_token() {
1648        let mut server = mockito::Server::new_async().await;
1649        let mock = server
1650            .mock("POST", "/api/2.1/unity-catalog/temporary-path-credentials")
1651            .match_header("x-forwarded-user", mockito::Matcher::Missing)
1652            .match_header("authorization", "Bearer proxy-service-token")
1653            .with_status(200)
1654            .with_header("content-type", "application/json")
1655            .with_body(azurite_credential_body())
1656            .expect(1)
1657            .create_async()
1658            .await;
1659
1660        let factory = factory_at(&server.url()).await;
1661        // `None` returns the base factory unchanged.
1662        let scoped = factory
1663            .with_forwarded_user("x-forwarded-user", None)
1664            .unwrap();
1665        let url = Url::parse("abfss://mycontainer@devstoreaccount1/prefix/").unwrap();
1666        scoped
1667            .for_path(&url, PathOperation::Read)
1668            .await
1669            .expect("vend + store build should succeed against the mock");
1670
1671        mock.assert_async().await;
1672    }
1673
1674    /// The header name is configurable; a custom outgoing header carries the user.
1675    #[tokio::test]
1676    async fn forwarded_user_honors_custom_header_name() {
1677        let mut server = mockito::Server::new_async().await;
1678        let mock = server
1679            .mock("POST", "/api/2.1/unity-catalog/temporary-path-credentials")
1680            .match_header("x-remote-user", "bob")
1681            .with_status(200)
1682            .with_header("content-type", "application/json")
1683            .with_body(azurite_credential_body())
1684            .expect(1)
1685            .create_async()
1686            .await;
1687
1688        let factory = factory_at(&server.url()).await;
1689        let scoped = factory
1690            .with_forwarded_user("x-remote-user", Some("bob"))
1691            .unwrap();
1692        let url = Url::parse("abfss://mycontainer@devstoreaccount1/prefix/").unwrap();
1693        scoped.for_path(&url, PathOperation::Read).await.unwrap();
1694
1695        mock.assert_async().await;
1696    }
1697
1698    /// An invalid header name / value is a clean config error, not a panic.
1699    #[tokio::test]
1700    async fn forwarded_user_rejects_invalid_header() {
1701        let factory = factory_at("http://127.0.0.1:0").await;
1702        // Space is not a legal header-name character.
1703        assert!(
1704            factory
1705                .with_forwarded_user("bad header", Some("alice"))
1706                .is_err()
1707        );
1708        // A newline is not a legal header-value character.
1709        assert!(
1710            factory
1711                .with_forwarded_user("x-forwarded-user", Some("a\nb"))
1712                .is_err()
1713        );
1714    }
1715
1716    /// The `ForwardedHeaderService` decorator inserts the header (replacing any
1717    /// prior value of the same name) and delegates to its inner service.
1718    #[tokio::test]
1719    async fn forwarded_header_service_inserts_and_delegates() {
1720        use olai_http::service::{HttpService, ReqwestService};
1721        use reqwest::header::{HeaderName, HeaderValue};
1722
1723        let mut server = mockito::Server::new_async().await;
1724        let mock = server
1725            .mock("GET", "/echo")
1726            .match_header("x-forwarded-user", "carol")
1727            .with_status(204)
1728            .expect(1)
1729            .create_async()
1730            .await;
1731
1732        let client = reqwest::Client::new();
1733        let svc = ForwardedHeaderService {
1734            inner: Arc::new(ReqwestService::new(client.clone())),
1735            name: HeaderName::from_static("x-forwarded-user"),
1736            value: HeaderValue::from_static("carol"),
1737        };
1738        // A pre-existing value of the same header must be overwritten, not appended.
1739        let mut request = client
1740            .get(format!("{}/echo", server.url()))
1741            .build()
1742            .unwrap();
1743        request
1744            .headers_mut()
1745            .insert("x-forwarded-user", HeaderValue::from_static("stale"));
1746
1747        let resp = svc.call(request).await.unwrap();
1748        assert_eq!(resp.status(), 204);
1749        mock.assert_async().await;
1750    }
1751}