tustvold commented on code in PR #375: URL: https://github.com/apache/arrow-rs-object-store/pull/375#discussion_r2094617785
########## src/registry.rs: ########## @@ -0,0 +1,328 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Map object URLs to [`ObjectStore`] + +use crate::path::{InvalidPart, Path, PathPart}; +use crate::{parse_url_opts, ObjectStore}; +use parking_lot::RwLock; +use std::collections::HashMap; +use std::sync::Arc; +use url::Url; + +/// [`ObjectStoreRegistry`] maps a URL to an [`ObjectStore`] instance +pub trait ObjectStoreRegistry: Send + Sync + std::fmt::Debug + 'static { + /// Register a new store for the provided store URL + /// + /// If a store with the same URL existed before, it is replaced and returned + fn register(&self, url: Url, store: Arc<dyn ObjectStore>) -> Option<Arc<dyn ObjectStore>>; + + /// Resolve an object URL + /// + /// If [`ObjectStoreRegistry::register`] has been called with a URL with the same + /// scheme, and authority as the object URL, and a path that is a prefix of the object + /// URL's, it should be returned along with the trailing path. Paths should be matched + /// on a path segment basis, and in the event of multiple possibilities the longest + /// path match should be returned. + /// + /// If a store hasn't been registered, an [`ObjectStoreRegistry`] may lazily create + /// one if the URL is understood + /// + /// For example + /// + /// ``` + /// # use std::sync::Arc; + /// # use url::Url; + /// # use object_store::memory::InMemory; + /// # use object_store::ObjectStore; + /// # use object_store::prefix::PrefixStore; + /// # use object_store::registry::{DefaultObjectStoreRegistry, ObjectStoreRegistry}; + /// # + /// let registry = DefaultObjectStoreRegistry::new(); + /// + /// let bucket1 = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>; + /// let base = Url::parse("s3://bucket1/").unwrap(); + /// registry.register(base, bucket1.clone()); + /// + /// let url = Url::parse("s3://bucket1/path/to/object").unwrap(); + /// let (ret, path) = registry.resolve(&url).unwrap(); + /// assert_eq!(path.as_ref(), "path/to/object"); + /// assert!(Arc::ptr_eq(&ret, &bucket1)); + /// + /// let bucket2 = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>; + /// let base = Url::parse("https://s3.region.amazonaws.com/bucket";).unwrap(); + /// registry.register(base, bucket2.clone()); + /// + /// let url = Url::parse("https://s3.region.amazonaws.com/bucket/path/to/object";).unwrap(); + /// let (ret, path) = registry.resolve(&url).unwrap(); + /// assert_eq!(path.as_ref(), "path/to/object"); + /// assert!(Arc::ptr_eq(&ret, &bucket2)); + /// + /// let bucket3 = Arc::new(PrefixStore::new(InMemory::new(), "path")) as Arc<dyn ObjectStore>; + /// let base = Url::parse("https://s3.region.amazonaws.com/bucket/path";).unwrap(); + /// registry.register(base, bucket3.clone()); + /// + /// let url = Url::parse("https://s3.region.amazonaws.com/bucket/path/to/object";).unwrap(); + /// let (ret, path) = registry.resolve(&url).unwrap(); + /// assert_eq!(path.as_ref(), "to/object"); + /// assert!(Arc::ptr_eq(&ret, &bucket3)); + /// ``` + fn resolve(&self, url: &Url) -> crate::Result<(Arc<dyn ObjectStore>, Path)>; +} + +/// Error type for [`DefaultObjectStoreRegistry`] +#[derive(Debug, thiserror::Error)] +#[non_exhaustive] +enum Error { + #[error("ObjectStore not found")] + NotFound, + + #[error("Error parsing URL path segment")] + InvalidPart(#[from] InvalidPart), +} + +impl From<Error> for crate::Error { + fn from(value: Error) -> Self { + Self::Generic { + store: "ObjectStoreRegistry", + source: Box::new(value), + } + } +} + +/// An [`ObjectStoreRegistry`] that uses [`parse_url_opts`] to create stores based on the environment +#[derive(Debug, Default)] +pub struct DefaultObjectStoreRegistry { + /// Mapping from [`url_key`] to [`PathEntry`] + map: RwLock<HashMap<String, PathEntry>>, +} + +/// [`PathEntry`] construct a tree of path segments starting from the root +/// +/// For example the following paths +/// +/// * `/` => store1 +/// * `/foo/bar` => store2 +/// +/// Would be represented by +/// +/// ```yaml +/// store: Some(store1) +/// children: +/// foo: +/// store: None +/// children: +/// bar: +/// store: Some(store2) +/// ``` +/// +#[derive(Debug, Default)] +struct PathEntry { + /// Store, if defined at this path + store: Option<Arc<dyn ObjectStore>>, + /// Child [`PathEntry`], keyed by the next path segment in their path + children: HashMap<String, Self>, +} + +impl PathEntry { + /// Lookup a store based on URL path + /// + /// Returns the store and its path segment depth + fn lookup(&self, to_resolve: &Url) -> Option<(&Arc<dyn ObjectStore>, usize)> { + let mut current = self; + let mut ret = self.store.as_ref().map(|store| (store, 0)); + let mut depth = 0; + // Traverse the PathEntry tree to find the longest match + for segment in path_segments(to_resolve.path()) { + if let Some(e) = current.children.get(segment) { + current = e; + depth += 1; + if let Some(store) = ¤t.store { + ret = Some((store, depth)) + } + } + } + ret + } +} + +impl DefaultObjectStoreRegistry { + /// Create a new [`DefaultObjectStoreRegistry`] + pub fn new() -> Self { + Self::default() + } +} + +impl ObjectStoreRegistry for DefaultObjectStoreRegistry { + fn register(&self, url: Url, store: Arc<dyn ObjectStore>) -> Option<Arc<dyn ObjectStore>> { + let mut map = self.map.write(); + let key = url_key(&url); + let mut entry = map.entry(key.to_string()).or_default(); + + for segment in path_segments(url.path()) { + entry = entry.children.entry(segment.to_string()).or_default(); + } + entry.store.replace(store) + } + + fn resolve(&self, to_resolve: &Url) -> crate::Result<(Arc<dyn ObjectStore>, Path)> { + let key = url_key(to_resolve); + { + let map = self.map.read(); + + if let Some((store, depth)) = map.get(key).and_then(|entry| entry.lookup(to_resolve)) { + let path = path_suffix(to_resolve, depth)?; + return Ok((Arc::clone(store), path)); + } + } + + if let Ok((store, path)) = parse_url_opts(to_resolve, std::env::vars()) { + let depth = num_segments(to_resolve.path()) - num_segments(path.as_ref()); + + let mut map = self.map.write(); + let mut entry = map.entry(key.to_string()).or_default(); + for segment in path_segments(to_resolve.path()).take(depth) { + entry = entry.children.entry(segment.to_string()).or_default(); + } + let store = Arc::clone(match &entry.store { + None => entry.store.insert(Arc::from(store)), + Some(x) => x, // Racing creation - use existing + }); + + let path = path_suffix(to_resolve, depth)?; + return Ok((store, path)); + } + + Err(Error::NotFound.into()) + } +} + +/// Extracts the scheme and authority of a URL (components before the Path) +fn url_key(url: &Url) -> &str { + &url[..url::Position::AfterPort] Review Comment: It should be noted this will consider credentials unlike the version in DF - IMO this keeps things simple, and ultimately people shouldn't really be putting credentials in URLs -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: github-unsubscr...@arrow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
