diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5fa6a1714a1..ac0275ea1e2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -246,7 +246,7 @@ jobs: target: wasm32-unknown-unknown env: rustflags: "RUSTFLAGS='-A warnings --cfg getrandom_backend=\"unsupported\"'" - args: "--target wasm32-unknown-unknown --exclude vortex --exclude vortex-cuda --exclude vortex-cub --exclude vortex-nvcomp --exclude vortex-datafusion --exclude vortex-duckdb --exclude vortex-tui --exclude vortex-zstd --exclude vortex-test-e2e-cuda --exclude vortex-python-cuda --exclude vortex-sqllogictest --exclude vortex-parquet-variant" + args: "--target wasm32-unknown-unknown --exclude vortex --exclude vortex-cuda --exclude vortex-cub --exclude vortex-nvcomp --exclude vortex-datafusion --exclude vortex-duckdb --exclude vortex-tui --exclude vortex-zstd --exclude vortex-test-e2e-cuda --exclude vortex-python-cuda --exclude vortex-sqllogictest --exclude vortex-parquet-variant --exclude vortex-object-store-opendal" steps: - uses: runs-on/action@v2 if: github.repository == 'vortex-data/vortex' diff --git a/Cargo.lock b/Cargo.lock index 16b0a83f9ac..0801580069c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2893,6 +2893,15 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "dlv-list" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "442039f5147480ba31067cb00ada1adae6892028e40e45fc5de7b7df6dcc1b5f" +dependencies = [ + "const-random", +] + [[package]] name = "document-features" version = "0.2.12" @@ -6014,6 +6023,7 @@ checksum = "96c9c85ce253ff87225e7669979d877a20c98a06604ec9d6dd5f4473e08f1ae1" dependencies = [ "opendal-core", "opendal-service-cos", + "opendal-service-oss", ] [[package]] @@ -6061,6 +6071,23 @@ dependencies = [ "serde", ] +[[package]] +name = "opendal-service-oss" +version = "0.57.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "328fa55e8888cbdfe00826bfea2a79042422b720e8369e9e021e46121dea5ace" +dependencies = [ + "bytes", + "http", + "log", + "opendal-core", + "quick-xml", + "reqsign-aliyun-oss", + "reqsign-core", + "reqsign-file-read-tokio", + "serde", +] + [[package]] name = "openssl-probe" version = "0.2.1" @@ -6161,6 +6188,16 @@ dependencies = [ "num-traits", ] +[[package]] +name = "ordered-multimap" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49203cdcae0030493bad186b28da2fa25645fa276a51b6fec8010d281e02ef79" +dependencies = [ + "dlv-list", + "hashbrown 0.14.5", +] + [[package]] name = "owo-colors" version = "4.3.0" @@ -7325,6 +7362,23 @@ dependencies = [ "bytecheck", ] +[[package]] +name = "reqsign-aliyun-oss" +version = "3.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f2a76ce7588780351c0f68eb8feaba8d23a99b40655195edbf50bfbce2af09b" +dependencies = [ + "anyhow", + "form_urlencoded", + "http", + "log", + "percent-encoding", + "reqsign-core", + "rust-ini", + "serde", + "serde_json", +] + [[package]] name = "reqsign-core" version = "3.2.0" @@ -7569,6 +7623,16 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "rust-ini" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "796e8d2b6696392a43bea58116b667fb4c29727dc5abd27d6acf338bb4f688c7" +dependencies = [ + "cfg-if", + "ordered-multimap", +] + [[package]] name = "rust-stemmers" version = "1.2.0" @@ -9944,10 +10008,12 @@ dependencies = [ "tempfile", "tokio", "tracing", + "url", "vortex-array", "vortex-buffer", "vortex-error", "vortex-metrics", + "vortex-object-store-opendal", "vortex-session", "vortex-utils", "wasm-bindgen-futures", @@ -10188,7 +10254,6 @@ dependencies = [ "vortex-object-store-opendal", "vortex-python-abi", "vortex-tui", - "vortex-utils", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index a9d8bfefb85..52d00d9c5e2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -95,7 +95,7 @@ rust-version = "1.91.0" version = "0.1.0" [workspace.dependencies] -anyhow = "1.0.97" +anyhow = "1.0.100" arbitrary = "1.3.2" arc-swap = "1.9" arcref = "0.2.0" @@ -173,7 +173,7 @@ indicatif = "0.18.0" insta = "1.43" inventory = "0.3.20" itertools = "0.14.0" -jiff = "0.2.0" +jiff = "0.2.17" jni = { version = "0.22.0" } kanal = "0.1.1" lasso = { version = "0.7", features = ["multi-threaded", "inline-more"] } @@ -191,9 +191,11 @@ noodles-vcf = { version = "0.88.0", features = ["async"] } num-traits = "0.2.19" num_enum = { version = "0.7.3", default-features = false } object_store = { version = "0.13.2", default-features = false } +object_store_opendal = "0.57.0" once_cell = "1.21" oneshot = { version = "0.2.0", features = ["async"] } onpair = "0.1.1" +opendal = { version = "0.57.0", default-features = false } opentelemetry = "0.32.0" opentelemetry-otlp = "0.32.0" opentelemetry_sdk = "0.32.0" @@ -232,7 +234,7 @@ rstest = "0.26.1" rstest_reuse = "0.7.0" rustc-hash = "2.1.1" rustix = { version = "1.1", features = ["fs"] } -serde = "1.0.220" +serde = "1.0.221" serde_json = "1.0.138" serde_test = "1.0.176" sha2 = "0.11.0" @@ -275,7 +277,7 @@ tracing-perfetto = "0.1.5" tracing-subscriber = "0.3" url = "2.5.7" uuid = { version = "1.23", features = ["js"] } -wasm-bindgen-futures = "0.4.54" +wasm-bindgen-futures = "0.4.58" wkb = "0.9.2" xshell = "0.2.6" zigzag = "0.1.0" @@ -310,6 +312,7 @@ vortex-json = { version = "0.1.0", path = "./vortex-json", default-features = fa vortex-layout = { version = "0.1.0", path = "./vortex-layout", default-features = false } vortex-mask = { version = "0.1.0", path = "./vortex-mask", default-features = false } vortex-metrics = { version = "0.1.0", path = "./vortex-metrics", default-features = false } +vortex-object-store-opendal = { version = "0.1.0", path = "./vortex-object-store-opendal", default-features = false } vortex-onpair = { version = "0.1.0", path = "./encodings/onpair", default-features = false } vortex-parquet-variant = { version = "0.1.0", path = "./encodings/parquet-variant" } vortex-pco = { version = "0.1.0", path = "./encodings/pco", default-features = false } diff --git a/docs/api/python/store/opendal.rst b/docs/api/python/store/opendal.rst index bac62f993be..a6bf038c809 100644 --- a/docs/api/python/store/opendal.rst +++ b/docs/api/python/store/opendal.rst @@ -1,13 +1,29 @@ -============= -OpenDAL (COS) -============= +================== +OpenDAL (COS, OSS) +================== -Vortex can read from and write to Tencent Cloud COS through +Vortex can read from and write to Tencent Cloud COS and Alibaba Cloud OSS through `OpenDAL `_, which provides native service support. -This store is available only when Vortex is built with the ``opendal`` feature +These stores are available only when Vortex is built with the ``opendal`` feature (e.g. ``maturin develop --features opendal`` or ``cargo build -p vortex-jni --features opendal``). +.. list-table:: + :header-rows: 1 + + * - Scheme + - Service + - Endpoint variable + - Credential variables + * - ``cos://`` + - Tencent Cloud COS + - ``COS_ENDPOINT`` + - ``TENCENTCLOUD_SECRET_ID``, ``TENCENTCLOUD_SECRET_KEY`` + * - ``oss://`` + - Alibaba Cloud OSS + - ``OSS_ENDPOINT`` + - ``ALIBABA_CLOUD_ACCESS_KEY_ID``, ``ALIBABA_CLOUD_ACCESS_KEY_SECRET`` + :class:`vortex.store.CosStore` ============================== @@ -85,3 +101,16 @@ exactly like the built-in S3/Azure/GCS stores: # When `store=` is supplied, the path is resolved as a key within the store, so the scheme # and bucket are not part of the path passed to read_url. a = read_url("path/to/dataset.vortex", store=store) + +Reading from OSS +================ + +Alibaba Cloud OSS is reachable by URL. There is no standalone ``OssStore`` class yet, so +configuration comes from the environment variables OpenDAL's OSS builder reads +(``ALIBABA_CLOUD_ACCESS_KEY_ID``, ``ALIBABA_CLOUD_ACCESS_KEY_SECRET`` and ``OSS_ENDPOINT``): + +.. code-block:: python + + import vortex as vx + + a = vx.io.read_url("oss://my-bucket/path/to/dataset.vortex") diff --git a/vortex-duckdb/Cargo.toml b/vortex-duckdb/Cargo.toml index 59dbe8d4003..734dbd5507e 100644 --- a/vortex-duckdb/Cargo.toml +++ b/vortex-duckdb/Cargo.toml @@ -38,10 +38,20 @@ static_assertions = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } url = { workspace = true } -vortex = { workspace = true, features = ["files", "tokio", "object_store"] } +vortex = { workspace = true, features = [ + "files", + "tokio", + "object_store", + "object_store_registry", +] } vortex-geo = { workspace = true } vortex-utils = { workspace = true, features = ["dashmap"] } +[features] +# OpenDAL-backed schemes (`cos://`, `oss://`) in the object store registry. Pulls in `opendal`, so +# it is opt-in. +opendal = ["vortex/opendal"] + [dev-dependencies] anyhow = { workspace = true } divan = { workspace = true } diff --git a/vortex-duckdb/src/multi_file.rs b/vortex-duckdb/src/multi_file.rs index 247b6a19909..756609ad2cc 100644 --- a/vortex-duckdb/src/multi_file.rs +++ b/vortex-duckdb/src/multi_file.rs @@ -2,11 +2,10 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::sync::Arc; +use std::sync::LazyLock; use itertools::Itertools; -use object_store::ObjectStore; -use object_store::aws::AmazonS3Builder; -use object_store::local::LocalFileSystem; +use object_store::registry::ObjectStoreRegistry; use url::Url; use vortex::error::VortexResult; use vortex::error::vortex_bail; @@ -16,6 +15,7 @@ use vortex::file::multi::parse_uri_or_path; use vortex::io::compat::Compat; use vortex::io::filesystem::FileSystemRef; use vortex::io::object_store::ObjectStoreFileSystem; +use vortex::io::object_store::Registry; use vortex::io::runtime::BlockingRuntime; use vortex::layout::scan::multi::MultiLayoutDataSource; use vortex_utils::aliases::hash_map::HashMap; @@ -25,18 +25,15 @@ use crate::SESSION; use crate::duckdb::BindInputRef; use crate::duckdb::ExtractedValue; +/// Process-wide registry, so repeated scans against the same bucket share one client. +static REGISTRY: LazyLock = LazyLock::new(Registry::new); + fn resolve_filesystem(base_url: &Url) -> VortexResult { - let object_store: Arc = match base_url.scheme() { - "file" => Arc::new(LocalFileSystem::new()), - "s3" => Arc::new( - AmazonS3Builder::from_env() - .with_bucket_name(base_url.host_str().ok_or_else(|| { - vortex_err!("Failed to extract bucket name from URL: {base_url}") - })?) - .build()?, - ), - other => vortex_bail!("Unsupported URL scheme '{other}'"), - }; + // `base_url` has its path cleared by the caller, so the resolved path is empty and only the + // store matters here. Going through the shared registry means DuckDB resolves the same set of + // schemes as the Python and Java bindings, including the OpenDAL-backed ones when the + // `opendal` feature is on. + let (object_store, _) = REGISTRY.resolve(base_url)?; Ok(Arc::new(ObjectStoreFileSystem::new( Arc::new(Compat::new(object_store)), diff --git a/vortex-io/Cargo.toml b/vortex-io/Cargo.toml index b3f6448484a..4c940561f38 100644 --- a/vortex-io/Cargo.toml +++ b/vortex-io/Cargo.toml @@ -34,10 +34,12 @@ tokio = { workspace = true, features = [ "sync", ] } # this is the maximum subset of fetaures that is safe for wasm32 targets tracing = { workspace = true } +url = { workspace = true, optional = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } vortex-metrics = { workspace = true } +vortex-object-store-opendal = { workspace = true, optional = true } vortex-session = { workspace = true } vortex-utils = { workspace = true } @@ -62,6 +64,23 @@ tokio = { workspace = true, features = ["full"] } [features] object_store = ["dep:object_store", "vortex-error/object_store"] +# The URL -> ObjectStore registry, plus the cloud backends it resolves URLs to. Kept separate from +# `object_store` so that plain file IO does not pull in the cloud HTTP stack. +object_store_registry = [ + "object_store", + "dep:url", + "object_store/aws", + "object_store/azure", + "object_store/gcp", + "object_store/http", +] +# OpenDAL-backed schemes (`cos://`, `oss://`) in the object store registry. +opendal = [ + "object_store_registry", + "dep:vortex-object-store-opendal", + "vortex-object-store-opendal/cos", + "vortex-object-store-opendal/oss", +] tokio = ["tokio/fs", "tokio/rt-multi-thread"] [lints] diff --git a/vortex-io/src/object_store/mod.rs b/vortex-io/src/object_store/mod.rs index 3345089a023..a9fb6476a42 100644 --- a/vortex-io/src/object_store/mod.rs +++ b/vortex-io/src/object_store/mod.rs @@ -3,8 +3,12 @@ mod filesystem; mod read_at; +#[cfg(feature = "object_store_registry")] +mod registry; mod write; pub use filesystem::*; pub use read_at::*; +#[cfg(feature = "object_store_registry")] +pub use registry::*; pub use write::*; diff --git a/vortex-io/src/object_store/registry/mod.rs b/vortex-io/src/object_store/registry/mod.rs new file mode 100644 index 00000000000..3f96331fadc --- /dev/null +++ b/vortex-io/src/object_store/registry/mod.rs @@ -0,0 +1,287 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Apache Software Foundation (ASF) + +//! A URL-to-[`ObjectStore`] registry. +//! +//! This is an adapted version of the `DefaultObjectStoreRegistry` from the `object_store` crate, +//! modified in two ways: +//! +//! 1. configuration is resolved out of environment variables case-insensitively, matching how the +//! various `Store::from_env` builders behave (see +//! ); +//! 2. schemes that `object_store` does not recognize natively — the OpenDAL-backed `cos://` and +//! `oss://` — are served under the `opendal` feature. +//! +//! Every Vortex language binding resolves URLs through this one type, so a scheme added here is +//! reachable from Python, Java and DuckDB alike. + +use std::sync::Arc; + +use object_store::ObjectStore; +use object_store::parse_url_opts; +use object_store::path::Path; +use object_store::path::PathPart; +use object_store::registry::ObjectStoreRegistry; +use parking_lot::RwLock; +use url::Url; +use vortex_utils::aliases::hash_map::HashMap; + +#[derive(Debug, Default)] +struct PathEntry { + /// Store, if defined at this path + store: Option>, + /// Child [`PathEntry`], keyed by the next path segment in their path + children: HashMap, +} + +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, 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()) { + match current.children.get(segment) { + Some(e) => { + current = e; + depth += 1; + if let Some(store) = ¤t.store { + ret = Some((store, depth)) + } + } + None => break, + } + } + ret + } +} + +/// An [`ObjectStoreRegistry`] that normalizes environment variables before doing lookups, and +/// resolves OpenDAL-backed schemes under the `opendal` feature. +/// +/// Stores are cached per URL authority, so repeated resolves against the same bucket share one +/// client. A store can also be registered explicitly at a path prefix with +/// [`ObjectStoreRegistry::register`], in which case it wins over the built-and-cached fallback. +/// +/// ``` +/// # use object_store::registry::ObjectStoreRegistry; +/// # use vortex_io::object_store::Registry; +/// # use url::Url; +/// # fn main() -> Result<(), Box> { +/// let registry = Registry::default(); +/// let store = object_store::memory::InMemory::new(); +/// registry.register(Url::parse("memory://bucket/")?, std::sync::Arc::new(store)); +/// +/// let (_store, path) = registry.resolve(&Url::parse("memory://bucket/a/b.vortex")?)?; +/// assert_eq!(path.as_ref(), "a/b.vortex"); +/// # Ok(()) +/// # } +/// ``` +#[derive(Debug, Default)] +pub struct Registry { + /// Mapping from [`url_key`] to [`PathEntry`] + map: RwLock>, + /// Where store configuration is read from when building a store. + env: EnvSource, +} + +/// Source of the configuration variables consulted when building a store. +/// +/// Tests construct a registry over a fixed set of variables rather than mutating the process +/// environment, which is unsound when tests run on multiple threads within one process (the +/// `std::env::set_var` block became `unsafe` in Rust 2024 for exactly this reason). +#[derive(Debug, Default)] +enum EnvSource { + /// Read from the process environment. + #[default] + Process, + /// A fixed set of variables. + #[cfg(test)] + Fixed(Vec<(String, String)>), +} + +impl EnvSource { + /// Case-insensitive lookup of a single configuration variable. + #[cfg(feature = "opendal")] + fn lookup(&self, key: &str) -> Option { + match self { + EnvSource::Process => std::env::var(key).ok(), + #[cfg(test)] + EnvSource::Fixed(vars) => vars + .iter() + .find(|(k, _)| k.eq_ignore_ascii_case(key)) + .map(|(_, v)| v.clone()), + } + } + + /// The configuration variables, with keys lowercased so that lookups are case-insensitive. + fn normalized_vars(&self) -> Vec<(String, String)> { + match self { + EnvSource::Process => std::env::vars() + .map(|(k, v)| (k.to_ascii_lowercase(), v)) + .collect(), + #[cfg(test)] + EnvSource::Fixed(vars) => vars + .iter() + .map(|(k, v)| (k.to_ascii_lowercase(), v.clone())) + .collect(), + } + } +} + +impl ObjectStoreRegistry for Registry { + fn register(&self, url: Url, store: Arc) -> Option> { + let mut map = self.map.write(); + let entry = entry_at( + &mut map, + url_key(&url), + url.path(), + num_segments(url.path()), + ); + entry.store.replace(store) + } + + fn resolve(&self, to_resolve: &Url) -> object_store::Result<(Arc, Path)> { + let key = url_key(to_resolve); + + // 1. Look up the user-registered map first. Every other scheme does this, so an explicit + // `Registry::register("cos://...", store)` should also win over the build-and-cache + // fallback below. This also means a previously-resolved store (built by us for the + // same URL) is served from the cache here, without rebuilding the client. + { + 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)); + } + } + + // 2. Build the store and the path *within* that store. `build_store` reports the path + // the way `parse_url_opts` does — i.e. the path is what the store will see as the + // key — which is what keeps every scheme on the one caching rule in + // `cache_and_resolve`. The depth is then just the difference in segment counts. + let (store, path) = self.build_store(to_resolve)?; + let depth = num_segments(to_resolve.path()) - num_segments(path.as_ref()); + + // 3. Cache the store and return it alongside the path that lives inside the store. + self.cache_and_resolve(to_resolve, store, depth) + } +} + +impl Registry { + /// Create a registry that reads store configuration from the process environment. + pub fn new() -> Self { + Self::default() + } + + /// Create a registry over a fixed set of configuration variables. + #[cfg(test)] + fn with_env(vars: I) -> Self + where + I: IntoIterator, + { + Self { + map: RwLock::default(), + env: EnvSource::Fixed(vars.into_iter().collect()), + } + } + + /// Caches `store` as the store serving the first `depth` path segments of `to_resolve`, and + /// returns it alongside the remaining path. + /// + /// If a racing `resolve` cached a store at the same position first, that store wins and the + /// one just built is dropped, so all callers of a given prefix share a single client. + fn cache_and_resolve( + &self, + to_resolve: &Url, + store: Arc, + depth: usize, + ) -> object_store::Result<(Arc, Path)> { + let path = path_suffix(to_resolve, depth)?; + + let mut map = self.map.write(); + let entry = entry_at(&mut map, url_key(to_resolve), to_resolve.path(), depth); + let stored = Arc::clone(match &entry.store { + None => entry.store.insert(store), + Some(existing) => existing, // Racing creation - use existing + }); + + Ok((stored, path)) + } + + /// Builds the [`ObjectStore`] that serves `to_resolve`, with the path within that store. + /// + /// This mirrors [`parse_url_opts`], extended with schemes the `object_store` crate does not + /// recognize. Reporting the path the way `parse_url_opts` does is what keeps every scheme on + /// the one caching rule in [`Registry::resolve`]: a scheme says where its store is rooted, and + /// the registry decides how to cache it. + fn build_store(&self, to_resolve: &Url) -> object_store::Result<(Arc, Path)> { + // OpenDAL-backed schemes (Tencent COS, Alibaba OSS) are not recognized by `object_store`, + // so build them from OpenDAL's own environment-variable configuration (e.g. + // `TENCENTCLOUD_SECRET_ID`). The operator is rooted at the bucket, which lives in the URL + // authority, so — exactly as for `s3://bucket/path` — the whole URL path is the object key. + #[cfg(feature = "opendal")] + if vortex_object_store_opendal::supports_scheme(to_resolve.scheme()) { + let store = vortex_object_store_opendal::make_opendal_store_with_env( + to_resolve, + &HashMap::new(), + |key| self.env.lookup(key), + )?; + return Ok((store, Path::from_url_path(to_resolve.path())?)); + } + + let (store, path) = parse_url_opts(to_resolve, self.env.normalized_vars())?; + Ok((Arc::from(store), path)) + } +} + +/// Extracts the scheme and authority of a URL (components before the Path) +fn url_key(url: &Url) -> &str { + &url[..url::Position::AfterPort] +} + +/// Returns the non-empty segments of a path +/// +/// Note: We don't use [`Url::path_segments`] as we only want non-empty paths +fn path_segments(s: &str) -> impl Iterator { + s.split('/').filter(|x| !x.is_empty()) +} + +/// Returns the number of non-empty path segments in a path +fn num_segments(s: &str) -> usize { + path_segments(s).count() +} + +/// Returns the path of `url` skipping the first `depth` segments +fn path_suffix(url: &Url, depth: usize) -> Result { + let segments = path_segments(url.path()).skip(depth); + let path = segments + .map(PathPart::parse) + .collect::>() + .map_err(|e| object_store::Error::Generic { + store: "ObjectStoreRegistry", + source: Box::new(e), + })?; + Ok(path) +} + +/// Walks to the [`PathEntry`] for `key` sitting `depth` segments into `path`, creating the +/// intermediate entries as needed. +fn entry_at<'a>( + map: &'a mut HashMap, + key: &str, + path: &str, + depth: usize, +) -> &'a mut PathEntry { + let mut current = map.entry(key.to_string()).or_default(); + for segment in path_segments(path).take(depth) { + current = current.children.entry(segment.to_string()).or_default(); + } + current +} + +#[cfg(test)] +mod tests; diff --git a/vortex-io/src/object_store/registry/tests.rs b/vortex-io/src/object_store/registry/tests.rs new file mode 100644 index 00000000000..ae9d21ba464 --- /dev/null +++ b/vortex-io/src/object_store/registry/tests.rs @@ -0,0 +1,146 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Write; +use std::sync::Arc; + +use object_store::ObjectStore; +use object_store::registry::ObjectStoreRegistry; +use url::Url; + +use super::Registry; + +/// A registry whose S3 configuration comes from a fixed map rather than the process environment, +/// so these tests neither read nor mutate global state. +fn registry() -> Registry { + Registry::with_env([("AWS_REGION".to_string(), "us-east-3".to_string())]) +} + +#[test] +#[expect(clippy::use_debug)] +fn test_resolve_url() -> Result<(), Box> { + let (store, _) = registry().resolve(&Url::parse("s3://my-bucket/test")?)?; + + // NOTE(aduffy): object_store doesn't let us downcast stores, the only way to verify + // that a configuration was added was to validate that it ends up in the Debug + // output :/ + let mut debug_str = String::new(); + write!(&mut debug_str, "{store:?}")?; + + assert!(debug_str.contains("us-east-3")); + Ok(()) +} + +/// Two resolves of the same URL must return the same `Arc` (the cached client) and the same +/// path, and two different keys must get distinct stores. This pins the symmetry the registry +/// promises: cache hit returns the same client and the same key. +#[test] +fn test_resolve_url_caches_per_key() -> Result<(), Box> { + let registry = registry(); + let first = Url::parse("s3://my-bucket/first/second")?; + let second = Url::parse("s3://my-bucket/first/second")?; + + let (store_a, path_a) = registry.resolve(&first)?; + let (store_b, path_b) = registry.resolve(&second)?; + assert!(Arc::ptr_eq(&store_a, &store_b)); + assert_eq!(path_a, path_b); + + // A different bucket gets its own client. + let other = Url::parse("s3://other-bucket/first/second")?; + let (store_c, _) = registry.resolve(&other)?; + assert!(!Arc::ptr_eq(&store_a, &store_c)); + Ok(()) +} + +/// Two objects in the same bucket must share a cached client. This is the regression that +/// the bespoke "walk every segment then return the whole key" path caused: by walking all +/// segments it stored the store at full depth, so the next resolve saw the whole key +/// already consumed and returned an empty path. Pinning the path here guards against that. +#[test] +fn test_resolve_url_shared_client_same_bucket() -> Result<(), Box> { + let registry = registry(); + let a = Url::parse("s3://my-bucket/path/to/data.vortex")?; + let b = Url::parse("s3://my-bucket/other/data.vortex")?; + + let (store_a, path_a) = registry.resolve(&a)?; + let (store_b, path_b) = registry.resolve(&b)?; + assert!(Arc::ptr_eq(&store_a, &store_b)); + assert_eq!(path_a.as_ref(), "path/to/data.vortex"); + assert_eq!(path_b.as_ref(), "other/data.vortex"); + Ok(()) +} + +/// Pinning the `entry_at` helper at depth > 0 (the path-prefix case) protects the shared +/// `entry_at` / `cache_and_resolve` helpers from regressing prefix registration. +#[test] +fn test_register_at_prefix_shares_store() -> Result<(), Box> { + let registry = registry(); + let prefix = Url::parse("s3://my-bucket/prefix/")?; + let (inner, _) = registry.resolve(&prefix)?; + let cached: Arc = Arc::clone(&inner); + let replaced = registry.register(prefix.clone(), Arc::clone(&cached)); + // First registration at this prefix replaces nothing. + assert!(replaced.is_none()); + + // A resolve at a deeper key still walks through the registered prefix's store. + let deeper = Url::parse("s3://my-bucket/prefix/inner/key")?; + let (store_deeper, path_deeper) = registry.resolve(&deeper)?; + assert!(Arc::ptr_eq(&store_deeper, &cached)); + assert_eq!(path_deeper.as_ref(), "inner/key"); + Ok(()) +} + +/// A `file://` URL with its path cleared must resolve to a local store. `vortex-duckdb` resolves +/// exactly this shape — it clears the path off the glob's base URL before asking for a +/// filesystem — so this pins the most common local-file case. +#[test] +fn test_resolve_file_url_with_empty_path() -> Result<(), Box> { + let registry = registry(); + + let mut base = Url::parse("file:///data/part-0.vortex")?; + base.set_path(""); + let (store, path) = registry.resolve(&base)?; + assert!(format!("{store:?}").contains("LocalFileSystem")); + assert_eq!(path.as_ref(), ""); + + // The same store is reused for a second glob rooted at the same authority. + let (store_again, _) = registry.resolve(&base)?; + assert!(Arc::ptr_eq(&store, &store_again)); + Ok(()) +} + +/// An explicitly registered store must win over the build-and-cache fallback, including for +/// schemes `object_store` does not recognize natively. `memory://` stands in for any such scheme +/// so this holds regardless of which service features are enabled. +#[test] +fn test_registered_store_wins_over_build() -> Result<(), Box> { + let registry = registry(); + let registered: Arc = Arc::new(object_store::memory::InMemory::new()); + registry.register(Url::parse("memory://bucket/")?, Arc::clone(®istered)); + + let (store, path) = registry.resolve(&Url::parse("memory://bucket/a/b.vortex")?)?; + assert!(Arc::ptr_eq(&store, ®istered)); + assert_eq!(path.as_ref(), "a/b.vortex"); + Ok(()) +} + +/// The OpenDAL-backed schemes must resolve through the registry rather than falling through to +/// `parse_url_opts`, which does not recognize them. Without an endpoint the build fails, but the +/// error must come from the OpenDAL builder ("missing required OpenDAL store configuration"), not +/// from `object_store`'s unrecognized-scheme path. +#[cfg(feature = "opendal")] +#[test] +fn test_opendal_schemes_reach_the_opendal_builder() -> Result<(), Box> { + for scheme in vortex_object_store_opendal::SUPPORTED_SCHEMES { + let url = Url::parse(&format!("{scheme}://bucket/key.vortex"))?; + let err = registry() + .resolve(&url) + .expect_err("no endpoint is configured, so the build must fail"); + let message = err.to_string(); + assert!( + message.contains("OpenDAL"), + "{scheme} did not reach the OpenDAL builder: {message}" + ); + } + Ok(()) +} diff --git a/vortex-jni/Cargo.toml b/vortex-jni/Cargo.toml index 80f46a67b82..b961352a228 100644 --- a/vortex-jni/Cargo.toml +++ b/vortex-jni/Cargo.toml @@ -34,15 +34,18 @@ url = { workspace = true } vortex = { workspace = true, features = ["object_store", "files"] } vortex-arrow = { workspace = true } vortex-geo = { workspace = true } -vortex-object-store-opendal = { path = "../vortex-object-store-opendal", optional = true } +vortex-object-store-opendal = { workspace = true, optional = true, features = [ + "cos", + "oss", +] } vortex-parquet-variant = { workspace = true } [dev-dependencies] jni = { workspace = true, features = ["invocation"] } [features] -# Enable OpenDAL-backed object stores (Tencent COS) for `cos://` URLs. This pulls in the -# `opendal` dependency, so it is opt-in. +# Enable OpenDAL-backed object stores (Tencent COS, Alibaba OSS) for `cos://` and `oss://` URLs. +# This pulls in the `opendal` dependency, so it is opt-in. opendal = ["dep:vortex-object-store-opendal"] [lib] diff --git a/vortex-jni/src/object_store.rs b/vortex-jni/src/object_store.rs index 8e4938d661a..a2c020cf3c8 100644 --- a/vortex-jni/src/object_store.rs +++ b/vortex-jni/src/object_store.rs @@ -62,11 +62,12 @@ pub(crate) fn make_object_store( // guard dropped at close of scope } - // OpenDAL-backed stores (Tencent COS) use schemes that `object_store` does not recognize - // natively. Resolve them via the optional `opendal` feature, and cache the result so - // subsequent calls for the same URL share a single client. + // OpenDAL-backed stores (Tencent COS, Alibaba OSS) use schemes that `object_store` does not + // recognize natively. Resolve them via the optional `opendal` feature, and cache the result so + // subsequent calls for the same URL share a single client. Asking `supports_scheme` rather + // than matching scheme strings here keeps this call site correct as services are added. #[cfg(feature = "opendal")] - if url.scheme() == "cos" { + if vortex_object_store_opendal::supports_scheme(url.scheme()) { let store = vortex_object_store_opendal::make_opendal_store(url, properties) .map_err(|e| VortexError::from(object_store::Error::from(e)))?; return cache_and_return(store, url, properties, &start); diff --git a/vortex-object-store-opendal/Cargo.toml b/vortex-object-store-opendal/Cargo.toml index 79005633382..424b8a255fb 100644 --- a/vortex-object-store-opendal/Cargo.toml +++ b/vortex-object-store-opendal/Cargo.toml @@ -1,7 +1,6 @@ [package] name = "vortex-object-store-opendal" -description = "OpenDAL-backed object store (Tencent COS) for Vortex" -publish = false +description = "OpenDAL-backed object stores (Tencent COS, Alibaba OSS) for Vortex" version = { workspace = true } homepage = { workspace = true } repository = { workspace = true } @@ -10,18 +9,25 @@ license = { workspace = true } keywords = { workspace = true } include = { workspace = true } edition = { workspace = true } +readme = { workspace = true } rust-version = { workspace = true } categories = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [dependencies] object_store = { workspace = true, features = ["cloud"] } -object_store_opendal = "0.57.0" -opendal = { version = "0.57.0", default-features = false, features = [ - "services-cos", -] } +object_store_opendal = { workspace = true } +opendal = { workspace = true } tracing = { workspace = true } url = { workspace = true } vortex-utils = { workspace = true } +[features] +default = ["cos", "oss"] +cos = ["opendal/services-cos"] +oss = ["opendal/services-oss"] + [lints] workspace = true diff --git a/vortex-object-store-opendal/src/cos.rs b/vortex-object-store-opendal/src/cos.rs new file mode 100644 index 00000000000..b873855403a --- /dev/null +++ b/vortex-object-store-opendal/src/cos.rs @@ -0,0 +1,231 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tencent Cloud COS, served over OpenDAL's `services::Cos`. + +use std::sync::Arc; + +use object_store::ObjectStore; +use object_store_opendal::OpendalStore; +use opendal::services; +use url::Url; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::OpenDALStoreError; +use crate::build_operator; +use crate::property_or_env; +use crate::warn_on_unknown_properties; + +/// The URL scheme served by Tencent Cloud COS. +pub const COS_SCHEME: &str = "cos"; + +/// Property keys recognized for `cos://` URLs. Anything else is warned about and dropped. +const KNOWN_PROPERTIES: &[&str] = &[ + "bucket", + "endpoint", + "secret_id", + "secret_key", + "root", + "disable_config_load", +]; + +/// Strongly-typed configuration for building an OpenDAL store against Tencent Cloud COS. +/// +/// The fields mirror the keyword arguments of the `CosStore` Python class. Building from a +/// `CosConfig` avoids the URL-round-trip that the [`crate::make_opendal_store`] entry point uses, +/// and is the preferred way to construct a COS store. +#[derive(Debug, Clone, Default)] +pub struct CosConfig { + /// COS bucket name. May be overridden by `properties["bucket"]` when adapting a URL. + pub bucket: String, + /// COS endpoint, e.g. `https://cos.ap-guangzhou.myqcloud.com`. + pub endpoint: String, + /// Tencent Cloud secret id (mapped to `TENCENTCLOUD_SECRET_ID`). + pub secret_id: Option, + /// Tencent Cloud secret key (mapped to `TENCENTCLOUD_SECRET_KEY`). + pub secret_key: Option, + /// Optional root prefix applied to all operations. + pub root: Option, + /// When `true`, disable OpenDAL's automatic config loading (so only the explicit + /// configuration is used). + pub disable_config_load: bool, +} + +/// Build an [`object_store::ObjectStore`] for Tencent Cloud COS directly from a [`CosConfig`]. +/// +/// This is the preferred entry point for callers that have a strongly-typed configuration object +/// (such as the `CosStore` pyclass in `vortex-python`). It does not synthesize a URL and so is not +/// fragile against reordering of `bucket` vs URL host precedence. +pub fn make_cos_store(config: CosConfig) -> Result, OpenDALStoreError> { + if config.bucket.is_empty() { + return Err(OpenDALStoreError::MissingConfig("bucket")); + } + if config.endpoint.is_empty() { + return Err(OpenDALStoreError::MissingConfig("endpoint")); + } + + let mut builder = services::Cos::default() + .bucket(&config.bucket) + .endpoint(&config.endpoint); + + if let Some(root) = config.root.as_deref() { + builder = builder.root(root); + } + if let Some(secret_id) = config.secret_id.as_deref() { + builder = builder.secret_id(secret_id); + } + if let Some(secret_key) = config.secret_key.as_deref() { + builder = builder.secret_key(secret_key); + } + if config.disable_config_load { + builder = builder.disable_config_load(); + } + + let operator = build_operator(builder)?; + Ok(Arc::new(OpendalStore::new(operator))) +} + +/// Translate a (`cos://` URL, properties) pair into a strongly-typed [`CosConfig`]. +/// +/// Bucket is taken from `properties["bucket"]` first and falls back to the URL host; endpoint is +/// taken from `properties["endpoint"]` first and falls back to `COS_ENDPOINT`; credentials fall +/// back to the variables OpenDAL's COS builder reads. +/// +/// `env_lookup` is the source of truth for environment-variable fallbacks. The production entry +/// point passes the real environment; tests pass a closure that returns from a fixed map, so they +/// do not race against the process environment. +pub(crate) fn url_and_properties_to_config( + url: &Url, + properties: &HashMap, + env_lookup: F, +) -> Result +where + F: Fn(&str) -> Option, +{ + warn_on_unknown_properties(properties, KNOWN_PROPERTIES); + + let bucket = properties + .get("bucket") + .cloned() + .or_else(|| url.host_str().map(str::to_string)) + .ok_or(OpenDALStoreError::MissingConfig("bucket"))?; + + let endpoint = property_or_env(properties, "endpoint", "COS_ENDPOINT", &env_lookup) + .ok_or(OpenDALStoreError::MissingConfig("endpoint"))?; + + Ok(CosConfig { + bucket, + endpoint, + secret_id: property_or_env( + properties, + "secret_id", + "TENCENTCLOUD_SECRET_ID", + &env_lookup, + ), + secret_key: property_or_env( + properties, + "secret_key", + "TENCENTCLOUD_SECRET_KEY", + &env_lookup, + ), + root: properties.get("root").cloned(), + disable_config_load: properties.get("disable_config_load").map(String::as_str) + == Some("true"), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The endpoint is required. With a fixed env-lookup that returns `None`, the call must + /// fail with `MissingConfig("endpoint")` regardless of what the process environment + /// happens to contain. This makes the test deterministic under `cargo test`'s default + /// multi-threaded runner (and avoids the `unsafe std::env::set_var` that became unsound in + /// Rust 2024). + #[test] + fn cos_requires_endpoint() { + let url = Url::parse("cos://my-bucket/path").unwrap(); + let props = HashMap::new(); + let result = url_and_properties_to_config(&url, &props, |_| None).unwrap_err(); + assert!(matches!( + result, + OpenDALStoreError::MissingConfig("endpoint") + )); + } + + /// When `properties` does not contain `endpoint`, the env-lookup should be consulted. + #[test] + fn cos_falls_back_to_cos_endpoint_env() { + let url = Url::parse("cos://my-bucket/path").unwrap(); + let env = |key: &str| match key { + "COS_ENDPOINT" => Some("https://example.com".to_string()), + _ => None, + }; + let props = HashMap::new(); + let config = url_and_properties_to_config(&url, &props, env).expect("config"); + assert_eq!(config.endpoint, "https://example.com"); + } + + /// An explicit property must win over the environment fallback. + #[test] + fn cos_property_overrides_env() { + let url = Url::parse("cos://url-bucket/path").unwrap(); + let env = |key: &str| match key { + "COS_ENDPOINT" => Some("https://from-env.example.com".to_string()), + _ => None, + }; + let mut props = HashMap::new(); + props.insert("bucket".to_string(), "prop-bucket".to_string()); + props.insert( + "endpoint".to_string(), + "https://from-prop.example.com".to_string(), + ); + + let config = url_and_properties_to_config(&url, &props, env).expect("config"); + assert_eq!(config.bucket, "prop-bucket"); + assert_eq!(config.endpoint, "https://from-prop.example.com"); + } + + /// With a fixed env-lookup that returns explicit credentials, the strongly-typed + /// `make_cos_store` should build a store successfully. + #[test] + fn cos_builds_with_explicit_config() { + let url = Url::parse("cos://my-bucket/path/to/dataset.vortex").unwrap(); + let env = |key: &str| match key { + "COS_ENDPOINT" => Some("https://example.com".to_string()), + "TENCENTCLOUD_SECRET_ID" => Some("AKID".to_string()), + "TENCENTCLOUD_SECRET_KEY" => Some("secret".to_string()), + _ => None, + }; + let mut props = HashMap::new(); + props.insert("disable_config_load".to_string(), "true".to_string()); + + let config = url_and_properties_to_config(&url, &props, env).expect("config"); + assert_eq!(config.secret_id.as_deref(), Some("AKID")); + let store = make_cos_store(config).expect("store should build"); + // Sanity: the returned store is a non-null `Arc`. + assert!(Arc::strong_count(&store) >= 1); + } + + /// The strongly-typed `make_cos_store` entry point must reject an empty bucket or endpoint + /// with a `MissingConfig` error before consulting any environment or builder. + #[test] + fn cos_config_rejects_empty_fields() { + assert!(matches!( + make_cos_store(CosConfig { + bucket: String::new(), + ..CosConfig::default() + }), + Err(OpenDALStoreError::MissingConfig("bucket")) + )); + assert!(matches!( + make_cos_store(CosConfig { + bucket: "b".to_string(), + endpoint: String::new(), + ..CosConfig::default() + }), + Err(OpenDALStoreError::MissingConfig("endpoint")) + )); + } +} diff --git a/vortex-object-store-opendal/src/lib.rs b/vortex-object-store-opendal/src/lib.rs index 92caee88092..5682ee6b001 100644 --- a/vortex-object-store-opendal/src/lib.rs +++ b/vortex-object-store-opendal/src/lib.rs @@ -2,36 +2,67 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! OpenDAL-backed [`object_store::ObjectStore`] implementations for cloud providers that are not -//! natively supported by the `object_store` crate, such as Tencent Cloud COS. +//! natively supported by the `object_store` crate: Tencent Cloud COS and Alibaba Cloud OSS. //! //! OpenDAL exposes each service as an `Operator`. We adapt an `Operator` into an //! `object_store::ObjectStore` via the `object_store_opendal::OpendalStore` bridge, which is built //! against the same `object_store 0.13.x` version the rest of Vortex uses. This lets Vortex consume -//! COS through its existing `ObjectStoreFileSystem` abstraction without any changes to -//! `vortex-io`. +//! these services through its existing `ObjectStoreFileSystem` abstraction. +//! +//! Callers that dispatch on a URL scheme should ask [`supports_scheme`] rather than comparing +//! against [`COS_SCHEME`] / [`OSS_SCHEME`] themselves, so that enabling another service does not +//! require touching every call site. +//! +//! # Cargo features +//! +//! * `cos` (default) — Tencent Cloud COS, the `cos://` scheme. +//! * `oss` (default) — Alibaba Cloud OSS, the `oss://` scheme. +//! +//! With both features disabled the crate still compiles: [`supports_scheme`] returns `false` for +//! every scheme and [`make_opendal_store`] always reports +//! [`OpenDALStoreError::UnsupportedScheme`]. //! //! # Limitations //! -//! The [`OpendalStore`] bridge owns its own HTTP request client. Configuration that the JNI/Python +//! The `OpendalStore` bridge owns its own HTTP request client. Configuration that the JNI/Python //! layers normally pass through [`object_store::ClientOptions`] — connect/request timeouts, -//! retries, proxy settings, `allow_http` — has no effect on COS URLs handled here. Properties -//! that are not recognized by this crate are dropped without a warning; callers that need strict +//! retries, proxy settings, `allow_http` — has no effect on URLs handled here. Properties that a +//! service does not recognize are logged at `warn` and dropped; callers that need strict //! validation must pre-filter their property maps. +#[cfg(feature = "cos")] +mod cos; +#[cfg(feature = "oss")] +mod oss; + use std::sync::Arc; use object_store::ObjectStore; -use object_store_opendal::OpendalStore; +#[cfg(any(feature = "cos", feature = "oss"))] use opendal::Operator; -use opendal::services; +#[cfg(any(feature = "cos", feature = "oss"))] use tracing::warn; use url::Url; use vortex_utils::aliases::hash_map::HashMap; +#[cfg(feature = "cos")] +pub use crate::cos::COS_SCHEME; +#[cfg(feature = "cos")] +pub use crate::cos::CosConfig; +#[cfg(feature = "cos")] +pub use crate::cos::make_cos_store; +#[cfg(feature = "oss")] +pub use crate::oss::OSS_SCHEME; +#[cfg(feature = "oss")] +pub use crate::oss::OssConfig; +#[cfg(feature = "oss")] +pub use crate::oss::make_oss_store; + /// Error type for building an OpenDAL-backed object store. #[derive(Debug)] pub enum OpenDALStoreError { - /// The URL scheme is not one this crate handles (e.g. `s3`, `gs`, ...). + /// The URL scheme is not one this crate handles (e.g. `s3`, `gs`, ...), or its Cargo feature + /// is not enabled. UnsupportedScheme(String), /// A required configuration value (bucket and/or endpoint) was missing. MissingConfig(&'static str), @@ -64,89 +95,72 @@ impl From for object_store::Error { } } -/// Schemes handled by this crate. -pub const COS_SCHEME: &str = "cos"; +/// The URL schemes this crate can build stores for, given the enabled Cargo features. +pub const SUPPORTED_SCHEMES: &[&str] = &[ + #[cfg(feature = "cos")] + COS_SCHEME, + #[cfg(feature = "oss")] + OSS_SCHEME, +]; -/// Strongly-typed configuration for building a [`OpendalStore`] against Tencent Cloud COS. +/// Returns `true` if `scheme` is served by an OpenDAL-backed store in this build. /// -/// The fields mirror the keyword arguments of the `CosStore` Python class. Building from a -/// `CosConfig` avoids the URL-round-trip that the legacy `make_opendal_store(url, properties)` -/// entry point used, and is the preferred way to construct a COS store. -#[derive(Debug, Clone, Default)] -pub struct CosConfig { - /// COS bucket name. May be overridden by `properties["bucket"]` when adapting a URL. - pub bucket: String, - /// COS endpoint, e.g. `https://cos.ap-guangzhou.myqcloud.com`. - pub endpoint: String, - /// Tencent Cloud secret id (mapped to `TENCENTCLOUD_SECRET_ID`). - pub secret_id: Option, - /// Tencent Cloud secret key (mapped to `TENCENTCLOUD_SECRET_KEY`). - pub secret_key: Option, - /// Optional root prefix applied to all operations. - pub root: Option, - /// When `true`, disable OpenDAL's automatic config loading (so only the explicit - /// configuration is used). - pub disable_config_load: bool, -} - -/// Build an [`object_store::ObjectStore`] for Tencent Cloud COS directly from a [`CosConfig`]. +/// This is the dispatch predicate every caller should use: it tracks the enabled Cargo features, +/// so a consumer that adds a service feature picks it up without changing its own scheme matching. /// -/// This is the preferred entry point for callers that have a strongly-typed configuration object -/// (such as the `CosStore` pyclass in `vortex-python`). It does not synthesize a URL and so is not -/// fragile against reordering of `bucket` vs URL host precedence. -pub fn make_cos_store(config: CosConfig) -> Result, OpenDALStoreError> { - if config.bucket.is_empty() { - return Err(OpenDALStoreError::MissingConfig("bucket")); - } - if config.endpoint.is_empty() { - return Err(OpenDALStoreError::MissingConfig("endpoint")); - } - - let mut builder = services::Cos::default() - .bucket(&config.bucket) - .endpoint(&config.endpoint); - - if let Some(root) = config.root.as_deref() { - builder = builder.root(root); - } - if let Some(secret_id) = config.secret_id.as_deref() { - builder = builder.secret_id(secret_id); - } - if let Some(secret_key) = config.secret_key.as_deref() { - builder = builder.secret_key(secret_key); - } - if config.disable_config_load { - builder = builder.disable_config_load(); - } - - let operator = build_operator(builder)?; - Ok(Arc::new(OpendalStore::new(operator))) +/// ``` +/// # use vortex_object_store_opendal::supports_scheme; +/// assert!(!supports_scheme("s3")); +/// ``` +pub fn supports_scheme(scheme: &str) -> bool { + SUPPORTED_SCHEMES.contains(&scheme) } -/// Build an [`object_store::ObjectStore`] for a `cos://` URL. +/// Build an [`object_store::ObjectStore`] for an OpenDAL-backed URL (`cos://`, `oss://`). /// /// `properties` are per-request configuration overrides (matching the `HashMap` -/// passed through the JNI/Python layers). Missing values fall back to environment variables that -/// OpenDAL's builders read automatically (e.g. `TENCENTCLOUD_SECRET_ID`). +/// passed through the JNI/Python layers). Missing values fall back to the environment variables +/// the corresponding service reads (e.g. `TENCENTCLOUD_SECRET_ID`, `ALIBABA_CLOUD_ACCESS_KEY_ID`). /// -/// Returns [`OpenDALStoreError::UnsupportedScheme`] if `url` does not use `cos://`. +/// Returns [`OpenDALStoreError::UnsupportedScheme`] if `url` uses a scheme this build does not +/// serve; test it up-front with [`supports_scheme`]. pub fn make_opendal_store( url: &Url, properties: &HashMap, ) -> Result, OpenDALStoreError> { - if url.scheme() != COS_SCHEME { - return Err(OpenDALStoreError::UnsupportedScheme( - url.scheme().to_string(), - )); - } + make_opendal_store_with_env(url, properties, env_var_lookup) +} - // Translate the (URL, properties) pair into a strongly-typed `CosConfig` and delegate to - // `make_cos_store`. Bucket is taken from `properties["bucket"]` first (matching the historical - // precedence) and falls back to the URL host; endpoint is taken from `properties["endpoint"]` - // first and falls back to `COS_ENDPOINT`; credentials fall back to the variables that - // OpenDAL's COS builder reads. - let config = url_and_properties_to_cos_config(url, properties, env_var_lookup)?; - make_cos_store(config) +/// Build an OpenDAL-backed store, resolving environment fallbacks through `env_lookup` instead of +/// the process environment. +/// +/// Callers that already own a configuration source — such as a registry that resolves variables +/// case-insensitively — should use this so that store construction does not silently depend on +/// global state. [`make_opendal_store`] is the same call against the real environment. +pub fn make_opendal_store_with_env( + url: &Url, + properties: &HashMap, + env_lookup: F, +) -> Result, OpenDALStoreError> +where + F: Fn(&str) -> Option, +{ + match url.scheme() { + #[cfg(feature = "cos")] + COS_SCHEME => make_cos_store(cos::url_and_properties_to_config( + url, properties, env_lookup, + )?), + #[cfg(feature = "oss")] + OSS_SCHEME => make_oss_store(oss::url_and_properties_to_config( + url, properties, env_lookup, + )?), + other => { + // Consumes the arguments so they count as used even in a build with no service + // features enabled, where every scheme lands here. + drop((properties, env_lookup)); + Err(OpenDALStoreError::UnsupportedScheme(other.to_string())) + } + } } /// Default environment-variable lookup: reads `key` from the process environment. @@ -159,67 +173,33 @@ fn env_var_lookup(key: &str) -> Option { std::env::var(key).ok() } -/// Translate the (URL, properties) pair into a strongly-typed [`CosConfig`]. -/// -/// `env_lookup` is the source of truth for environment-variable fallbacks. The production -/// entry points pass [`env_var_lookup`]; tests pass a closure that returns from a fixed map, so -/// they do not race against the process environment. -fn url_and_properties_to_cos_config( - url: &Url, +/// Take `key` from `properties`, falling back to `env_lookup(env_var)`. +#[cfg(any(feature = "cos", feature = "oss"))] +pub(crate) fn property_or_env( properties: &HashMap, - env_lookup: F, -) -> Result + key: &str, + env_var: &str, + env_lookup: &F, +) -> Option where F: Fn(&str) -> Option, { - warn_on_unknown_properties(properties); - - let bucket = properties - .get("bucket") - .cloned() - .or_else(|| url.host_str().map(str::to_string)) - .ok_or(OpenDALStoreError::MissingConfig("bucket"))?; - - let endpoint = properties - .get("endpoint") - .cloned() - .or_else(|| env_lookup("COS_ENDPOINT")) - .ok_or(OpenDALStoreError::MissingConfig("endpoint"))?; - - Ok(CosConfig { - bucket, - endpoint, - secret_id: properties - .get("secret_id") - .cloned() - .or_else(|| env_lookup("TENCENTCLOUD_SECRET_ID")), - secret_key: properties - .get("secret_key") - .cloned() - .or_else(|| env_lookup("TENCENTCLOUD_SECRET_KEY")), - root: properties.get("root").cloned(), - disable_config_load: properties.get("disable_config_load").map(String::as_str) - == Some("true"), - }) + properties.get(key).cloned().or_else(|| env_lookup(env_var)) } -fn warn_on_unknown_properties(properties: &HashMap) { - const KNOWN: &[&str] = &[ - "bucket", - "endpoint", - "secret_id", - "secret_key", - "root", - "disable_config_load", - ]; +/// Log a warning for every property key the service does not recognize. +#[cfg(any(feature = "cos", feature = "oss"))] +pub(crate) fn warn_on_unknown_properties(properties: &HashMap, known: &[&str]) { for key in properties.keys() { - if !KNOWN.contains(&key.as_str()) { + if !known.contains(&key.as_str()) { warn!("ignoring unknown OpenDAL store property: {key}"); } } } -fn build_operator(builder: B) -> Result +/// Finish an OpenDAL builder into an [`Operator`], mapping builder errors into our error type. +#[cfg(any(feature = "cos", feature = "oss"))] +pub(crate) fn build_operator(builder: B) -> Result where B: opendal::Builder, { @@ -243,74 +223,29 @@ mod tests { )); } - /// The endpoint is required. With a fixed env-lookup that returns `None`, the call must - /// fail with `MissingConfig("endpoint")` regardless of what the process environment - /// happens to contain. This makes the test deterministic under `cargo test`'s default - /// multi-threaded runner (and avoids the `unsafe std::env::set_var` that became unsound in - /// Rust 2024). #[test] - fn cos_requires_endpoint() { - let url = Url::parse("cos://my-bucket/path").unwrap(); - let props = HashMap::new(); - let result = url_and_properties_to_cos_config(&url, &props, |_| None).unwrap_err(); - assert!(matches!( - result, - OpenDALStoreError::MissingConfig("endpoint") - )); + fn supports_scheme_tracks_enabled_features() { + assert!(!supports_scheme("s3")); + assert_eq!(supports_scheme("cos"), cfg!(feature = "cos")); + assert_eq!(supports_scheme("oss"), cfg!(feature = "oss")); } - /// When `properties` does not contain `endpoint`, the env-lookup should be consulted. - /// The fixed lookup returns `"https://example.com"`, so the build must succeed. + /// Every scheme advertised by [`SUPPORTED_SCHEMES`] must actually reach a builder rather than + /// falling through to the `UnsupportedScheme` arm of [`make_opendal_store`]. Without a + /// configured endpoint the build fails, but it must fail with `MissingConfig`, not + /// `UnsupportedScheme` — that difference is what pins the dispatch table to the feature set. #[test] - fn cos_falls_back_to_cos_endpoint_env() { - let url = Url::parse("cos://my-bucket/path").unwrap(); - let env = |key: &str| match key { - "COS_ENDPOINT" => Some("https://example.com".to_string()), - _ => None, - }; - let props = HashMap::new(); - let config = url_and_properties_to_cos_config(&url, &props, env).expect("config"); - assert_eq!(config.endpoint, "https://example.com"); - } - - /// With a fixed env-lookup that returns explicit credentials, the strongly-typed - /// `make_cos_store` should build a store successfully. - #[test] - fn cos_builds_with_explicit_config() { - let url = Url::parse("cos://my-bucket/path/to/dataset.vortex").unwrap(); - let env = |key: &str| match key { - "COS_ENDPOINT" => Some("https://example.com".to_string()), - "TENCENTCLOUD_SECRET_ID" => Some("AKID".to_string()), - "TENCENTCLOUD_SECRET_KEY" => Some("secret".to_string()), - _ => None, - }; - let mut props = HashMap::new(); - props.insert("disable_config_load".to_string(), "true".to_string()); - - let config = url_and_properties_to_cos_config(&url, &props, env).expect("config"); - let store = make_cos_store(config).expect("store should build"); - // Sanity: the returned store is a non-null `Arc`. - assert!(Arc::strong_count(&store) >= 1); - } - - /// The strongly-typed `make_cos_store` entry point must reject an empty bucket or endpoint - /// with a `MissingConfig` error before consulting any environment or builder. - #[test] - fn cos_config_rejects_empty_fields() { - assert!(matches!( - make_cos_store(CosConfig { - bucket: String::new(), - ..CosConfig::default() - }), - Err(OpenDALStoreError::MissingConfig("bucket")) - )); - assert!(matches!( - make_cos_store(CosConfig { - bucket: "b".to_string(), - endpoint: String::new(), - ..CosConfig::default() - }), - Err(OpenDALStoreError::MissingConfig("endpoint")) - )); + fn every_supported_scheme_dispatches() { + for scheme in SUPPORTED_SCHEMES { + let url = Url::parse(&format!("{scheme}://bucket/path")).unwrap(); + let props = HashMap::new(); + assert!( + !matches!( + make_opendal_store(&url, &props), + Err(OpenDALStoreError::UnsupportedScheme(_)) + ), + "{scheme} is advertised but not dispatched" + ); + } } } diff --git a/vortex-object-store-opendal/src/oss.rs b/vortex-object-store-opendal/src/oss.rs new file mode 100644 index 00000000000..3ff3b049728 --- /dev/null +++ b/vortex-object-store-opendal/src/oss.rs @@ -0,0 +1,257 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Alibaba Cloud OSS, served over OpenDAL's `services::Oss`. + +use std::sync::Arc; + +use object_store::ObjectStore; +use object_store_opendal::OpendalStore; +use opendal::services; +use url::Url; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::OpenDALStoreError; +use crate::build_operator; +use crate::property_or_env; +use crate::warn_on_unknown_properties; + +/// The URL scheme served by Alibaba Cloud OSS. +pub const OSS_SCHEME: &str = "oss"; + +/// Property keys recognized for `oss://` URLs. Anything else is warned about and dropped. +const KNOWN_PROPERTIES: &[&str] = &[ + "bucket", + "endpoint", + "access_key_id", + "access_key_secret", + "security_token", + "root", + "skip_signature", +]; + +/// Strongly-typed configuration for building an OpenDAL store against Alibaba Cloud OSS. +/// +/// This mirrors [`crate::CosConfig`], but OSS names its credentials `access_key_id` / +/// `access_key_secret` and has no equivalent of COS's `disable_config_load`; the closest knob is +/// [`OssConfig::skip_signature`], which sends unsigned requests for public buckets. +#[derive(Debug, Clone, Default)] +pub struct OssConfig { + /// OSS bucket name. May be overridden by `properties["bucket"]` when adapting a URL. + pub bucket: String, + /// OSS endpoint, e.g. `https://oss-cn-hangzhou.aliyuncs.com`. + pub endpoint: String, + /// Alibaba Cloud access key id (mapped to `ALIBABA_CLOUD_ACCESS_KEY_ID`). + pub access_key_id: Option, + /// Alibaba Cloud access key secret (mapped to `ALIBABA_CLOUD_ACCESS_KEY_SECRET`). + pub access_key_secret: Option, + /// STS security token for temporary credentials (mapped to `ALIBABA_CLOUD_SECURITY_TOKEN`). + pub security_token: Option, + /// Optional root prefix applied to all operations. + pub root: Option, + /// When `true`, send unsigned requests. Use this for public, read-only buckets. + pub skip_signature: bool, +} + +/// Build an [`object_store::ObjectStore`] for Alibaba Cloud OSS directly from an [`OssConfig`]. +/// +/// This is the preferred entry point for callers that have a strongly-typed configuration object. +/// It does not synthesize a URL and so is not fragile against reordering of `bucket` vs URL host +/// precedence. +pub fn make_oss_store(config: OssConfig) -> Result, OpenDALStoreError> { + if config.bucket.is_empty() { + return Err(OpenDALStoreError::MissingConfig("bucket")); + } + if config.endpoint.is_empty() { + return Err(OpenDALStoreError::MissingConfig("endpoint")); + } + + let mut builder = services::Oss::default() + .bucket(&config.bucket) + .endpoint(&config.endpoint); + + if let Some(root) = config.root.as_deref() { + builder = builder.root(root); + } + if let Some(access_key_id) = config.access_key_id.as_deref() { + builder = builder.access_key_id(access_key_id); + } + if let Some(access_key_secret) = config.access_key_secret.as_deref() { + builder = builder.access_key_secret(access_key_secret); + } + if let Some(security_token) = config.security_token.as_deref() { + builder = builder.security_token(security_token); + } + if config.skip_signature { + builder = builder.skip_signature(); + } + + let operator = build_operator(builder)?; + Ok(Arc::new(OpendalStore::new(operator))) +} + +/// Translate an (`oss://` URL, properties) pair into a strongly-typed [`OssConfig`]. +/// +/// Bucket is taken from `properties["bucket"]` first and falls back to the URL host; endpoint is +/// taken from `properties["endpoint"]` first and falls back to `OSS_ENDPOINT`; credentials fall +/// back to the `ALIBABA_CLOUD_*` variables OpenDAL's OSS builder reads. +/// +/// `env_lookup` is the source of truth for environment-variable fallbacks. The production entry +/// point passes the real environment; tests pass a closure that returns from a fixed map, so they +/// do not race against the process environment. +pub(crate) fn url_and_properties_to_config( + url: &Url, + properties: &HashMap, + env_lookup: F, +) -> Result +where + F: Fn(&str) -> Option, +{ + warn_on_unknown_properties(properties, KNOWN_PROPERTIES); + + let bucket = properties + .get("bucket") + .cloned() + .or_else(|| url.host_str().map(str::to_string)) + .ok_or(OpenDALStoreError::MissingConfig("bucket"))?; + + let endpoint = property_or_env(properties, "endpoint", "OSS_ENDPOINT", &env_lookup) + .ok_or(OpenDALStoreError::MissingConfig("endpoint"))?; + + Ok(OssConfig { + bucket, + endpoint, + access_key_id: property_or_env( + properties, + "access_key_id", + "ALIBABA_CLOUD_ACCESS_KEY_ID", + &env_lookup, + ), + access_key_secret: property_or_env( + properties, + "access_key_secret", + "ALIBABA_CLOUD_ACCESS_KEY_SECRET", + &env_lookup, + ), + security_token: property_or_env( + properties, + "security_token", + "ALIBABA_CLOUD_SECURITY_TOKEN", + &env_lookup, + ), + root: properties.get("root").cloned(), + skip_signature: properties.get("skip_signature").map(String::as_str) == Some("true"), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The endpoint is required, and the failure must be deterministic regardless of what the + /// process environment happens to contain — hence the fixed `|_| None` lookup. + #[test] + fn oss_requires_endpoint() { + let url = Url::parse("oss://my-bucket/path").unwrap(); + let props = HashMap::new(); + let result = url_and_properties_to_config(&url, &props, |_| None).unwrap_err(); + assert!(matches!( + result, + OpenDALStoreError::MissingConfig("endpoint") + )); + } + + /// OpenDAL's OSS service reads `ALIBABA_CLOUD_*` (with the underscore), not `ALIBABACLOUD_*`. + /// Pinning the exact names here keeps the fallbacks honest. + #[test] + fn oss_falls_back_to_env() { + let url = Url::parse("oss://my-bucket/path").unwrap(); + let env = |key: &str| match key { + "OSS_ENDPOINT" => Some("https://oss-cn-hangzhou.aliyuncs.com".to_string()), + "ALIBABA_CLOUD_ACCESS_KEY_ID" => Some("LTAI".to_string()), + "ALIBABA_CLOUD_ACCESS_KEY_SECRET" => Some("secret".to_string()), + "ALIBABA_CLOUD_SECURITY_TOKEN" => Some("token".to_string()), + _ => None, + }; + let props = HashMap::new(); + let config = url_and_properties_to_config(&url, &props, env).expect("config"); + + assert_eq!(config.bucket, "my-bucket"); + assert_eq!(config.endpoint, "https://oss-cn-hangzhou.aliyuncs.com"); + assert_eq!(config.access_key_id.as_deref(), Some("LTAI")); + assert_eq!(config.access_key_secret.as_deref(), Some("secret")); + assert_eq!(config.security_token.as_deref(), Some("token")); + } + + /// An explicit property must win over the environment fallback. + #[test] + fn oss_property_overrides_env() { + let url = Url::parse("oss://url-bucket/path").unwrap(); + let env = |key: &str| match key { + "OSS_ENDPOINT" => Some("https://from-env.example.com".to_string()), + "ALIBABA_CLOUD_ACCESS_KEY_ID" => Some("from-env".to_string()), + _ => None, + }; + let mut props = HashMap::new(); + props.insert("bucket".to_string(), "prop-bucket".to_string()); + props.insert( + "endpoint".to_string(), + "https://from-prop.example.com".to_string(), + ); + props.insert("access_key_id".to_string(), "from-prop".to_string()); + + let config = url_and_properties_to_config(&url, &props, env).expect("config"); + assert_eq!(config.bucket, "prop-bucket"); + assert_eq!(config.endpoint, "https://from-prop.example.com"); + assert_eq!(config.access_key_id.as_deref(), Some("from-prop")); + } + + /// A full config must reach the OpenDAL builder and produce a store. + #[test] + fn oss_builds_with_explicit_config() { + let url = Url::parse("oss://my-bucket/path/to/dataset.vortex").unwrap(); + let env = |key: &str| match key { + "OSS_ENDPOINT" => Some("https://oss-cn-hangzhou.aliyuncs.com".to_string()), + "ALIBABA_CLOUD_ACCESS_KEY_ID" => Some("LTAI".to_string()), + "ALIBABA_CLOUD_ACCESS_KEY_SECRET" => Some("secret".to_string()), + _ => None, + }; + let props = HashMap::new(); + let config = url_and_properties_to_config(&url, &props, env).expect("config"); + let store = make_oss_store(config).expect("store should build"); + assert!(Arc::strong_count(&store) >= 1); + } + + /// Public buckets need no credentials when `skip_signature` is set. + #[test] + fn oss_builds_anonymous() { + let store = make_oss_store(OssConfig { + bucket: "public-bucket".to_string(), + endpoint: "https://oss-cn-hangzhou.aliyuncs.com".to_string(), + skip_signature: true, + ..OssConfig::default() + }) + .expect("anonymous store should build"); + assert!(Arc::strong_count(&store) >= 1); + } + + /// `make_oss_store` must reject an empty bucket or endpoint before touching the builder. + #[test] + fn oss_config_rejects_empty_fields() { + assert!(matches!( + make_oss_store(OssConfig { + bucket: String::new(), + ..OssConfig::default() + }), + Err(OpenDALStoreError::MissingConfig("bucket")) + )); + assert!(matches!( + make_oss_store(OssConfig { + bucket: "b".to_string(), + endpoint: String::new(), + ..OssConfig::default() + }), + Err(OpenDALStoreError::MissingConfig("endpoint")) + )); + } +} diff --git a/vortex-python/Cargo.toml b/vortex-python/Cargo.toml index 3447912861a..b42ffd32c4e 100644 --- a/vortex-python/Cargo.toml +++ b/vortex-python/Cargo.toml @@ -28,9 +28,9 @@ default = ["extension-module", "tui"] # duplicate PyInit__lib symbols. extension-module = [] tui = ["dep:tokio", "dep:vortex-tui"] -# OpenDAL-backed object stores (Tencent COS) for `cos://` URLs. Pulls in `opendal`, so it is -# opt-in. -opendal = ["dep:vortex-object-store-opendal", "dep:vortex-utils"] +# OpenDAL-backed object stores (Tencent COS, Alibaba OSS) for `cos://` and `oss://` URLs. Pulls in +# `opendal`, so it is opt-in. +opendal = ["dep:vortex-object-store-opendal", "vortex/opendal"] [dependencies] arrow-array = { workspace = true } @@ -54,12 +54,17 @@ pyo3-log = { workspace = true } pyo3-object_store = { workspace = true } tokio = { workspace = true, features = ["rt-multi-thread"], optional = true } url = { workspace = true } -vortex = { workspace = true, features = ["object_store"] } +vortex = { workspace = true, features = [ + "object_store", + "object_store_registry", +] } vortex-arrow = { workspace = true } -vortex-object-store-opendal = { path = "../vortex-object-store-opendal", optional = true } +vortex-object-store-opendal = { workspace = true, optional = true, features = [ + "cos", + "oss", +] } vortex-python-abi = { path = "../vortex-python-abi" } vortex-tui = { workspace = true, optional = true } -vortex-utils = { workspace = true, optional = true } [dev-dependencies] vortex-array = { path = "../vortex-array", features = ["_test-harness"] } diff --git a/vortex-python/src/object_store/mod.rs b/vortex-python/src/object_store/mod.rs index d43e32ce137..034d9546f56 100644 --- a/vortex-python/src/object_store/mod.rs +++ b/vortex-python/src/object_store/mod.rs @@ -1,5 +1,4 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -pub(crate) mod registry; pub(crate) mod resolve; diff --git a/vortex-python/src/object_store/registry.rs b/vortex-python/src/object_store/registry.rs deleted file mode 100644 index 144a0d0e536..00000000000 --- a/vortex-python/src/object_store/registry.rs +++ /dev/null @@ -1,309 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Apache Software Foundation (ASF) - -//! This file is an adapted version of the `DefaultObjectStoreRegistry` from the object_store crate, -//! but modified to resolve configurations out of environment variables case-insensitively. This -//! is similar to how all the `Store::from_env` builders work for the various object stores. -//! -//! See also - -#![expect(clippy::disallowed_types)] - -use std::collections::HashMap; -use std::sync::Arc; - -use object_store::ObjectStore; -use object_store::parse_url_opts; -use object_store::path::Path; -use object_store::path::PathPart; -use object_store::registry::ObjectStoreRegistry; -use parking_lot::RwLock; -use url::Url; -#[cfg(feature = "opendal")] -use vortex_utils::aliases::hash_map::HashMap as VortexHashMap; - -#[derive(Debug, Default)] -struct PathEntry { - /// Store, if defined at this path - store: Option>, - /// Child [`PathEntry`], keyed by the next path segment in their path - children: HashMap, -} - -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, 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()) { - match current.children.get(segment) { - Some(e) => { - current = e; - depth += 1; - if let Some(store) = ¤t.store { - ret = Some((store, depth)) - } - } - None => break, - } - } - ret - } -} - -/// An implementation of the [`ObjectStoreRegistry`] that normalizes environment variables -/// before doing lookups. -#[derive(Debug, Default)] -pub(crate) struct Registry { - /// Mapping from [`url_key`] to [`PathEntry`] - map: RwLock>, -} - -impl ObjectStoreRegistry for Registry { - fn register(&self, url: Url, store: Arc) -> Option> { - let mut map = self.map.write(); - let entry = entry_at( - &mut map, - url_key(&url), - url.path(), - num_segments(url.path()), - ); - entry.store.replace(store) - } - - fn resolve(&self, to_resolve: &Url) -> object_store::Result<(Arc, Path)> { - let key = url_key(to_resolve); - - // 1. Look up the user-registered map first. Every other scheme does this, so an explicit - // `Registry::register("cos://...", store)` should also win over the build-and-cache - // fallback below. This also means a previously-resolved store (built by us for the - // same URL) is served from the cache here, without rebuilding the client. - { - 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)); - } - } - - // 2. Build the store and the path *within* that store. `build_store` reports the path - // the way `parse_url_opts` does — i.e. the path is what the store will see as the - // key — which is what keeps every scheme on the one caching rule in - // `cache_and_resolve`. The depth is then just the difference in segment counts. - let (store, path) = build_store(to_resolve)?; - let depth = num_segments(to_resolve.path()) - num_segments(path.as_ref()); - - // 3. Cache the store and return it alongside the path that lives inside the store. - self.cache_and_resolve(to_resolve, store, depth) - } -} - -impl Registry { - /// Caches `store` as the store serving the first `depth` path segments of `to_resolve`, and - /// returns it alongside the remaining path. - /// - /// If a racing `resolve` cached a store at the same position first, that store wins and the - /// one just built is dropped, so all callers of a given prefix share a single client. - fn cache_and_resolve( - &self, - to_resolve: &Url, - store: Arc, - depth: usize, - ) -> object_store::Result<(Arc, Path)> { - let path = path_suffix(to_resolve, depth)?; - - let mut map = self.map.write(); - let entry = entry_at(&mut map, url_key(to_resolve), to_resolve.path(), depth); - let stored = Arc::clone(match &entry.store { - None => entry.store.insert(store), - Some(existing) => existing, // Racing creation - use existing - }); - - Ok((stored, path)) - } -} - -/// Builds the [`ObjectStore`] that serves `to_resolve`, with the path within that store. -/// -/// This mirrors [`parse_url_opts`], extended with schemes the `object_store` crate does not -/// recognize. Reporting the path the way `parse_url_opts` does is what keeps every scheme on -/// the one caching rule in [`Registry::resolve`]: a scheme says where its store is rooted, and -/// the registry decides how to cache it. -fn build_store(to_resolve: &Url) -> object_store::Result<(Arc, Path)> { - // OpenDAL-backed schemes (Tencent COS) are not recognized by `object_store`, so build them - // from OpenDAL's own environment-variable configuration (e.g. `TENCENTCLOUD_SECRET_ID`). - // The operator is rooted at the bucket, which lives in the URL authority, so — exactly as - // for `s3://bucket/path` — the whole URL path is the object key. - #[cfg(feature = "opendal")] - if to_resolve.scheme() == vortex_object_store_opendal::COS_SCHEME { - let store = - vortex_object_store_opendal::make_opendal_store(to_resolve, &VortexHashMap::new())?; - return Ok((store, Path::from_url_path(to_resolve.path())?)); - } - - let normalized_env = std::env::vars().map(|(k, v)| (k.to_ascii_lowercase(), v)); - let (store, path) = parse_url_opts(to_resolve, normalized_env)?; - Ok((Arc::from(store), path)) -} - -/// Extracts the scheme and authority of a URL (components before the Path) -fn url_key(url: &Url) -> &str { - &url[..url::Position::AfterPort] -} - -/// Returns the non-empty segments of a path -/// -/// Note: We don't use [`Url::path_segments`] as we only want non-empty paths -fn path_segments(s: &str) -> impl Iterator { - s.split('/').filter(|x| !x.is_empty()) -} - -/// Returns the number of non-empty path segments in a path -fn num_segments(s: &str) -> usize { - path_segments(s).count() -} - -/// Returns the path of `url` skipping the first `depth` segments -fn path_suffix(url: &Url, depth: usize) -> Result { - let segments = path_segments(url.path()).skip(depth); - let path = segments - .map(PathPart::parse) - .collect::>() - .map_err(|e| object_store::Error::Generic { - store: "ObjectStoreRegistry", - source: Box::new(e), - })?; - Ok(path) -} - -/// Walks to the [`PathEntry`] for `key` sitting `depth` segments into `path`, creating the -/// intermediate entries as needed. -fn entry_at<'a>( - map: &'a mut HashMap, - key: &str, - path: &str, - depth: usize, -) -> &'a mut PathEntry { - let mut current = map.entry(key.to_string()).or_default(); - for segment in path_segments(path).take(depth) { - current = current.children.entry(segment.to_string()).or_default(); - } - current -} - -#[cfg(test)] -mod tests { - use std::fmt::Write; - use std::sync::Arc; - - use object_store::ObjectStore; - use object_store::registry::ObjectStoreRegistry; - use url::Url; - - use crate::object_store::registry::Registry; - - fn with_var(key: &str, value: &str, func: F) - where - F: FnOnce(), - { - let old_val = std::env::var(key).ok(); - - // SAFETY: these unit tests run single-threaded. - unsafe { std::env::set_var(key, value) }; - - func(); - - // Set the variable back to its original value - match old_val { - None => { - unsafe { std::env::remove_var(key) }; - } - Some(val) => { - unsafe { std::env::set_var(key, val) }; - } - } - } - - #[test] - #[expect(clippy::use_debug)] - fn test_resolve_url() { - with_var("AWS_REGION", "us-east-3", || { - let registry = Registry::default(); - let (store, _) = registry - .resolve(&Url::parse("s3://my-bucket/test").unwrap()) - .unwrap(); - - // NOTE(aduffy): object_store doesn't let us downcast stores, the only way to verify - // that a configuration was added was to validate that it ends up in the Debug - // output :/ - let mut debug_str = String::new(); - write!(&mut debug_str, "{store:?}").unwrap(); - - assert!(debug_str.contains("us-east-3")); - }); - } - - /// Two resolves of the same URL must return the same `Arc` (the cached client) and the same - /// path, and two different keys must get distinct stores. This pins the symmetry the registry - /// promises: cache hit returns the same client and the same key. - #[test] - fn test_resolve_url_caches_per_key() { - with_var("AWS_REGION", "us-east-3", || { - let registry = Registry::default(); - let first = Url::parse("s3://my-bucket/first/second").unwrap(); - let second = Url::parse("s3://my-bucket/first/second").unwrap(); - - let (store_a, path_a) = registry.resolve(&first).unwrap(); - let (store_b, path_b) = registry.resolve(&second).unwrap(); - assert!(Arc::ptr_eq(&store_a, &store_b)); - assert_eq!(path_a, path_b); - - // A different bucket gets its own client. - let other = Url::parse("s3://other-bucket/first/second").unwrap(); - let (store_c, _) = registry.resolve(&other).unwrap(); - assert!(!Arc::ptr_eq(&store_a, &store_c)); - }); - } - - /// Two objects in the same bucket must share a cached client. This is the regression that - /// the bespoke "walk every segment then return the whole key" path caused: by walking all - /// segments it stored the store at full depth, so the next resolve saw the whole key - /// already consumed and returned an empty path. Pinning the path here guards against that. - #[test] - fn test_resolve_url_shared_client_same_bucket() { - with_var("AWS_REGION", "us-east-3", || { - let registry = Registry::default(); - let a = Url::parse("s3://my-bucket/path/to/data.vortex").unwrap(); - let b = Url::parse("s3://my-bucket/other/data.vortex").unwrap(); - - let (store_a, path_a) = registry.resolve(&a).unwrap(); - let (store_b, path_b) = registry.resolve(&b).unwrap(); - assert!(Arc::ptr_eq(&store_a, &store_b)); - assert_eq!(path_a.as_ref(), "path/to/data.vortex"); - assert_eq!(path_b.as_ref(), "other/data.vortex"); - }); - } - - /// Pinning the `entry_at` helper at depth > 0 (the path-prefix case) protects the shared - /// `entry_at` / `cache_and_resolve` helpers from regressing prefix registration. - #[test] - fn test_register_at_prefix_shares_store() { - let registry = Registry::default(); - let prefix = Url::parse("s3://my-bucket/prefix/").unwrap(); - let (inner, _) = registry.resolve(&prefix).unwrap(); - let cached: Arc = Arc::clone(&inner); - let replaced = registry.register(prefix.clone(), Arc::clone(&cached)); - // First registration at this prefix replaces nothing. - assert!(replaced.is_none()); - - // A resolve at a deeper key still walks through the registered prefix's store. - let deeper = Url::parse("s3://my-bucket/prefix/inner/key").unwrap(); - let (store_deeper, path_deeper) = registry.resolve(&deeper).unwrap(); - assert!(Arc::ptr_eq(&store_deeper, &cached)); - assert_eq!(path_deeper.as_ref(), "inner/key"); - } -} diff --git a/vortex-python/src/object_store/resolve.rs b/vortex-python/src/object_store/resolve.rs index 4b4943bb9c6..ea16c1aeea0 100644 --- a/vortex-python/src/object_store/resolve.rs +++ b/vortex-python/src/object_store/resolve.rs @@ -12,8 +12,7 @@ use url::Url; use vortex::error::VortexResult; use vortex::error::vortex_err; use vortex::io::compat::Compat; - -use crate::object_store::registry::Registry; +use vortex::io::object_store::Registry; static REGISTRY: LazyLock = LazyLock::new(Registry::default); diff --git a/vortex-tui/Cargo.toml b/vortex-tui/Cargo.toml index 09c929796d7..8e763480c5d 100644 --- a/vortex-tui/Cargo.toml +++ b/vortex-tui/Cargo.toml @@ -77,9 +77,9 @@ vortex-datafusion = { workspace = true, optional = true } [target.'cfg(target_arch = "wasm32")'.dependencies] console_error_panic_hook = "0.1.7" ratzilla = "0.3" -wasm-bindgen = "0.2.104" +wasm-bindgen = "0.2.108" wasm-bindgen-futures = { workspace = true } -web-sys = { version = "0.3.81", features = [ +web-sys = { version = "0.3.85", features = [ "console", "DataTransfer", "Document", diff --git a/vortex-web/crate/Cargo.toml b/vortex-web/crate/Cargo.toml index b4de2003c28..5fcdeb7152c 100644 --- a/vortex-web/crate/Cargo.toml +++ b/vortex-web/crate/Cargo.toml @@ -23,11 +23,11 @@ vortex-arrow = { workspace = true } [target.'cfg(target_arch = "wasm32")'.dependencies] futures = { workspace = true } -wasm-bindgen = "0.2.104" +wasm-bindgen = "0.2.108" wasm-bindgen-futures = { workspace = true } console_error_panic_hook = "0.1.7" -js-sys = "0.3.81" -web-sys = { version = "0.3.81", features = [ +js-sys = "0.3.85" +web-sys = { version = "0.3.85", features = [ "Blob", "console", "File", diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index 3cc3e40d65b..b2fd6282007 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -73,6 +73,10 @@ default = ["files", "zstd"] files = ["dep:vortex-file"] memmap2 = ["vortex-buffer/memmap2"] object_store = ["vortex-file?/object_store", "vortex-io/object_store"] +# The URL -> ObjectStore registry, plus the cloud backends it resolves URLs to. +object_store_registry = ["object_store", "vortex-io/object_store_registry"] +# OpenDAL-backed schemes (`cos://`, `oss://`) in the object store registry. +opendal = ["object_store_registry", "vortex-io/opendal"] tokio = [ "vortex-error/tokio", "vortex-file?/tokio",