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}