unitycatalog_server/services/
location.rs

1use itertools::Itertools;
2use object_store::ObjectStoreScheme;
3use unitycatalog_common::{Error, Result};
4use url::Url;
5
6pub enum StorageLocationScheme {
7    ObjectStore(ObjectStoreScheme),
8    Azurite,
9}
10
11impl AsRef<str> for StorageLocationScheme {
12    fn as_ref(&self) -> &str {
13        match self {
14            Self::ObjectStore(ObjectStoreScheme::Local) => "file",
15            Self::ObjectStore(ObjectStoreScheme::Memory) => "memory",
16            Self::ObjectStore(ObjectStoreScheme::AmazonS3) => "s3",
17            Self::ObjectStore(ObjectStoreScheme::GoogleCloudStorage) => "gs",
18            Self::ObjectStore(ObjectStoreScheme::MicrosoftAzure) => "az",
19            Self::ObjectStore(ObjectStoreScheme::Http) => "http",
20            // NB: ObjectStoreScheme is non exhaustive, so we need to handle the unknown case.
21            Self::ObjectStore(_) => "unknown",
22
23            // Custom schemes
24            Self::Azurite => "azurite",
25        }
26    }
27}
28
29impl StorageLocationScheme {
30    pub fn parse(url: &Url) -> Result<Self> {
31        match ObjectStoreScheme::parse(url) {
32            Ok((ObjectStoreScheme::Http, _)) if is_azurite(url) => Ok(Self::Azurite),
33            Ok((scheme, _)) => Ok(Self::ObjectStore(scheme)),
34            Err(e) => {
35                if is_azurite(url) {
36                    Ok(Self::Azurite)
37                } else {
38                    Err(Error::invalid_argument(e.to_string()))
39                }
40            }
41        }
42    }
43}
44
45/// A URL representing a storage location.
46///
47/// This struct provides a cerntalized place to parse various URLs into more specific
48/// service references along with the semantic information that can be extracted from
49/// the URL.
50pub struct StorageLocationUrl {
51    /// The raw, unaltered URL.
52    url: Url,
53    /// The normalized `scheme://authority/` base of the location, used to derive
54    /// [`location`](Self::location). A plain [`Url`] — this is only ever built
55    /// from an authority string and read back as a URL/str.
56    store_url: Url,
57    scheme: StorageLocationScheme,
58    location: Url,
59}
60
61impl StorageLocationUrl {
62    pub fn try_new(url: Url) -> Result<Self> {
63        let (store_url, scheme, location) = get_store_url(&url)?;
64        Ok(Self {
65            url,
66            store_url,
67            scheme,
68            location,
69        })
70    }
71
72    pub fn parse(url: impl AsRef<str>) -> Result<Self> {
73        let url = Url::parse(url.as_ref())?;
74        Self::try_new(url)
75    }
76
77    pub fn raw(&self) -> &Url {
78        &self.url
79    }
80
81    pub fn location(&self) -> &Url {
82        &self.location
83    }
84
85    pub fn store_url(&self) -> &Url {
86        &self.store_url
87    }
88
89    pub fn scheme(&self) -> &StorageLocationScheme {
90        &self.scheme
91    }
92
93    /// Returns `(bucket_or_container, object_prefix)` extracted from the URL.
94    ///
95    /// - S3: `s3://bucket/prefix/path` → `("bucket", "prefix/path")`
96    /// - Azure HTTPS: `https://account.blob.core.windows.net/container/path` → `("container", "path")`
97    /// - Azurite HTTP: `http://localhost:10000/account/container/path` → `("container", "path")`
98    /// - Azurite custom scheme: `azurite://container/path` → `("container", "path")`
99    pub fn bucket_and_prefix(&self) -> Result<(String, String)> {
100        match &self.scheme {
101            StorageLocationScheme::ObjectStore(ObjectStoreScheme::AmazonS3) => {
102                let bucket = self
103                    .url
104                    .host_str()
105                    .ok_or_else(|| Error::invalid_argument("S3 URL missing bucket"))?
106                    .to_owned();
107                let prefix = self.url.path().trim_start_matches('/').to_owned();
108                Ok((bucket, prefix))
109            }
110            StorageLocationScheme::ObjectStore(ObjectStoreScheme::MicrosoftAzure) => {
111                // abfss://container@account.dfs.core.windows.net/path
112                // or https://account.blob.core.windows.net/container/path
113                let host = self.url.host_str().unwrap_or_default();
114                if host.contains('@') {
115                    // ABFSS: container@account.dfs.core.windows.net
116                    let container = host
117                        .split('@')
118                        .next()
119                        .ok_or_else(|| Error::invalid_argument("Invalid ABFSS URL"))?
120                        .to_owned();
121                    let prefix = self.url.path().trim_start_matches('/').to_owned();
122                    Ok((container, prefix))
123                } else {
124                    // https://account.blob.core.windows.net/container/path
125                    let mut segments = self
126                        .url
127                        .path_segments()
128                        .ok_or_else(|| Error::invalid_argument("Azure URL has no path"))?;
129                    let container = segments
130                        .next()
131                        .filter(|s| !s.is_empty())
132                        .ok_or_else(|| Error::invalid_argument("Azure URL missing container"))?
133                        .to_owned();
134                    let prefix = segments.collect::<Vec<_>>().join("/");
135                    Ok((container, prefix))
136                }
137            }
138            StorageLocationScheme::Azurite => {
139                if self.url.scheme() != "azurite" {
140                    // http://localhost:10000/account/container/path
141                    let parts: Vec<_> = self.url.path().splitn(4, '/').collect();
142                    // parts: ["", "account", "container", "path"]
143                    let container = parts
144                        .get(2)
145                        .filter(|s| !s.is_empty())
146                        .ok_or_else(|| Error::invalid_argument("Azurite URL missing container"))?
147                        .to_string();
148                    let prefix = parts.get(3).copied().unwrap_or("").to_owned();
149                    Ok((container, prefix))
150                } else {
151                    // azurite://container/path
152                    let container = self
153                        .url
154                        .host_str()
155                        .ok_or_else(|| {
156                            Error::invalid_argument("Azurite URL missing container in host")
157                        })?
158                        .to_owned();
159                    let prefix = self.url.path().trim_start_matches('/').to_owned();
160                    Ok((container, prefix))
161                }
162            }
163            _ => Err(Error::invalid_argument(
164                "bucket_and_prefix not supported for this URL scheme",
165            )),
166        }
167    }
168
169    /// Returns the Azure storage account name parsed from the URL, if present.
170    ///
171    /// - HTTPS: `https://account.blob.core.windows.net/…` → `Some("account")`
172    /// - ABFSS: `abfss://container@account.dfs.core.windows.net/…` → `Some("account")`
173    /// - Azurite HTTP: `http://localhost:10000/account/container/path` → `Some("account")`
174    /// - Azurite custom: `azurite://container/path` → `None` (no account concept)
175    pub fn azure_account(&self) -> Option<String> {
176        match &self.scheme {
177            StorageLocationScheme::ObjectStore(ObjectStoreScheme::MicrosoftAzure) => {
178                let host = self.url.host_str()?;
179                if host.contains('@') {
180                    // abfss://container@account.dfs.core.windows.net
181                    host.split('@')
182                        .nth(1)
183                        .map(|h| h.split('.').next().unwrap_or(h).to_owned())
184                } else {
185                    // https://account.blob.core.windows.net
186                    host.split('.').next().map(ToOwned::to_owned)
187                }
188            }
189            StorageLocationScheme::Azurite if self.url.scheme() != "azurite" => {
190                // http://localhost:10000/account/container/path
191                self.url
192                    .path_segments()
193                    .and_then(|mut s| s.next())
194                    .filter(|s| !s.is_empty())
195                    .map(ToOwned::to_owned)
196            }
197            _ => None,
198        }
199    }
200}
201
202fn is_azurite(url: &Url) -> bool {
203    // for now we assume that azurite is using default values.
204    // since this is only for local development, this may be indefinitely the case.
205    url.scheme() == "azurite"
206        || (url.scheme() == "http"
207            && matches!(url.host_str(), Some("localhost") | Some("127.0.0.1"))
208            && url.port() == Some(10000))
209}
210
211fn get_store_url(url: &url::Url) -> Result<(Url, StorageLocationScheme, Url)> {
212    let scheme = StorageLocationScheme::parse(url)?;
213    let store_url = match &scheme {
214        StorageLocationScheme::ObjectStore(_) => Url::parse(&format!(
215            "{}://{}",
216            scheme.as_ref(),
217            &url[url::Position::BeforeHost..url::Position::AfterPort]
218        ))
219        .map_err(|e| Error::Generic(e.to_string()))?,
220        StorageLocationScheme::Azurite => {
221            if url.scheme() != "azurite" {
222                // split the path into 3 parts:
223                // 1. account name
224                // 2. container name
225                // 3. path
226                let parts = url.path().splitn(4, "/").collect::<Vec<_>>();
227                if parts.len() != 4 {
228                    return Err(Error::invalid_argument(format!(
229                        "Invalid azurite path: {}",
230                        url.path()
231                    )));
232                }
233                Url::parse(&format!("azurite://{}", parts[2]))
234                    .map_err(|e| Error::Generic(e.to_string()))?
235            } else {
236                Url::parse(&format!(
237                    "{}://{}/",
238                    url.scheme(),
239                    &url[url::Position::BeforeHost..url::Position::AfterPort]
240                ))
241                .map_err(|e| Error::Generic(e.to_string()))?
242            }
243        }
244    };
245    let location = match &scheme {
246        StorageLocationScheme::ObjectStore(_) => store_url.join(url.path())?,
247        StorageLocationScheme::Azurite if url.scheme() != "azurite" => {
248            let path = url
249                .path_segments()
250                .ok_or_else(|| Error::invalid_argument("Invalid azurite path"))?
251                .skip(2)
252                .join("/");
253            url::Url::parse(&format!("{}{}", store_url.as_str(), path))?
254        }
255        StorageLocationScheme::Azurite => url.clone(),
256    };
257    Ok((store_url, scheme, location))
258}
unitycatalog_server/services/location.rs

unitycatalog_server/services/
location.rs
