diff --git a/Cargo.lock b/Cargo.lock index 56b62132af6..a431dc86231 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5999,6 +5999,7 @@ dependencies = [ "opendal-service-hf", "opendal-service-oss", "opendal-service-s3", + "opendal-service-tos", ] [[package]] @@ -6216,6 +6217,23 @@ dependencies = [ "url", ] +[[package]] +name = "opendal-service-tos" +version = "0.57.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f2f7a4c32e5202eb4ac72e76c4b5e30c86ab60762811172f4111103b9d673a1" +dependencies = [ + "bytes", + "http 1.4.1", + "opendal-core", + "quick-xml 0.39.4", + "reqsign-core", + "reqsign-file-read-tokio", + "reqsign-volcengine-tos", + "serde", + "serde_json", +] + [[package]] name = "openssl" version = "0.10.80" @@ -7403,6 +7421,19 @@ dependencies = [ "serde_json", ] +[[package]] +name = "reqsign-volcengine-tos" +version = "3.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9d757602a7ef2b6025c0da77e6d2e23fbdef35930fa466b15ffbf0a3f13acf7" +dependencies = [ + "anyhow", + "http 1.4.1", + "log", + "percent-encoding", + "reqsign-core", +] + [[package]] name = "reqwest" version = "0.12.28" diff --git a/docs/src/guide/object_store.md b/docs/src/guide/object_store.md index 1710e3b5100..182b93c0574 100644 --- a/docs/src/guide/object_store.md +++ b/docs/src/guide/object_store.md @@ -218,3 +218,33 @@ ds = lance.dataset( | `oss_secret_access_key` | Access key secret used for OSS authentication. Optional if credentials are provided by environment. | | `oss_region` | OSS region (for example, `cn-hangzhou`). Optional. | | `oss_security_token` | Security token for temporary credentials (STS). Optional. | + +## Volcengine TOS Configuration + +TOS credentials can be set in the environment variables `TOS_ACCESS_KEY_ID`, +`TOS_SECRET_ACCESS_KEY`, `TOS_ENDPOINT`, `TOS_REGION`, and `TOS_SECURITY_TOKEN`. +Lance also accepts the corresponding `VOLCENGINE_` environment variable prefix. +Alternatively, credentials can be passed as parameters to the `storage_options` +parameter; explicit `storage_options` override environment variables: + +```python +import lance +ds = lance.dataset( + "tos://bucket/path", + storage_options={ + "tos_endpoint": "https://tos-cn-beijing.volces.com", + "tos_region": "cn-beijing", + "tos_access_key_id": "my-access-key", + "tos_secret_access_key": "my-secret-key", + "tos_security_token": "my-session-token", + } +) +``` + +| Key | Description | +|-----|-------------| +| `tos_endpoint` | TOS endpoint. Required (for example, `https://tos-cn-beijing.volces.com`). | +| `tos_region` | TOS signing region (for example, `cn-beijing`). Optional. | +| `tos_access_key_id` | Access key ID used for TOS authentication. Optional if credentials are provided by environment. | +| `tos_secret_access_key` | Secret access key used for TOS authentication. Optional if credentials are provided by environment. | +| `tos_security_token` | Security token for temporary credentials. Optional. | diff --git a/python/Cargo.lock b/python/Cargo.lock index 162a221907f..2fe2e44d75c 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -5334,6 +5334,7 @@ dependencies = [ "opendal-service-hf", "opendal-service-oss", "opendal-service-s3", + "opendal-service-tos", ] [[package]] @@ -5551,6 +5552,23 @@ dependencies = [ "url", ] +[[package]] +name = "opendal-service-tos" +version = "0.57.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f2f7a4c32e5202eb4ac72e76c4b5e30c86ab60762811172f4111103b9d673a1" +dependencies = [ + "bytes", + "http 1.4.1", + "opendal-core", + "quick-xml 0.39.4", + "reqsign-core", + "reqsign-file-read-tokio", + "reqsign-volcengine-tos", + "serde", + "serde_json", +] + [[package]] name = "openssl-probe" version = "0.2.1" @@ -6692,6 +6710,19 @@ dependencies = [ "serde_json", ] +[[package]] +name = "reqsign-volcengine-tos" +version = "3.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9d757602a7ef2b6025c0da77e6d2e23fbdef35930fa466b15ffbf0a3f13acf7" +dependencies = [ + "anyhow", + "http 1.4.1", + "log", + "percent-encoding", + "reqsign-core", +] + [[package]] name = "reqwest" version = "0.12.28" diff --git a/rust/lance-io/Cargo.toml b/rust/lance-io/Cargo.toml index 9b89489b0f7..ec099d66889 100644 --- a/rust/lance-io/Cargo.toml +++ b/rust/lance-io/Cargo.toml @@ -76,6 +76,8 @@ azure = ["object_store/azure", "dep:opendal", "opendal/services-azblob", "openda oss = ["dep:opendal", "opendal/services-oss", "dep:object_store_opendal"] tencent = ["dep:opendal", "opendal/services-cos", "dep:object_store_opendal"] huggingface = ["dep:opendal", "opendal/services-huggingface", "dep:object_store_opendal"] +tos = ["dep:opendal", "opendal/services-tos", "dep:object_store_opendal"] +tos-test = ["tos"] test-util = [] [lints] diff --git a/rust/lance-io/src/object_store.rs b/rust/lance-io/src/object_store.rs index 698378c5b9e..b1ba3d45d8a 100644 --- a/rust/lance-io/src/object_store.rs +++ b/rust/lance-io/src/object_store.rs @@ -36,7 +36,7 @@ use super::local::LocalObjectReader; use crate::uring::{UringCurrentThreadReader, UringReader}; #[cfg(any(feature = "aws", feature = "azure", feature = "gcp"))] pub(crate) mod dynamic_credentials; -#[cfg(any(feature = "oss", feature = "huggingface"))] +#[cfg(any(feature = "oss", feature = "huggingface", feature = "tos"))] pub(crate) mod dynamic_opendal; mod list_retry; pub mod providers; @@ -61,7 +61,15 @@ pub const DEFAULT_LOCAL_IO_PARALLELISM: usize = 8; pub const DEFAULT_CLOUD_IO_PARALLELISM: usize = 64; const DEFAULT_LOCAL_BLOCK_SIZE: usize = 4 * 1024; // 4KB block size -#[cfg(any(feature = "aws", feature = "gcp", feature = "azure"))] +#[cfg(any( + feature = "aws", + feature = "gcp", + feature = "azure", + feature = "oss", + feature = "tencent", + feature = "huggingface", + feature = "tos", +))] const DEFAULT_CLOUD_BLOCK_SIZE: usize = 64 * 1024; // 64KB block size pub static DEFAULT_MAX_IOP_SIZE: std::sync::LazyLock = std::sync::LazyLock::new(|| { diff --git a/rust/lance-io/src/object_store/providers.rs b/rust/lance-io/src/object_store/providers.rs index 20fa251a0c5..aafe665cfe1 100644 --- a/rust/lance-io/src/object_store/providers.rs +++ b/rust/lance-io/src/object_store/providers.rs @@ -33,6 +33,8 @@ pub mod oss; pub mod shared_memory; #[cfg(feature = "tencent")] pub mod tencent; +#[cfg(feature = "tos")] +pub mod tos; #[async_trait::async_trait] pub trait ObjectStoreProvider: std::fmt::Debug + Sync + Send { @@ -95,6 +97,7 @@ pub struct ObjectStoreRegistryStats { /// - `s3+ddb`: An S3 object store with DynamoDB for metadata. /// - `az`: An Azure Blob Storage object store. /// - `gs`: A Google Cloud Storage object store. +/// - `tos`: A Volcengine TOS object store. /// /// Use [`Self::empty()`] to create an empty registry, with no providers registered. /// @@ -330,6 +333,8 @@ impl Default for ObjectStoreRegistry { providers.insert("cos".into(), Arc::new(tencent::TencentStoreProvider)); #[cfg(feature = "huggingface")] providers.insert("hf".into(), Arc::new(huggingface::HuggingfaceStoreProvider)); + #[cfg(feature = "tos")] + providers.insert("tos".into(), Arc::new(tos::TosStoreProvider)); Self { providers: RwLock::new(providers), active_stores: RwLock::new(HashMap::new()), diff --git a/rust/lance-io/src/object_store/providers/tos.rs b/rust/lance-io/src/object_store/providers/tos.rs new file mode 100644 index 00000000000..923186484c6 --- /dev/null +++ b/rust/lance-io/src/object_store/providers/tos.rs @@ -0,0 +1,301 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::collections::HashMap; +use std::sync::Arc; + +use object_store::ObjectStore as OSObjectStore; +use object_store_opendal::OpendalStore; +use opendal::{Operator, services::Tos}; +use url::Url; + +use crate::object_store::dynamic_opendal::DynamicOpenDalStore; +use crate::object_store::{ + DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, + ObjectStoreParams, ObjectStoreProvider, StorageOptions, +}; +use lance_core::error::{Error, Result}; + +#[derive(Default, Debug)] +pub struct TosStoreProvider; + +impl TosStoreProvider { + fn tos_env_options_from_iter(vars: I) -> HashMap + where + I: IntoIterator, + K: Into, + V: Into, + { + let vars = vars + .into_iter() + .map(|(key, value)| (key.into(), value.into())) + .collect::>(); + let mut config_map = HashMap::new(); + + for prefix in ["VOLCENGINE_", "TOS_"] { + for (key, value) in &vars { + if let Some(stripped_key) = key.strip_prefix(prefix) { + config_map.insert(stripped_key.to_ascii_lowercase(), value.clone()); + } + } + } + + config_map + } + + fn base_tos_options( + base_path: &Url, + storage_options: &StorageOptions, + ) -> Result> { + let bucket = base_path + .host_str() + .ok_or_else(|| Error::invalid_input("TOS URL must contain bucket name"))? + .to_string(); + + let prefix = base_path.path().trim_start_matches('/').to_string(); + + let mut config_map = Self::tos_env_options_from_iter(std::env::vars()); + + config_map.extend(storage_options.0.clone()); + + config_map.insert("bucket".to_string(), bucket); + if prefix.is_empty() { + config_map.remove("root"); + } else { + config_map.insert("root".to_string(), "/".to_string()); + } + + Ok(config_map) + } + + /// Normalize TOS storage options, resolving aliases for well-known keys + /// while passing through all other options so that OpenDAL can use them. + fn normalize_tos_config(options: &HashMap) -> Result> { + let mut config_map = options.clone(); + + let alias_groups: &[(&str, &[&str])] = &[ + ("endpoint", &["tos_endpoint"]), + ("region", &["tos_region"]), + ("access_key_id", &["tos_access_key_id"]), + ("secret_access_key", &["tos_secret_access_key"]), + ("security_token", &["tos_security_token"]), + ]; + + for (canonical, aliases) in alias_groups { + for alias in *aliases { + if let Some(value) = config_map.remove(*alias) { + config_map.insert(canonical.to_string(), value); + break; + } + } + } + + if !config_map.contains_key("endpoint") { + return Err(Error::invalid_input( + "TOS endpoint is required. Please provide 'tos_endpoint' in storage options or set TOS_ENDPOINT environment variable", + )); + } + + Ok(config_map) + } + + fn build_tos_store(config_map: HashMap) -> Result { + let operator = Operator::from_iter::(config_map) + .map_err(|e| Error::invalid_input(format!("Failed to create TOS operator: {:?}", e)))? + .finish(); + + Ok(OpendalStore::new(operator)) + } +} + +#[async_trait::async_trait] +impl ObjectStoreProvider for TosStoreProvider { + async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result { + let block_size = params.block_size.unwrap_or(DEFAULT_CLOUD_BLOCK_SIZE); + let storage_options = StorageOptions(params.storage_options().cloned().unwrap_or_default()); + + let base_options = Self::base_tos_options(&base_path, &storage_options)?; + let accessor = params.get_accessor(); + + let inner: Arc = + if let Some(accessor) = accessor.filter(|a| a.has_provider()) { + Arc::new( + DynamicOpenDalStore::new( + format!("tos:{}", base_path), + base_options, + accessor, + Self::normalize_tos_config, + Self::build_tos_store, + ) + .with_protected_keys(["bucket", "root"]), + ) + } else { + Arc::new(Self::build_tos_store(Self::normalize_tos_config( + &base_options, + )?)?) + }; + + let mut url = base_path; + if !url.path().ends_with('/') { + url.set_path(&format!("{}/", url.path())); + } + + Ok(ObjectStore { + scheme: "tos".to_string(), + inner, + block_size, + max_iop_size: *DEFAULT_MAX_IOP_SIZE, + use_constant_size_upload_parts: params.use_constant_size_upload_parts, + list_is_lexically_ordered: params.list_is_lexically_ordered.unwrap_or(true), + io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, + download_retry_count: storage_options.download_retry_count(), + io_tracker: Default::default(), + store_prefix: self.calculate_object_store_prefix(&url, params.storage_options())?, + }) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use super::TosStoreProvider; + use crate::object_store::dynamic_opendal::DynamicOpenDalStore; + use crate::object_store::test_utils::StaticMockStorageOptionsProvider; + use crate::object_store::{ObjectStoreProvider, StorageOptionsAccessor}; + use url::Url; + + #[test] + fn test_tos_store_path() { + let provider = TosStoreProvider; + + let url = Url::parse("tos://bucket/path/to/file").unwrap(); + let path = provider.extract_path(&url).unwrap(); + let expected_path = object_store::path::Path::from("path/to/file"); + assert_eq!(path, expected_path); + } + + #[test] + fn test_tos_env_options_normalize_supported_prefixes() { + let config = TosStoreProvider::tos_env_options_from_iter([ + ("VOLCENGINE_ENDPOINT", "https://tos-cn-beijing.volces.com"), + ("TOS_ACCESS_KEY_ID", "tos-akid"), + ("TOS_SECRET_ACCESS_KEY", "tos-secret"), + ]); + + assert_eq!( + config.get("endpoint").unwrap(), + "https://tos-cn-beijing.volces.com" + ); + assert_eq!(config.get("access_key_id").unwrap(), "tos-akid"); + assert_eq!(config.get("secret_access_key").unwrap(), "tos-secret"); + } + + #[test] + fn test_tos_alias_options_override_canonical_env_options() { + let config = TosStoreProvider::normalize_tos_config(&HashMap::from([ + ( + "endpoint".to_string(), + "https://env.example.com".to_string(), + ), + ( + "tos_endpoint".to_string(), + "https://user.example.com".to_string(), + ), + ("region".to_string(), "env-region".to_string()), + ("tos_region".to_string(), "user-region".to_string()), + ("access_key_id".to_string(), "env-akid".to_string()), + ("tos_access_key_id".to_string(), "user-akid".to_string()), + ("secret_access_key".to_string(), "env-secret".to_string()), + ( + "tos_secret_access_key".to_string(), + "user-secret".to_string(), + ), + ("security_token".to_string(), "env-token".to_string()), + ("tos_security_token".to_string(), "user-token".to_string()), + ("bucket".to_string(), "bucket".to_string()), + ])) + .unwrap(); + + assert_eq!(config.get("endpoint").unwrap(), "https://user.example.com"); + assert_eq!(config.get("region").unwrap(), "user-region"); + assert_eq!(config.get("access_key_id").unwrap(), "user-akid"); + assert_eq!(config.get("secret_access_key").unwrap(), "user-secret"); + assert_eq!(config.get("security_token").unwrap(), "user-token"); + assert!(!config.contains_key("tos_endpoint")); + assert!(!config.contains_key("tos_secret_access_key")); + assert!(!config.contains_key("tos_security_token")); + } + + #[test] + fn test_tos_url_bucket_and_root_are_authoritative() { + let storage_options = crate::object_store::StorageOptions(HashMap::from([ + ( + "tos_endpoint".to_string(), + "https://tos-cn-beijing.volces.com".to_string(), + ), + ("bucket".to_string(), "storage-options-bucket".to_string()), + ("root".to_string(), "/storage-options-root".to_string()), + ])); + let base_options = TosStoreProvider::base_tos_options( + &Url::parse("tos://url-bucket/path").unwrap(), + &storage_options, + ) + .unwrap(); + let config = TosStoreProvider::normalize_tos_config(&base_options).unwrap(); + + assert_eq!(config.get("bucket").unwrap(), "url-bucket"); + assert_eq!(config.get("root").unwrap(), "/"); + + let base_options = TosStoreProvider::base_tos_options( + &Url::parse("tos://url-bucket").unwrap(), + &storage_options, + ) + .unwrap(); + let config = TosStoreProvider::normalize_tos_config(&base_options).unwrap(); + + assert_eq!(config.get("bucket").unwrap(), "url-bucket"); + assert!(!config.contains_key("root")); + } + + #[tokio::test] + async fn test_dynamic_opendal_tos_store_uses_provider_credentials() { + let accessor = Arc::new(StorageOptionsAccessor::with_provider(Arc::new( + StaticMockStorageOptionsProvider { + options: HashMap::from([ + ( + "tos_endpoint".to_string(), + "https://tos-cn-beijing.volces.com".to_string(), + ), + ("tos_region".to_string(), "cn-beijing".to_string()), + ("tos_access_key_id".to_string(), "akid".to_string()), + ("tos_secret_access_key".to_string(), "secret".to_string()), + ("tos_security_token".to_string(), "token".to_string()), + ]), + }, + ))); + + let base_options = TosStoreProvider::base_tos_options( + &Url::parse("tos://url-bucket/path").unwrap(), + &crate::object_store::StorageOptions(HashMap::new()), + ) + .unwrap(); + + let store = DynamicOpenDalStore::new( + "tos", + base_options, + accessor, + TosStoreProvider::normalize_tos_config, + TosStoreProvider::build_tos_store, + ) + .with_protected_keys(["bucket", "root"]); + + let current_store = store + .current_store() + .await + .expect("dynamic OpenDAL TOS store should build"); + + assert!(current_store.to_string().contains("Opendal")); + } +} diff --git a/rust/lance-io/tests/tos_integration.rs b/rust/lance-io/tests/tos_integration.rs new file mode 100644 index 00000000000..97eca30090e --- /dev/null +++ b/rust/lance-io/tests/tos_integration.rs @@ -0,0 +1,81 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors +//! These integration tests can only be run against a real Volcengine TOS bucket. + +#![cfg(feature = "tos-test")] + +use futures::TryStreamExt; +use lance_io::object_store::ObjectStore; +use object_store::ObjectStoreExt; +use object_store::path::Path; +use tokio::io::AsyncWriteExt; + +fn tos_bucket() -> String { + std::env::var("TOS_BUCKET").expect("TOS_BUCKET must be set") +} + +async fn delete_prefix(store: &ObjectStore, prefix: &str) { + let prefix_path = Path::from(prefix); + let locations = store + .inner + .list(Some(&prefix_path)) + .map_ok(|meta| meta.location) + .try_collect::>() + .await + .unwrap_or_default(); + + for location in locations { + let _ = store.inner.delete(&location).await; + } +} + +#[ignore = "Must be run manually on Volcengine TOS"] +#[tokio::test] +async fn test_tos_write_read_list_delete() { + let prefix = format!("lance-tos-{}-{}", std::process::id(), rand::random::()); + let bucket = tos_bucket(); + let (store, base_path) = ObjectStore::from_uri(&format!("tos://{bucket}/{prefix}")) + .await + .unwrap(); + assert_eq!(base_path, Path::from(prefix.as_str())); + + let path = Path::from(format!("{prefix}/small.txt")); + delete_prefix(&store, &prefix).await; + + let result: Result<(), Box> = async { + let mut writer = store.create(&path).await?; + writer.write_all(b"hello").await?; + writer.write_all(b" tos").await?; + writer.shutdown().await?; + + let meta = store.inner.head(&path).await?; + if meta.size != 9 { + return Err(format!("expected object size 9, got {}", meta.size).into()); + } + + let data = store.inner.get(&path).await?.bytes().await?; + if data.as_ref() != b"hello tos" { + return Err("downloaded TOS object content did not match".into()); + } + + let listed = store + .inner + .list(Some(&Path::from(prefix.as_str()))) + .try_collect::>() + .await?; + if !listed.iter().any(|meta| meta.location == path) { + return Err("uploaded TOS object was not returned by list".into()); + } + + store.inner.delete(&path).await?; + if store.exists(&path).await? { + return Err("deleted TOS object still exists".into()); + } + + Ok(()) + } + .await; + + delete_prefix(&store, &prefix).await; + result.unwrap(); +} diff --git a/rust/lance/Cargo.toml b/rust/lance/Cargo.toml index 5bfaddd9d12..fbffcad1a9f 100644 --- a/rust/lance/Cargo.toml +++ b/rust/lance/Cargo.toml @@ -139,7 +139,7 @@ parquet = { version = "58", default-features = false, features = ["arrow", "asyn reqwest = { version = "0.12", default-features = false, features = ["rustls-tls", "json"] } [features] -default = ["aws", "azure", "gcp", "oss", "huggingface", "tencent", "geo"] +default = ["aws", "azure", "gcp", "oss", "huggingface", "tencent", "tos", "geo"] fp16kernels = ["lance-linalg/fp16kernels"] # Prevent dynamic linking of lzma, which comes from datafusion cli = ["dep:clap", "lzma-sys/static"] @@ -159,6 +159,7 @@ azure = ["lance-io/azure"] oss = ["lance-io/oss"] tencent = ["lance-io/tencent"] huggingface = ["lance-io/huggingface"] +tos = ["lance-io/tos"] geo = ["lance-datafusion/geo", "lance-index/geo"] # Enable slow integration tests (disabled by default in CI) slow_tests = []