diff --git a/Cargo.lock b/Cargo.lock index dd35d2f4725..16b0a83f9ac 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4358,10 +4358,12 @@ dependencies = [ "jiff-core", "jiff-static", "jiff-tzdb-platform", + "js-sys", "log", "portable-atomic", "portable-atomic-util", "serde_core", + "wasm-bindgen", "windows-link", ] @@ -5487,6 +5489,15 @@ dependencies = [ "digest 0.11.3", ] +[[package]] +name = "mea" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2640d335e7273dacdcf51044026139b2e269c3bb0dfc3f8cb3496b85e3f6a42c" +dependencies = [ + "slab", +] + [[package]] name = "memchr" version = "2.8.3" @@ -5941,6 +5952,23 @@ dependencies = [ "web-time", ] +[[package]] +name = "object_store_opendal" +version = "0.57.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eb12a624a41fce745838d0ef3701ff6c47797c13cd18ad3612fd2a3134fdbd8" +dependencies = [ + "async-trait", + "bytes", + "chrono", + "futures", + "mea", + "object_store", + "opendal", + "pin-project", + "tokio", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -5978,6 +6006,61 @@ version = "11.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" +[[package]] +name = "opendal" +version = "0.57.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96c9c85ce253ff87225e7669979d877a20c98a06604ec9d6dd5f4473e08f1ae1" +dependencies = [ + "opendal-core", + "opendal-service-cos", +] + +[[package]] +name = "opendal-core" +version = "0.57.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4f8607c90e2c963a91467f50fb49fbc7fb3d573f88cea219ca59ccd3740b309" +dependencies = [ + "anyhow", + "base64", + "bytes", + "futures", + "http", + "http-body", + "jiff", + "log", + "md-5 0.11.0", + "mea", + "percent-encoding", + "quick-xml", + "reqsign-core", + "reqwest 0.13.4", + "serde", + "serde_json", + "tokio", + "url", + "uuid", + "web-time", +] + +[[package]] +name = "opendal-service-cos" +version = "0.57.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa8cafe9729213375c7331019b0cb756ad3e1aff7f45cd32c45eae91ebde8901" +dependencies = [ + "bytes", + "http", + "log", + "opendal-core", + "quick-xml", + "reqsign-core", + "reqsign-file-read-tokio", + "reqsign-tencent-cos", + "serde", +] + [[package]] name = "openssl-probe" version = "0.2.1" @@ -7242,6 +7325,53 @@ dependencies = [ "bytecheck", ] +[[package]] +name = "reqsign-core" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e38b44697c60a823705ccef85cb04d8e0527c9d16ed7c58bf1c6395bdd24ceb" +dependencies = [ + "anyhow", + "base64", + "bytes", + "futures", + "hex", + "hmac", + "http", + "jiff", + "log", + "percent-encoding", + "sha1", + "sha2", + "windows-sys 0.61.2", +] + +[[package]] +name = "reqsign-file-read-tokio" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "688ff0ae421b8d4b92b53fdafaf53df2de28f428a9962edcf21702990b26f74b" +dependencies = [ + "anyhow", + "reqsign-core", + "tokio", +] + +[[package]] +name = "reqsign-tencent-cos" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6497dd9f6e3d1349b420521484099b284f95e8d3a65f088fccef42493a7b644" +dependencies = [ + "anyhow", + "http", + "log", + "percent-encoding", + "reqsign-core", + "serde", + "serde_json", +] + [[package]] name = "reqwest" version = "0.12.28" @@ -9861,6 +9991,7 @@ dependencies = [ "vortex", "vortex-arrow", "vortex-geo", + "vortex-object-store-opendal", "vortex-parquet-variant", ] @@ -9955,6 +10086,18 @@ dependencies = [ "vortex-cuda-macros", ] +[[package]] +name = "vortex-object-store-opendal" +version = "0.1.0" +dependencies = [ + "object_store", + "object_store_opendal", + "opendal", + "tracing", + "url", + "vortex-utils", +] + [[package]] name = "vortex-onpair" version = "0.1.0" @@ -10042,8 +10185,10 @@ dependencies = [ "vortex", "vortex-array", "vortex-arrow", + "vortex-object-store-opendal", "vortex-python-abi", "vortex-tui", + "vortex-utils", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index c2a3d1c20e8..a9d8bfefb85 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,6 +12,7 @@ members = [ "vortex-flatbuffers", "vortex-metrics", "vortex-io", + "vortex-object-store-opendal", "vortex-proto", "vortex-array", "vortex-arrow", diff --git a/docs/api/python/store.rst b/docs/api/python/store.rst index 305657c65f1..ce4a6f75b2d 100644 --- a/docs/api/python/store.rst +++ b/docs/api/python/store.rst @@ -13,6 +13,7 @@ Vortex arrays support reading and writing to many object storage systems: store/http store/local store/memory + store/opendal store/config .. autofunction:: vortex.store.from_url diff --git a/docs/api/python/store/opendal.rst b/docs/api/python/store/opendal.rst new file mode 100644 index 00000000000..bac62f993be --- /dev/null +++ b/docs/api/python/store/opendal.rst @@ -0,0 +1,87 @@ +============= +OpenDAL (COS) +============= + +Vortex can read from and write to Tencent Cloud COS through +`OpenDAL `_, which provides native service support. + +This store is available only when Vortex is built with the ``opendal`` feature +(e.g. ``maturin develop --features opendal`` or ``cargo build -p vortex-jni --features opendal``). + +:class:`vortex.store.CosStore` +============================== + +.. py:class:: vortex.store.CosStore(bucket, endpoint, *, secret_id=None, secret_key=None, root=None, disable_config_load=False) + + A Tencent Cloud COS object store, backed by OpenDAL. Construct it with explicit + configuration and pass it to + :func:`vortex.io.read_url` / :func:`vortex.io.write` via the ``store=`` argument, + exactly like the built-in S3/Azure/GCS stores. + + The class is only available when Vortex is built with the ``opendal`` feature; on + a default build, instantiating it raises :class:`ImportError`. + + :param bucket: COS bucket name (e.g. ``"my-bucket"``). + :param endpoint: COS endpoint (e.g. ``"https://cos.ap-guangzhou.myqcloud.com"``). + :param secret_id: Optional Tencent Cloud secret id. Maps to the ``TENCENTCLOUD_SECRET_ID`` + environment variable when unset. + :param secret_key: Optional Tencent Cloud secret key. Maps to the + ``TENCENTCLOUD_SECRET_KEY`` environment variable when unset. + :param root: Optional key prefix applied to every operation. + :param disable_config_load: When ``True``, disable OpenDAL's automatic config loading + and rely only on the explicit configuration. Defaults to ``False``. + +Reading from COS +================ + +Pass a ``cos://`` URL directly. Credentials and the endpoint are picked up from the environment +variables OpenDAL's COS builder reads (``TENCENTCLOUD_SECRET_ID``, ``TENCENTCLOUD_SECRET_KEY`` and +``COS_ENDPOINT``): + +.. code-block:: python + + import vortex as vx + + a = vx.io.read_url("cos://my-bucket/path/to/dataset.vortex") + +Or configure explicitly with :class:`~vortex.store.CosStore` and pass it to +:func:`vortex.io.read_url` via ``store=``: + +.. code-block:: python + + from vortex.io import read_url + from vortex.store import CosStore + + store = CosStore( + bucket="my-bucket", + endpoint="https://cos.ap-guangzhou.myqcloud.com", + secret_id="AKID...", + secret_key="...", + ) + + # When `store=` is supplied, the path is 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) + +Passing a store object directly +=============================== + +``CosStore`` is a concrete store object. You can build one once and pass it +directly to :func:`vortex.io.read_url` / :func:`vortex.io.write` via the ``store=`` argument, +exactly like the built-in S3/Azure/GCS stores: + +.. code-block:: python + + from vortex.io import read_url + from vortex.store import CosStore + + store = CosStore( + bucket="my-bucket", + endpoint="https://cos.ap-guangzhou.myqcloud.com", + secret_id="AKID...", + secret_key="...", + ) + + # 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) diff --git a/docs/conf.py b/docs/conf.py index a2cadbb255f..177f9a22fe4 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -54,7 +54,12 @@ git_root = Path(__file__).parent.parent nitpicky = True # ensures all :class:, :obj:, etc. links are valid -nitpick_ignore = [] +nitpick_ignore = [ + # `vortex.store.CosStore` is re-exported through the private `vortex.store._cos` module, + # and the `ObjectStore` type alias resolves to the private path. The public class is + # fully documented in `opendal.rst`; the private path is intentionally not. + ("py:class", "vortex.store._cos.CosStore"), +] doctest_global_setup = "import pyarrow; import vortex; import vortex as vx; import random; random.seed(a=0)" doctest_default_flags = ( diff --git a/vortex-jni/Cargo.toml b/vortex-jni/Cargo.toml index 50e30bb5aab..80f46a67b82 100644 --- a/vortex-jni/Cargo.toml +++ b/vortex-jni/Cargo.toml @@ -34,11 +34,17 @@ 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-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. +opendal = ["dep:vortex-object-store-opendal"] + [lib] crate-type = ["cdylib"] diff --git a/vortex-jni/src/object_store.rs b/vortex-jni/src/object_store.rs index cfc29f04d20..8e4938d661a 100644 --- a/vortex-jni/src/object_store.rs +++ b/vortex-jni/src/object_store.rs @@ -38,19 +38,21 @@ pub(crate) fn object_store_fs( Ok(Arc::new(ObjectStoreFileSystem::new(object_store, handle))) } +/// Process-wide cache of constructed object stores, keyed by URL + properties so that repeated +/// requests against the same bucket/configuration share a single client. +static OBJECT_STORES: LazyLock>>> = + LazyLock::new(|| Mutex::new(HashMap::new())); + #[expect(clippy::cognitive_complexity)] pub(crate) fn make_object_store( url: &Url, properties: &HashMap, ) -> VortexResult> { - static OBJECT_STORES: LazyLock>>> = - LazyLock::new(|| Mutex::new(HashMap::new())); - let start = std::time::Instant::now(); - let (scheme, _) = ObjectStoreScheme::parse(url) - .map_err(|error| VortexError::from(object_store::Error::from(error)))?; - + // The cache key depends only on the URL authority + sorted properties, so we can hoist it + // above the scheme dispatch. This lets every store (including OpenDAL-backed ones) share + // a single client across repeated requests against the same bucket/configuration. let cache_key = url_cache_key(url, properties); { @@ -60,6 +62,19 @@ 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. + #[cfg(feature = "opendal")] + if url.scheme() == "cos" { + 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); + } + + let (scheme, _) = ObjectStoreScheme::parse(url) + .map_err(|error| VortexError::from(object_store::Error::from(error)))?; + // Configure extra properties on that scheme instead. let store: Arc = match scheme { ObjectStoreScheme::Local => { @@ -143,6 +158,18 @@ pub(crate) fn make_object_store( } }; + cache_and_return(store, url, properties, &start) +} + +/// Insert the built store into the process-wide cache (keyed by URL + properties) and return it, +/// logging the construction latency. +fn cache_and_return( + store: Arc, + url: &Url, + properties: &HashMap, + start: &std::time::Instant, +) -> VortexResult> { + let cache_key = url_cache_key(url, properties); OBJECT_STORES.lock().insert(cache_key, Arc::clone(&store)); let duration = start.elapsed(); diff --git a/vortex-object-store-opendal/Cargo.toml b/vortex-object-store-opendal/Cargo.toml new file mode 100644 index 00000000000..79005633382 --- /dev/null +++ b/vortex-object-store-opendal/Cargo.toml @@ -0,0 +1,27 @@ +[package] +name = "vortex-object-store-opendal" +description = "OpenDAL-backed object store (Tencent COS) for Vortex" +publish = false +version = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } +authors = { workspace = true } +license = { workspace = true } +keywords = { workspace = true } +include = { workspace = true } +edition = { workspace = true } +rust-version = { workspace = true } +categories = { workspace = 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", +] } +tracing = { workspace = true } +url = { workspace = true } +vortex-utils = { workspace = true } + +[lints] +workspace = true diff --git a/vortex-object-store-opendal/src/lib.rs b/vortex-object-store-opendal/src/lib.rs new file mode 100644 index 00000000000..92caee88092 --- /dev/null +++ b/vortex-object-store-opendal/src/lib.rs @@ -0,0 +1,316 @@ +// SPDX-License-Identifier: Apache-2.0 +// 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. +//! +//! 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`. +//! +//! # Limitations +//! +//! 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 +//! validation must pre-filter their property maps. + +use std::sync::Arc; + +use object_store::ObjectStore; +use object_store_opendal::OpendalStore; +use opendal::Operator; +use opendal::services; +use tracing::warn; +use url::Url; +use vortex_utils::aliases::hash_map::HashMap; + +/// 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`, ...). + UnsupportedScheme(String), + /// A required configuration value (bucket and/or endpoint) was missing. + MissingConfig(&'static str), + /// The OpenDAL builder rejected the provided configuration. + Build(opendal::Error), +} + +impl std::fmt::Display for OpenDALStoreError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + OpenDALStoreError::UnsupportedScheme(s) => { + write!(f, "unsupported OpenDAL scheme: {s}") + } + OpenDALStoreError::MissingConfig(k) => { + write!(f, "missing required OpenDAL store configuration: {k}") + } + OpenDALStoreError::Build(e) => write!(f, "failed to build OpenDAL store: {e}"), + } + } +} + +impl std::error::Error for OpenDALStoreError {} + +impl From for object_store::Error { + fn from(e: OpenDALStoreError) -> Self { + object_store::Error::Generic { + store: "OpenDAL", + source: Box::new(e), + } + } +} + +/// Schemes handled by this crate. +pub const COS_SCHEME: &str = "cos"; + +/// Strongly-typed configuration for building a [`OpendalStore`] 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 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 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))) +} + +/// Build an [`object_store::ObjectStore`] for a `cos://` URL. +/// +/// `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`). +/// +/// Returns [`OpenDALStoreError::UnsupportedScheme`] if `url` does not use `cos://`. +pub fn make_opendal_store( + url: &Url, + properties: &HashMap, +) -> Result, OpenDALStoreError> { + if url.scheme() != COS_SCHEME { + return Err(OpenDALStoreError::UnsupportedScheme( + url.scheme().to_string(), + )); + } + + // 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) +} + +/// Default environment-variable lookup: reads `key` from the process environment. +/// +/// The lookup is factored out so tests can pass a fixed map instead of mutating the global +/// environment, which is unsound when `cargo test` runs tests on multiple threads within one +/// process (the `unsafe std::env::set_var` block became `unsafe` in Rust 2024 for exactly this +/// reason). +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, + properties: &HashMap, + env_lookup: F, +) -> Result +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"), + }) +} + +fn warn_on_unknown_properties(properties: &HashMap) { + const KNOWN: &[&str] = &[ + "bucket", + "endpoint", + "secret_id", + "secret_key", + "root", + "disable_config_load", + ]; + for key in properties.keys() { + if !KNOWN.contains(&key.as_str()) { + warn!("ignoring unknown OpenDAL store property: {key}"); + } + } +} + +fn build_operator(builder: B) -> Result +where + B: opendal::Builder, +{ + let operator: Operator = Operator::new(builder) + .map_err(OpenDALStoreError::Build)? + .finish(); + Ok(operator) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rejects_unknown_scheme() { + let url = Url::parse("s3://bucket/path").unwrap(); + let props = HashMap::new(); + assert!(matches!( + make_opendal_store(&url, &props), + Err(OpenDALStoreError::UnsupportedScheme(_)) + )); + } + + /// 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") + )); + } + + /// When `properties` does not contain `endpoint`, the env-lookup should be consulted. + /// The fixed lookup returns `"https://example.com"`, so the build must succeed. + #[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")) + )); + } +} diff --git a/vortex-python/Cargo.toml b/vortex-python/Cargo.toml index a87548d69a9..3447912861a 100644 --- a/vortex-python/Cargo.toml +++ b/vortex-python/Cargo.toml @@ -28,6 +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"] [dependencies] arrow-array = { workspace = true } @@ -53,8 +56,10 @@ tokio = { workspace = true, features = ["rt-multi-thread"], optional = true } url = { workspace = true } vortex = { workspace = true, features = ["object_store"] } vortex-arrow = { workspace = true } +vortex-object-store-opendal = { path = "../vortex-object-store-opendal", optional = true } 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/python/vortex/_lib/__init__.pyi b/vortex-python/python/vortex/_lib/__init__.pyi index 79f62c4521f..09c9a520406 100644 --- a/vortex-python/python/vortex/_lib/__init__.pyi +++ b/vortex-python/python/vortex/_lib/__init__.pyi @@ -1,3 +1,23 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright the Vortex contributors +class CosStore: + """A Tencent Cloud COS object store, backed by OpenDAL. + + Construct it with explicit configuration and pass it to + ``vortex.io.read_url(url, store=cos_store)`` / + ``vortex.io.write(arrays, path, store=cos_store)``. + + This class is only available when Vortex is built with the ``opendal`` feature. + """ + + def __init__( + self, + bucket: str, + endpoint: str, + *, + secret_id: str | None = None, + secret_key: str | None = None, + root: str | None = None, + disable_config_load: bool = False, + ) -> None: ... diff --git a/vortex-python/python/vortex/_lib/store/__init__.pyi b/vortex-python/python/vortex/_lib/store/__init__.pyi index b394b118b80..95bc141a044 100644 --- a/vortex-python/python/vortex/_lib/store/__init__.pyi +++ b/vortex-python/python/vortex/_lib/store/__init__.pyi @@ -89,6 +89,9 @@ def from_url( # type: ignore[misc] # docstring in pyi file - `gs://bucket/path` -> [`GCSStore`][vortex.store.GCSStore] - `az://account/container/path` -> [`AzureStore`][vortex.store.AzureStore] (also supports `adl`, `azure`, `abfs`, `abfss`) + - `cos://bucket/path` -> OpenDAL-backed Tencent Cloud COS store (requires the + `opendal` feature; configure via environment variables such as + `TENCENTCLOUD_SECRET_ID` / `TENCENTCLOUD_SECRET_KEY` and `COS_ENDPOINT`) - `http://mydomain/path` -> [`HTTPStore`][vortex.store.HTTPStore] - `https://mydomain/path` -> [`HTTPStore`][vortex.store.HTTPStore] diff --git a/vortex-python/python/vortex/store/__init__.py b/vortex-python/python/vortex/store/__init__.py index 043f0bf9ab6..c9127087bf0 100644 --- a/vortex-python/python/vortex/store/__init__.py +++ b/vortex-python/python/vortex/store/__init__.py @@ -16,13 +16,14 @@ AzureStore, ) from ._client import ClientConfig +from ._cos import CosStore from ._gcs import GCSConfig, GCSCredential, GCSCredentialProvider, GCSStore from ._http import HTTPStore from ._local import LocalStore from ._memory import MemoryStore from ._retry import BackoffConfig, RetryConfig -ObjectStore: TypeAlias = AzureStore | GCSStore | HTTPStore | S3Store | LocalStore | MemoryStore +ObjectStore: TypeAlias = AzureStore | CosStore | GCSStore | HTTPStore | S3Store | LocalStore | MemoryStore """All supported ObjectStore implementations.""" @@ -89,6 +90,9 @@ def from_url( # type: ignore[misc] # docstring in pyi file - ``gs://bucket/path`` -> :class:`~vortex.store.GCSStore` - ``az://account/container/path`` -> :class:`~vortex.store.AzureStore` (also supports ``adl``, ``azure``, ``abfs``, ``abfss``) + - ``cos://bucket/path`` -> OpenDAL-backed Tencent Cloud COS store (requires the + ``opendal`` feature; configure via environment variables such as + ``TENCENTCLOUD_SECRET_ID`` / ``TENCENTCLOUD_SECRET_KEY`` and ``COS_ENDPOINT``) - ``http://mydomain/path`` -> :class:`~vortex.store.HTTPStore` - ``https://mydomain/path`` -> :class:`~vortex.store.HTTPStore` @@ -139,6 +143,8 @@ def from_url( # type: ignore[misc] # docstring in pyi file "BackoffConfig", "ClientConfig", "RetryConfig", + # COS (OpenDAL-backed) + "CosStore", # GCS "GCSConfig", "GCSCredential", diff --git a/vortex-python/python/vortex/store/_cos.py b/vortex-python/python/vortex/store/_cos.py new file mode 100644 index 00000000000..666e0c38b39 --- /dev/null +++ b/vortex-python/python/vortex/store/_cos.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +"""Tencent Cloud COS object store, backed by OpenDAL. + +This store is only available when Vortex is built with the ``opendal`` feature. +The class is re-exported from the native extension module; if the feature is +not enabled, instantiating :class:`CosStore` raises :class:`ImportError`. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + # `vortex._lib.__init__.pyi` declares `CosStore` so type checkers see a real + # signature; the runtime branch below provides the same symbol when the + # `opendal` feature is enabled, and a placeholder that raises on + # instantiation when it is not. + from vortex._lib import CosStore +else: + try: + from vortex._lib import CosStore as CosStore + except ImportError: + + class CosStore: + """Placeholder; the real implementation requires the ``opendal`` feature.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + raise ImportError( + "CosStore requires Vortex to be built with the 'opendal' feature; " + + "build with `maturin build --features opendal` " + + "or `maturin develop --features opendal`." + ) + + +__all__ = ["CosStore"] diff --git a/vortex-python/src/io.rs b/vortex-python/src/io.rs index 182ab7fbe74..a9dfa8c8f1a 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + use arrow_array::RecordBatchReader; use arrow_array::ffi_stream::ArrowArrayStreamReader; use async_fs::File; @@ -41,6 +43,8 @@ use crate::install_module; use crate::iter::PyArrayIterator; use crate::object_store::resolve::ResolvedStore; use crate::object_store::resolve::resolve_store; +#[cfg(feature = "opendal")] +use crate::opendal_store::CosStore; use crate::session::session; pub(crate) fn init(py: Python, parent: &Bound) -> PyResult<()> { @@ -62,7 +66,7 @@ pub(crate) fn init(py: Python, parent: &Bound) -> PyResult<()> { /// ---------- /// url : str /// The URL to read from. -/// store : vortex.store.AzureStore | vortex.store.GCSStore | vortex.store.HTTPStore | vortex.store.LocalStore | vortex.store.MemoryStore | vortex.store.S3Store | None +/// store : vortex.store.AzureStore | vortex.store.CosStore | vortex.store.GCSStore | vortex.store.HTTPStore | vortex.store.LocalStore | vortex.store.MemoryStore | vortex.store.S3Store | None /// Pre-configured object store with credentials and settings. /// If provided, uses this store's configuration. /// If None, checks session registry for matching URL pattern. @@ -115,6 +119,7 @@ pub(crate) fn init(py: Python, parent: &Bound) -> PyResult<()> { /// ... secret_access_key="..." /// ... ) /// >>> a = vx.io.read_url("s3://my-bucket/data.vortex", store=store) # doctest: +SKIP + #[pyfunction] #[pyo3(signature = (url, *, store = None, projection = None, row_filter = None, indices = None, row_range = None))] pub fn read_url<'py>( @@ -127,7 +132,7 @@ pub fn read_url<'py>( row_range: Option<(u64, u64)>, ) -> PyVortexResult { let store_arc = if let Some(store_obj) = store { - let py_store: PyObjectStore = store_obj.extract()?; + let py_store: AnyVortexStore = store_obj.extract()?; Some(py_store.into_inner()) } else { None @@ -137,6 +142,46 @@ pub fn read_url<'py>( dataset.to_array_inner(py, projection, row_filter, indices, row_range) } +/// A store object accepted by `read_url` / `write`. +/// +/// This recognizes both the built-in `pyo3-object_store` classes (S3, Azure, GCS, HTTP, +/// Local, Memory) and Vortex's own OpenDAL-backed classes (`CosStore`). +pub(crate) enum AnyVortexStore { + /// A store extracted from one of the built-in `pyo3-object_store` classes. + Builtin(PyObjectStore), + /// Vortex's OpenDAL-backed COS store. + #[cfg(feature = "opendal")] + Cos(CosStore), +} + +impl AnyVortexStore { + /// Consume self and return the underlying `Arc`. + fn into_inner(self) -> Arc { + match self { + AnyVortexStore::Builtin(s) => s.into_inner(), + #[cfg(feature = "opendal")] + AnyVortexStore::Cos(s) => s.to_arc(), + } + } +} + +impl<'py> FromPyObject<'_, 'py> for AnyVortexStore { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, PyAny>) -> PyResult { + if let Ok(builtin) = obj.extract::() { + return Ok(AnyVortexStore::Builtin(builtin)); + } + #[cfg(feature = "opendal")] + if let Ok(cos) = obj.extract::() { + return Ok(AnyVortexStore::Cos(cos)); + } + Err(PyTypeError::new_err( + "Expected an object store instance (S3/Azure/GCS/HTTP/Local/Memory/COS/OSS store)", + )) + } +} + /// Write an array to a Vortex file. /// /// Parameters @@ -148,7 +193,7 @@ pub fn read_url<'py>( /// path : str /// The file path. /// -/// store : vortex.store.AzureStore | vortex.store.GCSStore | vortex.store.HTTPStore | vortex.store.LocalStore | vortex.store.MemoryStore | vortex.store.S3Store | None +/// store : vortex.store.AzureStore | vortex.store.CosStore | vortex.store.GCSStore | vortex.store.HTTPStore | vortex.store.LocalStore | vortex.store.MemoryStore | vortex.store.S3Store | None /// An optional object store configuration to use for writing the output. /// /// Examples @@ -190,7 +235,7 @@ pub fn write( py: Python, iter: PyIntoArrayIterator, path: &str, - store: Option, + store: Option, ) -> PyVortexResult<()> { let session = session(); py.detach(|| { @@ -290,7 +335,7 @@ impl PyVortexWriteOptions { /// path : str /// The file path. /// - /// store : vortex.store.AzureStore | vortex.store.GCSStore | vortex.store.HTTPStore | vortex.store.LocalStore | vortex.store.MemoryStore | vortex.store.S3Store | None + /// store : vortex.store.AzureStore | vortex.store.CosStore | vortex.store.GCSStore | vortex.store.HTTPStore | vortex.store.LocalStore | vortex.store.MemoryStore | vortex.store.S3Store | None /// An optional object store configuration to use for writing the output. /// /// Examples @@ -319,7 +364,7 @@ impl PyVortexWriteOptions { py: Python, iter: PyIntoArrayIterator, path: &str, - store: Option, + store: Option, ) -> PyVortexResult<()> { let session = session(); py.detach(|| { diff --git a/vortex-python/src/lib.rs b/vortex-python/src/lib.rs index 6582a255955..0912426d10b 100644 --- a/vortex-python/src/lib.rs +++ b/vortex-python/src/lib.rs @@ -25,6 +25,8 @@ mod file; mod io; mod iter; mod object_store; +#[cfg(feature = "opendal")] +mod opendal_store; mod python_repr; mod registry; mod runtime; @@ -80,6 +82,8 @@ fn _lib(py: Python, m: &Bound) -> PyResult<()> { file::init(py, m)?; io::init(py, m)?; iter::init(py, m)?; + #[cfg(feature = "opendal")] + opendal_store::init(py, m)?; runtime::init(py, m)?; store::init(py, m)?; registry::init(py, m)?; diff --git a/vortex-python/src/object_store/registry.rs b/vortex-python/src/object_store/registry.rs index c9544e0558e..144a0d0e536 100644 --- a/vortex-python/src/object_store/registry.rs +++ b/vortex-python/src/object_store/registry.rs @@ -19,6 +19,8 @@ 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 { @@ -64,50 +66,88 @@ pub(crate) struct Registry { impl ObjectStoreRegistry for Registry { fn register(&self, url: Url, store: Arc) -> Option> { 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(); - } + 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)); } } - let normalized_env = std::env::vars().map(|(k, v)| (k.to_ascii_lowercase(), v)); + // 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()); - if let Ok((store, path)) = parse_url_opts(to_resolve, normalized_env) { - 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) + } +} - 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 - }); +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 path = path_suffix(to_resolve, depth)?; - return Ok((store, path)); - } + 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 + }); - Err(object_store::Error::Generic { - store: "ObjectStoreRegistry", - source: "URL could not be resolved".into(), - }) + 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) @@ -140,10 +180,27 @@ fn path_suffix(url: &Url, depth: usize) -> Result { 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; @@ -189,4 +246,64 @@ mod tests { 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/opendal_store.rs b/vortex-python/src/opendal_store.rs new file mode 100644 index 00000000000..f0c8d868b79 --- /dev/null +++ b/vortex-python/src/opendal_store.rs @@ -0,0 +1,115 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Python-facing wrapper for the OpenDAL-backed Tencent Cloud COS object store. +//! +//! This lets callers build a concrete store object in Python and pass it to +//! `vortex.io.read_url(store=...)` / `vortex.io.write(..., store=...)` exactly like any +//! of the built-in store classes (S3, Azure, ...). The store is constructed from the same +//! configuration the `cos://` URL registry uses, but materialized eagerly so it +//! can be handed around as a first-class value. +//! +//! This module is only compiled when the `opendal` feature is enabled. + +use std::sync::Arc; + +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use vortex_object_store_opendal::CosConfig; + +/// A Tencent Cloud COS object store, backed by OpenDAL. +/// +/// Construct it with explicit configuration and pass it to +/// ``vortex.io.read_url(url, store=cos_store)`` / ``vortex.io.write(arrays, path, store=cos_store)``. +#[pyclass(name = "CosStore", module = "vortex._lib", frozen, from_py_object)] +#[derive(Clone, Debug)] +pub struct CosStore { + store: Arc, +} + +impl CosStore { + /// Clone the underlying object store as an `Arc`. + pub fn to_arc(&self) -> Arc { + Arc::clone(&self.store) + } +} + +#[pymethods] +impl CosStore { + #[new] + #[pyo3(signature = ( + bucket, + endpoint, + *, + secret_id = None, + secret_key = None, + root = None, + disable_config_load = false, + ))] + fn new( + bucket: String, + endpoint: String, + secret_id: Option, + secret_key: Option, + root: Option, + disable_config_load: bool, + ) -> PyResult { + let config = CosConfig { + bucket, + endpoint, + secret_id, + secret_key, + root, + disable_config_load, + }; + let store = vortex_object_store_opendal::make_cos_store(config) + .map_err(|e| PyValueError::new_err(e.to_string()))?; + Ok(Self { store }) + } +} + +/// Register the OpenDAL-backed store classes on the `vortex._lib` module. +pub(crate) fn init(_py: Python, parent: &Bound) -> PyResult<()> { + parent.add_class::()?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use vortex_object_store_opendal::CosConfig; + + use super::*; + + /// `make_cos_store` (the strongly-typed entry point used by the `CosStore` pyclass) must + /// accept a fully-specified `CosConfig` and yield a usable `Arc`. + #[test] + fn cos_config_builds_object_store() { + let config = CosConfig { + bucket: "my-bucket".to_string(), + endpoint: "https://cos.ap-guangzhou.myqcloud.com".to_string(), + secret_id: Some("AKID".to_string()), + secret_key: Some("secret".to_string()), + root: Some("nested/prefix".to_string()), + disable_config_load: true, + }; + let store = + vortex_object_store_opendal::make_cos_store(config).expect("cos store should build"); + // Sanity-check the result is a non-null `Arc`. + assert!(Arc::strong_count(&store) >= 1); + } + + /// Constructing a `CosConfig` with an empty `bucket` must surface as a `MissingConfig` error + /// rather than silently building an invalid store. + #[test] + fn cos_config_rejects_empty_bucket() { + let result = vortex_object_store_opendal::make_cos_store(CosConfig { + bucket: String::new(), + endpoint: "https://cos.ap-guangzhou.myqcloud.com".to_string(), + ..CosConfig::default() + }); + assert!(matches!( + result, + Err(vortex_object_store_opendal::OpenDALStoreError::MissingConfig("bucket")) + )); + } +}