object_store/
registry.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Map object URLs to [`ObjectStore`]
19
20use crate::path::{InvalidPart, Path, PathPart};
21use crate::{ObjectStore, parse_url_opts};
22use parking_lot::RwLock;
23use std::collections::HashMap;
24use std::sync::Arc;
25use url::Url;
26
27/// [`ObjectStoreRegistry`] maps a URL to an [`ObjectStore`] instance
28pub trait ObjectStoreRegistry: Send + Sync + std::fmt::Debug + 'static {
29    /// Register a new store for the provided store URL
30    ///
31    /// If a store with the same URL existed before, it is replaced and returned
32    fn register(&self, url: Url, store: Arc<dyn ObjectStore>) -> Option<Arc<dyn ObjectStore>>;
33
34    /// Resolve an object URL
35    ///
36    /// If [`ObjectStoreRegistry::register`] has been called with a URL with the same
37    /// scheme, and authority as the object URL, and a path that is a prefix of the object
38    /// URL's, it should be returned along with the trailing path. Paths should be matched
39    /// on a path segment basis, and in the event of multiple possibilities the longest
40    /// path match should be returned.
41    ///
42    /// If a store hasn't been registered, an [`ObjectStoreRegistry`] may lazily create
43    /// one if the URL is understood
44    ///
45    /// For example
46    ///
47    /// ```
48    /// # use std::sync::Arc;
49    /// # use url::Url;
50    /// # use object_store::memory::InMemory;
51    /// # use object_store::ObjectStore;
52    /// # use object_store::prefix::PrefixStore;
53    /// # use object_store::registry::{DefaultObjectStoreRegistry, ObjectStoreRegistry};
54    /// #
55    /// let registry = DefaultObjectStoreRegistry::new();
56    ///
57    /// let bucket1 = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
58    /// let base = Url::parse("s3://bucket1/").unwrap();
59    /// registry.register(base, bucket1.clone());
60    ///
61    /// let url = Url::parse("s3://bucket1/path/to/object").unwrap();
62    /// let (ret, path) = registry.resolve(&url).unwrap();
63    /// assert_eq!(path.as_ref(), "path/to/object");
64    /// assert!(Arc::ptr_eq(&ret, &bucket1));
65    ///
66    /// let bucket2 = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
67    /// let base = Url::parse("https://s3.region.amazonaws.com/bucket").unwrap();
68    /// registry.register(base, bucket2.clone());
69    ///
70    /// let url = Url::parse("https://s3.region.amazonaws.com/bucket/path/to/object").unwrap();
71    /// let (ret, path) = registry.resolve(&url).unwrap();
72    /// assert_eq!(path.as_ref(), "path/to/object");
73    /// assert!(Arc::ptr_eq(&ret, &bucket2));
74    ///
75    /// let bucket3 = Arc::new(PrefixStore::new(InMemory::new(), "path")) as Arc<dyn ObjectStore>;
76    /// let base = Url::parse("https://s3.region.amazonaws.com/bucket/path").unwrap();
77    /// registry.register(base, bucket3.clone());
78    ///
79    /// let url = Url::parse("https://s3.region.amazonaws.com/bucket/path/to/object").unwrap();
80    /// let (ret, path) = registry.resolve(&url).unwrap();
81    /// assert_eq!(path.as_ref(), "to/object");
82    /// assert!(Arc::ptr_eq(&ret, &bucket3));
83    /// ```
84    fn resolve(&self, url: &Url) -> crate::Result<(Arc<dyn ObjectStore>, Path)>;
85
86    /// Deregister the store registered for exactly `url`, returning it if present.
87    ///
88    /// Unlike [`resolve`](Self::resolve), which matches the longest registered
89    /// path that is a prefix of the object URL, this removes the store
90    /// registered at the *exact* scheme, authority and path of `url` (the
91    /// inverse of [`register`](Self::register)). A store registered at a
92    /// different path under the same authority is left in place.
93    ///
94    /// The default implementation does nothing and returns `None`.
95    ///
96    /// ```
97    /// # use std::sync::Arc;
98    /// # use url::Url;
99    /// # use object_store::memory::InMemory;
100    /// # use object_store::ObjectStore;
101    /// # use object_store::registry::{DefaultObjectStoreRegistry, ObjectStoreRegistry};
102    /// #
103    /// let registry = DefaultObjectStoreRegistry::new();
104    /// let bucket = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
105    /// let url = Url::parse("s3://bucket/path").unwrap();
106    /// registry.register(url.clone(), bucket.clone());
107    ///
108    /// let removed = registry.deregister(&url).unwrap();
109    /// assert!(Arc::ptr_eq(&removed, &bucket));
110    /// ```
111    fn deregister(&self, url: &Url) -> Option<Arc<dyn ObjectStore>> {
112        let _ = url;
113        None
114    }
115}
116
117/// Error type for [`DefaultObjectStoreRegistry`]
118///
119/// Crate private/opaque type to make the error handling code more ergonomic.
120/// Always converted into `crate::Error` when reported externally.
121#[derive(Debug, thiserror::Error)]
122#[non_exhaustive]
123enum Error {
124    #[error("ObjectStore not found")]
125    NotFound,
126
127    #[error("Error parsing URL path segment")]
128    InvalidPart(#[from] InvalidPart),
129}
130
131impl From<Error> for crate::Error {
132    fn from(value: Error) -> Self {
133        Self::Generic {
134            store: "ObjectStoreRegistry",
135            source: Box::new(value),
136        }
137    }
138}
139
140/// An [`ObjectStoreRegistry`] that uses [`parse_url_opts`] to create stores based on the environment
141#[derive(Debug, Default)]
142pub struct DefaultObjectStoreRegistry {
143    /// Mapping from [`url_key`] to [`PathEntry`]
144    map: RwLock<HashMap<String, PathEntry>>,
145}
146
147/// [`PathEntry`] construct a tree of path segments starting from the root
148///
149/// For example the following paths
150///
151/// * `/` => store1
152/// * `/foo/bar` => store2
153///
154/// Would be represented by
155///
156/// ```yaml
157/// store: Some(store1)
158/// children:
159///   foo:
160///     store: None
161///     children:
162///       bar:
163///         store: Some(store2)
164/// ```
165///
166#[derive(Debug, Default)]
167struct PathEntry {
168    /// Store, if defined at this path
169    store: Option<Arc<dyn ObjectStore>>,
170    /// Child [`PathEntry`], keyed by the next path segment in their path
171    children: HashMap<String, Self>,
172}
173
174impl PathEntry {
175    /// Lookup a store based on URL path
176    ///
177    /// Returns the store and its path segment depth
178    fn lookup(&self, to_resolve: &Url) -> Option<(&Arc<dyn ObjectStore>, usize)> {
179        let mut current = self;
180        let mut ret = self.store.as_ref().map(|store| (store, 0));
181        let mut depth = 0;
182        // Traverse the PathEntry tree to find the longest match
183        for segment in path_segments(to_resolve.path()) {
184            match current.children.get(segment) {
185                Some(e) => {
186                    current = e;
187                    depth += 1;
188                    if let Some(store) = &current.store {
189                        ret = Some((store, depth))
190                    }
191                }
192                None => break,
193            }
194        }
195        ret
196    }
197
198    /// Remove the store registered at the exact `segments` path under this
199    /// entry, returning it. Child entries left empty are pruned.
200    fn remove(&mut self, segments: &[&str]) -> Option<Arc<dyn ObjectStore>> {
201        match segments.split_first() {
202            None => self.store.take(),
203            Some((head, rest)) => {
204                let child = self.children.get_mut(*head)?;
205                let removed = child.remove(rest);
206                if child.store.is_none() && child.children.is_empty() {
207                    self.children.remove(*head);
208                }
209                removed
210            }
211        }
212    }
213}
214
215impl DefaultObjectStoreRegistry {
216    /// Create a new [`DefaultObjectStoreRegistry`]
217    pub fn new() -> Self {
218        Self::default()
219    }
220}
221
222impl ObjectStoreRegistry for DefaultObjectStoreRegistry {
223    fn register(&self, url: Url, store: Arc<dyn ObjectStore>) -> Option<Arc<dyn ObjectStore>> {
224        let mut map = self.map.write();
225        let key = url_key(&url);
226        let mut entry = map.entry(key.to_string()).or_default();
227
228        for segment in path_segments(url.path()) {
229            entry = entry.children.entry(segment.to_string()).or_default();
230        }
231        entry.store.replace(store)
232    }
233
234    fn resolve(&self, to_resolve: &Url) -> crate::Result<(Arc<dyn ObjectStore>, Path)> {
235        let key = url_key(to_resolve);
236        {
237            let map = self.map.read();
238
239            if let Some((store, depth)) = map.get(key).and_then(|entry| entry.lookup(to_resolve)) {
240                let path = path_suffix(to_resolve, depth)?;
241                return Ok((Arc::clone(store), path));
242            }
243        }
244
245        if let Ok((store, path)) = parse_url_opts(to_resolve, std::env::vars()) {
246            let depth = num_segments(to_resolve.path()) - num_segments(path.as_ref());
247
248            let mut map = self.map.write();
249            let mut entry = map.entry(key.to_string()).or_default();
250            for segment in path_segments(to_resolve.path()).take(depth) {
251                entry = entry.children.entry(segment.to_string()).or_default();
252            }
253            let store = Arc::clone(match &entry.store {
254                None => entry.store.insert(Arc::from(store)),
255                Some(x) => x, // Racing creation - use existing
256            });
257
258            let path = path_suffix(to_resolve, depth)?;
259            return Ok((store, path));
260        }
261
262        Err(Error::NotFound.into())
263    }
264
265    fn deregister(&self, url: &Url) -> Option<Arc<dyn ObjectStore>> {
266        let key = url_key(url);
267        let segments: Vec<&str> = path_segments(url.path()).collect();
268        let mut map = self.map.write();
269        let entry = map.get_mut(key)?;
270        let removed = entry.remove(&segments);
271        // Drop the authority entry entirely once it holds no stores.
272        if entry.store.is_none() && entry.children.is_empty() {
273            map.remove(key);
274        }
275        removed
276    }
277}
278
279/// Extracts the scheme and authority of a URL (components before the Path)
280fn url_key(url: &Url) -> &str {
281    &url[..url::Position::AfterPort]
282}
283
284/// Returns the non-empty segments of a path
285///
286/// Note: We don't use [`Url::path_segments`] as we only want non-empty paths
287fn path_segments(s: &str) -> impl Iterator<Item = &str> {
288    s.split('/').filter(|x| !x.is_empty())
289}
290
291/// Returns the number of non-empty path segments in a path
292fn num_segments(s: &str) -> usize {
293    path_segments(s).count()
294}
295
296/// Returns the path of `url` skipping the first `depth` segments
297fn path_suffix(url: &Url, depth: usize) -> Result<Path, Error> {
298    let segments = path_segments(url.path()).skip(depth);
299    let path = segments.map(PathPart::parse).collect::<Result<_, _>>()?;
300    Ok(path)
301}
302
303#[cfg(test)]
304mod tests {
305    use super::*;
306    use crate::memory::InMemory;
307    use crate::prefix::PrefixStore;
308
309    #[test]
310    fn test_num_segments() {
311        assert_eq!(num_segments(""), 0);
312        assert_eq!(num_segments("/"), 0);
313        assert_eq!(num_segments("/banana"), 1);
314        assert_eq!(num_segments("banana"), 1);
315        assert_eq!(num_segments("/banana/crumble"), 2);
316        assert_eq!(num_segments("banana/crumble"), 2);
317    }
318
319    #[test]
320    fn test_default_registry() {
321        let registry = DefaultObjectStoreRegistry::new();
322
323        // Should automatically register in memory store
324        let banana_url = Url::parse("memory:///banana").unwrap();
325        let (resolved, path) = registry.resolve(&banana_url).unwrap();
326        assert_eq!(path.as_ref(), "banana");
327
328        // Should replace store
329        let url = Url::parse("memory:///").unwrap();
330        let root = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
331        let replaced = registry.register(url, Arc::clone(&root)).unwrap();
332        assert!(Arc::ptr_eq(&resolved, &replaced));
333
334        // Should not replace store
335        let banana = Arc::new(PrefixStore::new(InMemory::new(), "banana")) as Arc<dyn ObjectStore>;
336        assert!(
337            registry
338                .register(banana_url.clone(), Arc::clone(&banana))
339                .is_none()
340        );
341
342        // Should resolve to banana store
343        let (resolved, path) = registry.resolve(&banana_url).unwrap();
344        assert_eq!(path.as_ref(), "");
345        assert!(Arc::ptr_eq(&resolved, &banana));
346
347        // If we register another store it still resolves banana
348        let apples_url = Url::parse("memory:///apples").unwrap();
349        let apples = Arc::new(PrefixStore::new(InMemory::new(), "apples")) as Arc<dyn ObjectStore>;
350        assert!(registry.register(apples_url, Arc::clone(&apples)).is_none());
351
352        // Should still resolve to banana store
353        let (resolved, path) = registry.resolve(&banana_url).unwrap();
354        assert_eq!(path.as_ref(), "");
355        assert!(Arc::ptr_eq(&resolved, &banana));
356
357        // Should be path segment based
358        let banana_muffins_url = Url::parse("memory:///banana_muffins").unwrap();
359        let (resolved, path) = registry.resolve(&banana_muffins_url).unwrap();
360        assert_eq!(path.as_ref(), "banana_muffins");
361        assert!(Arc::ptr_eq(&resolved, &root));
362
363        // Should resolve to root even though path contains prefix of valid store
364        let to_resolve = Url::parse("memory:///foo/banana").unwrap();
365        let (resolved, path) = registry.resolve(&to_resolve).unwrap();
366        assert_eq!(path.as_ref(), "foo/banana");
367        assert!(Arc::ptr_eq(&resolved, &root));
368
369        let nested_url = Url::parse("memory:///apples/bananas").unwrap();
370        let nested =
371            Arc::new(PrefixStore::new(InMemory::new(), "apples/bananas")) as Arc<dyn ObjectStore>;
372        assert!(registry.register(nested_url, Arc::clone(&nested)).is_none());
373
374        let to_resolve = Url::parse("memory:///apples/bananas/muffins/cupcakes").unwrap();
375        let (resolved, path) = registry.resolve(&to_resolve).unwrap();
376        assert_eq!(path.as_ref(), "muffins/cupcakes");
377        assert!(Arc::ptr_eq(&resolved, &nested));
378
379        let nested_url2 = Url::parse("memory:///1/2/3").unwrap();
380        let nested2 = Arc::new(PrefixStore::new(InMemory::new(), "1/2/3")) as Arc<dyn ObjectStore>;
381        assert!(
382            registry
383                .register(nested_url2, Arc::clone(&nested2))
384                .is_none()
385        );
386
387        let to_resolve = Url::parse("memory:///1/2/3/4/5/6").unwrap();
388        let (resolved, path) = registry.resolve(&to_resolve).unwrap();
389        assert_eq!(path.as_ref(), "4/5/6");
390        assert!(Arc::ptr_eq(&resolved, &nested2));
391
392        let custom_scheme_url = Url::parse("custom:///").unwrap();
393        let custom_scheme = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
394        assert!(
395            registry
396                .register(custom_scheme_url, Arc::clone(&custom_scheme))
397                .is_none()
398        );
399
400        let to_resolve = Url::parse("custom:///6/7").unwrap();
401        let (resolved, path) = registry.resolve(&to_resolve).unwrap();
402        assert_eq!(path.as_ref(), "6/7");
403        assert!(Arc::ptr_eq(&resolved, &custom_scheme));
404    }
405
406    #[test]
407    fn test_deregister() {
408        let registry = DefaultObjectStoreRegistry::new();
409        let a = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
410        let b = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
411
412        // Use a scheme `parse_url_opts` cannot lazily create, so a failed
413        // `resolve` reflects deregistration rather than the environment.
414        let a_url = Url::parse("mock://bucket/a").unwrap();
415        let b_url = Url::parse("mock://bucket/b").unwrap();
416        registry.register(a_url.clone(), Arc::clone(&a));
417        registry.register(b_url.clone(), Arc::clone(&b));
418
419        // Removes exactly the store at `a`, leaving its sibling `b` in place.
420        let removed = registry.deregister(&a_url).unwrap();
421        assert!(Arc::ptr_eq(&removed, &a));
422        assert!(registry.resolve(&a_url).is_err());
423        let (resolved, _) = registry.resolve(&b_url).unwrap();
424        assert!(Arc::ptr_eq(&resolved, &b));
425
426        // Deregistering again returns None.
427        assert!(registry.deregister(&a_url).is_none());
428        // A never-registered URL returns None.
429        assert!(
430            registry
431                .deregister(&Url::parse("mock://other/x").unwrap())
432                .is_none()
433        );
434    }
435
436    #[test]
437    fn test_deregister_keeps_longer_prefix() {
438        let registry = DefaultObjectStoreRegistry::new();
439        let outer = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
440        let inner = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
441        let outer_url = Url::parse("mock://bucket/a").unwrap();
442        let inner_url = Url::parse("mock://bucket/a/b").unwrap();
443        registry.register(outer_url.clone(), Arc::clone(&outer));
444        registry.register(inner_url.clone(), Arc::clone(&inner));
445
446        // Removing the outer mount leaves the nested one resolvable.
447        assert!(Arc::ptr_eq(
448            &registry.deregister(&outer_url).unwrap(),
449            &outer
450        ));
451        let (resolved, path) = registry
452            .resolve(&Url::parse("mock://bucket/a/b/x").unwrap())
453            .unwrap();
454        assert!(Arc::ptr_eq(&resolved, &inner));
455        assert_eq!(path.as_ref(), "x");
456    }
457}
object_store/registry.rs

object_store/
registry.rs
