From f643e1257b8545cd849a55705f04f54be5eb8efc Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Tue, 22 Sep 2026 16:53:24 +0800 Subject: [PATCH] fix: use relative object keys for local snapshot and test stores Signed-off-by: jeremyhi --- src/cli/src/metadata/snapshot.rs | 139 +++++++++++------- src/common/datasource/src/file_format/csv.rs | 19 ++- src/common/datasource/src/file_format/json.rs | 11 +- .../datasource/src/file_format/tests.rs | 40 +++-- src/common/datasource/src/test_util.rs | 33 +++-- src/common/datasource/src/tests.rs | 14 +- 6 files changed, 151 insertions(+), 105 deletions(-) diff --git a/src/cli/src/metadata/snapshot.rs b/src/cli/src/metadata/snapshot.rs index 59b2599ad9c..2ba95225768 100644 --- a/src/cli/src/metadata/snapshot.rs +++ b/src/cli/src/metadata/snapshot.rs @@ -12,14 +12,18 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::path::Path; + use async_trait::async_trait; use clap::{Parser, Subcommand}; use common_error::ext::BoxedError; use common_meta::snapshot::MetadataSnapshotManager; -use object_store::{ObjectStore, services}; +use object_store::ObjectStore; +use snafu::OptionExt; use crate::Tool; use crate::common::{ObjectStoreConfig, StoreConfig, new_fs_object_store}; +use crate::error::InvalidFilePathSnafu; use crate::utils::resolve_relative_path_with_current_dir; /// Subcommand for metadata snapshot operations, including saving snapshots, restoring from snapshots, and viewing snapshot information. @@ -270,24 +274,51 @@ fn build_object_store_and_resolve_file_path( fs_root: &str, file_path: &str, ) -> Result<(ObjectStore, String), BoxedError> { - let object_store = object_store.build().map_err(BoxedError::new)?; - let object_store = match object_store { - Some(object_store) => object_store, - None => new_fs_object_store(fs_root)?, - }; + if let Some(object_store) = object_store.build().map_err(BoxedError::new)? { + return Ok((object_store, file_path.to_string())); + } - let file_path = if object_store.info().scheme() == services::FS_SCHEME { - resolve_relative_path_with_current_dir(file_path).map_err(BoxedError::new)? - } else { - file_path.to_string() - }; + if fs_root != "/" { + let root = resolve_relative_path_with_current_dir(fs_root).map_err(BoxedError::new)?; + let path = Path::new(file_path); + let key = if path.is_absolute() { + path.strip_prefix(&root).map_err(|_| { + BoxedError::new( + InvalidFilePathSnafu { + msg: format!("Snapshot path is outside root {root}: {file_path}"), + } + .build(), + ) + })? + } else { + path + }; + let key = key.to_string_lossy().into_owned(); + #[cfg(windows)] + let key = key.replace('\\', "/"); + return Ok((new_fs_object_store(&root)?, key)); + } - Ok((object_store, file_path)) + let path = resolve_relative_path_with_current_dir(file_path).map_err(BoxedError::new)?; + let path = Path::new(&path); + let parent = path + .parent() + .with_context(|| InvalidFilePathSnafu { + msg: format!("Snapshot path has no parent: {}", path.display()), + }) + .map_err(BoxedError::new)?; + let file_name = path + .file_name() + .with_context(|| InvalidFilePathSnafu { + msg: format!("Snapshot path has no file name: {}", path.display()), + }) + .map_err(BoxedError::new)?; + let object_store = new_fs_object_store(&parent.to_string_lossy())?; + Ok((object_store, file_name.to_string_lossy().into_owned())) } #[cfg(test)] mod tests { - use std::env; use std::path::Path; use std::sync::Arc; use std::time::Duration; @@ -311,44 +342,52 @@ mod tests { #[tokio::test] async fn test_cmd_resolve_file_path() { - common_telemetry::init_default_ut_logging(); - let cmd = RestoreCommand::parse_from([ - "", - "--file_name", - "metadata_snapshot.metadata.fb", - "--backend", - "memory-store", - "--store-addrs", - "memory://", - ]); - let tool = cmd.build().await.unwrap(); - let current_dir = env::current_dir().unwrap(); - let file_path = current_dir.join("metadata_snapshot.metadata.fb"); - assert_eq!(tool.file_path, file_path.to_string_lossy().to_string()); + let name = "metadata_snapshot.metadata.fb"; + let current_dir = std::env::current_dir().unwrap(); + let absolute = current_dir.join(name); + for file_path in [name, absolute.to_str().unwrap()] { + let cmd = RestoreCommand::parse_from([ + "", + "--file_name", + file_path, + "--backend", + "memory-store", + "--store-addrs", + "memory://", + ]); + let tool = cmd.build().await.unwrap(); + assert_eq!(tool.file_path, name); + } - let cmd = RestoreCommand::parse_from([ - "", - "--file_name", - "metadata_snapshot.metadata.fb", - "--backend", - "memory-store", - "--store-addrs", - "memory://", - ]); - let tool = cmd.build().await.unwrap(); - assert_eq!(tool.file_path, file_path.to_string_lossy().to_string()); - - let cmd = RestoreCommand::parse_from([ - "", - "--file_name", - "metadata_snapshot.metadata.fb", - "--backend", - "memory-store", - "--store-addrs", - "memory://", - ]); - let tool = cmd.build().await.unwrap(); - assert_eq!(tool.file_path, file_path.to_string_lossy().to_string()); + let dir = tempfile::tempdir().unwrap(); + for file_path in [ + "backup/snapshot.fb".to_string(), + dir.path() + .join("backup/snapshot.fb") + .to_string_lossy() + .into_owned(), + ] { + let (store, key) = build_object_store_and_resolve_file_path( + ObjectStoreConfig::default(), + dir.path().to_str().unwrap(), + &file_path, + ) + .unwrap(); + assert_eq!(key, "backup/snapshot.fb"); + store.write(&key, "snapshot").await.unwrap(); + assert_eq!( + std::fs::read(dir.path().join("backup/snapshot.fb")).unwrap(), + b"snapshot" + ); + } + assert!( + build_object_store_and_resolve_file_path( + ObjectStoreConfig::default(), + dir.path().to_str().unwrap(), + absolute.to_str().unwrap(), + ) + .is_err() + ); } async fn setup_backup_file(object_store: ObjectStore, file_path: &str) { diff --git a/src/common/datasource/src/file_format/csv.rs b/src/common/datasource/src/file_format/csv.rs index d016d4eabe9..d4cfb45088c 100644 --- a/src/common/datasource/src/file_format/csv.rs +++ b/src/common/datasource/src/file_format/csv.rs @@ -647,16 +647,15 @@ mod tests { let compressed_file_name = format!("test_compressed_csv.{}", compression_type.file_extension()); let compressed_file_path = temp_dir.path().join(&compressed_file_name); - let compressed_file_path_str = compressed_file_path.to_str().unwrap(); // Create a simple file store for testing - let store = test_store("/"); + let store = test_store(temp_dir.path().to_str().unwrap()); // Export CSV with compression let rows = stream_to_csv( Box::pin(DfRecordBatchStreamAdapter::new(recordbatches.as_stream())), store, - compressed_file_path_str, + &compressed_file_name, 1024, 1, &format, @@ -708,13 +707,13 @@ mod tests { } // Verify the compressed file can be decompressed and content matches original data - let store = test_store("/"); + let store = test_store(temp_dir.path().to_str().unwrap()); let schema = Arc::new( CsvFormat { compression_type, ..Default::default() } - .infer_schema(&store, compressed_file_path_str) + .infer_schema(&store, &compressed_file_name) .await .unwrap(), ); @@ -722,7 +721,7 @@ mod tests { let stream = file_to_stream( &store, - compressed_file_path_str, + &compressed_file_name, csv_source.clone(), None, compression_type, @@ -761,13 +760,13 @@ mod tests { ) .unwrap(); - let store = test_store("/"); + let store = test_store(temp_dir.path().to_str().unwrap()); let schema = Arc::new(arrow_schema::Schema::new(vec![ Field::new("id", DataType::UInt32, false), Field::new("name", DataType::Utf8, false), Field::new("value", DataType::Float64, false), ])); - let path = csv_file_path.to_str().unwrap(); + let path = "input.csv"; let stream = tolerant_csv_stream(&store, path, schema, vec![0, 1, 2], &CsvFormat::default()) @@ -793,13 +792,13 @@ mod tests { let csv_file_path = temp_dir.path().join("input.csv"); std::fs::write(&csv_file_path, "id,name,value\n1,Alice,10.5\n2,Bob\n").unwrap(); - let store = test_store("/"); + let store = test_store(temp_dir.path().to_str().unwrap()); let schema = Arc::new(arrow_schema::Schema::new(vec![ Field::new("id", DataType::UInt32, false), Field::new("name", DataType::Utf8, false), Field::new("value", DataType::Float64, false), ])); - let path = csv_file_path.to_str().unwrap(); + let path = "input.csv"; let stream = tolerant_csv_stream(&store, path, schema, vec![0, 1, 2], &CsvFormat::default()) diff --git a/src/common/datasource/src/file_format/json.rs b/src/common/datasource/src/file_format/json.rs index 3c24a4b7e27..23ce461fcb2 100644 --- a/src/common/datasource/src/file_format/json.rs +++ b/src/common/datasource/src/file_format/json.rs @@ -313,16 +313,15 @@ mod tests { let compressed_file_name = format!("test_compressed_json.{}", compression_type.file_extension()); let compressed_file_path = temp_dir.path().join(&compressed_file_name); - let compressed_file_path_str = compressed_file_path.to_str().unwrap(); // Create a simple file store for testing - let store = test_store("/"); + let store = test_store(temp_dir.path().to_str().unwrap()); // Export JSON with compression let rows = stream_to_json( Box::pin(DfRecordBatchStreamAdapter::new(recordbatches.as_stream())), store, - compressed_file_path_str, + &compressed_file_name, 1024, 1, &format, @@ -374,13 +373,13 @@ mod tests { } // Verify the compressed file can be decompressed and content matches original data - let store = test_store("/"); + let store = test_store(temp_dir.path().to_str().unwrap()); let schema = Arc::new( JsonFormat { compression_type, ..Default::default() } - .infer_schema(&store, compressed_file_path_str) + .infer_schema(&store, &compressed_file_name) .await .unwrap(), ); @@ -388,7 +387,7 @@ mod tests { let stream = file_to_stream( &store, - compressed_file_path_str, + &compressed_file_name, json_source.clone(), None, compression_type, diff --git a/src/common/datasource/src/file_format/tests.rs b/src/common/datasource/src/file_format/tests.rs index 27b86e1b114..f227318be3e 100644 --- a/src/common/datasource/src/file_format/tests.rs +++ b/src/common/datasource/src/file_format/tests.rs @@ -66,13 +66,15 @@ impl Test<'_> { #[tokio::test] async fn test_json_opener() { - let store = test_store("/"); + let store = test_store( + find_workspace_path("src/common/datasource/tests") + .to_str() + .unwrap(), + ); let schema = basic_schema_with_time_format(); let file_source = Arc::new(JsonSource::new(schema)).with_batch_size(test_util::TEST_BATCH_SIZE); - let path = &find_workspace_path("/src/common/datasource/tests/json/basic.json") - .display() - .to_string(); + let path = "json/basic.json"; let tests = [ Test { @@ -108,11 +110,13 @@ async fn test_json_opener() { #[tokio::test] async fn test_csv_opener() { - let store = test_store("/"); + let store = test_store( + find_workspace_path("src/common/datasource/tests") + .to_str() + .unwrap(), + ); let schema = basic_schema_with_time_format(); - let path = &find_workspace_path("/src/common/datasource/tests/csv/basic.csv") - .display() - .to_string(); + let path = "csv/basic.csv"; let file_source = CsvSource::new(schema).with_batch_size(test_util::TEST_BATCH_SIZE); @@ -150,13 +154,15 @@ async fn test_csv_opener() { #[tokio::test(flavor = "multi_thread")] async fn test_parquet_exec() { - let store = test_store("/"); + let store = test_store( + find_workspace_path("src/common/datasource/tests") + .to_str() + .unwrap(), + ); let schema = test_basic_schema(); - let path = &find_workspace_path("/src/common/datasource/tests/parquet/basic.parquet") - .display() - .to_string(); + let path = "parquet/basic.parquet"; let parquet_source = ParquetSource::new(schema) .with_parquet_file_reader_factory(Arc::new(DefaultParquetFileReaderFactory::new(store))); @@ -191,11 +197,13 @@ async fn test_parquet_exec() { #[tokio::test] async fn test_orc_opener() { - let path = &find_workspace_path("/src/common/datasource/tests/orc/test.orc") - .display() - .to_string(); + let path = "orc/test.orc"; - let store = test_store("/"); + let store = test_store( + find_workspace_path("src/common/datasource/tests") + .to_str() + .unwrap(), + ); let schema = Arc::new(OrcFormat.infer_schema(&store, path).await.unwrap()); let file_source = Arc::new(OrcSource::new(schema.into())); diff --git a/src/common/datasource/src/test_util.rs b/src/common/datasource/src/test_util.rs index 4cefda88890..de5945d3eef 100644 --- a/src/common/datasource/src/test_util.rs +++ b/src/common/datasource/src/test_util.rs @@ -15,6 +15,7 @@ use std::sync::Arc; use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; +use common_test_util::find_workspace_path; use common_test_util::temp_dir::{TempDir, create_temp_dir}; use datafusion::datasource::file_format::file_compression_type::FileCompressionType; use datafusion::datasource::listing::PartitionedFile; @@ -56,8 +57,8 @@ pub fn test_store(root: &str) -> ObjectStore { pub fn test_tmp_store(root: &str) -> (ObjectStore, TempDir) { let dir = create_temp_dir(root); - let builder = Fs::default(); - (ObjectStore::new(builder.root("/")).unwrap(), dir) + let store = test_store(dir.path().to_str().unwrap()); + (store, dir) } pub fn test_basic_schema() -> SchemaRef { @@ -95,7 +96,11 @@ pub(crate) fn scan_config( } pub async fn setup_stream_to_json_test(origin_path: &str, threshold: impl Fn(usize) -> usize) { - let store = test_store("/"); + let store = test_store( + find_workspace_path("src/common/datasource/tests") + .to_str() + .unwrap(), + ); let schema = basic_schema_with_time_format(); @@ -117,9 +122,9 @@ pub async fn setup_stream_to_json_test(origin_path: &str, threshold: impl Fn(usi .build() .unwrap(); - let (tmp_store, dir) = test_tmp_store("test_stream_to_json"); + let (tmp_store, _dir) = test_tmp_store("test_stream_to_json"); - let output_path = format!("{}/{}", dir.path().display(), "output"); + let output_path = "output"; let json_format = JsonFormat::default(); @@ -127,7 +132,7 @@ pub async fn setup_stream_to_json_test(origin_path: &str, threshold: impl Fn(usi stream_to_json( Box::pin(stream), tmp_store.clone(), - &output_path, + output_path, threshold(size), 8, &json_format, @@ -136,7 +141,7 @@ pub async fn setup_stream_to_json_test(origin_path: &str, threshold: impl Fn(usi .is_ok() ); - let written = tmp_store.read(&output_path).await.unwrap(); + let written = tmp_store.read(output_path).await.unwrap(); let origin = store.read(origin_path).await.unwrap(); assert_eq_lines(written.to_vec(), origin.to_vec()); } @@ -146,7 +151,11 @@ pub async fn setup_stream_to_csv_test( format_path: &str, threshold: impl Fn(usize) -> usize, ) { - let store = test_store("/"); + let store = test_store( + find_workspace_path("src/common/datasource/tests") + .to_str() + .unwrap(), + ); let schema = basic_schema_with_time_format(); @@ -168,9 +177,9 @@ pub async fn setup_stream_to_csv_test( .build() .unwrap(); - let (tmp_store, dir) = test_tmp_store("test_stream_to_csv"); + let (tmp_store, _dir) = test_tmp_store("test_stream_to_csv"); - let output_path = format!("{}/{}", dir.path().display(), "output"); + let output_path = "output"; let csv_format = CsvFormat { timestamp_format: Some("%m-%d-%Y".to_string()), @@ -183,7 +192,7 @@ pub async fn setup_stream_to_csv_test( stream_to_csv( Box::pin(stream), tmp_store.clone(), - &output_path, + output_path, threshold(size), 8, &csv_format, @@ -192,7 +201,7 @@ pub async fn setup_stream_to_csv_test( .is_ok() ); - let written = tmp_store.read(&output_path).await.unwrap(); + let written = tmp_store.read(output_path).await.unwrap(); let format_expect = store.read(format_path).await.unwrap(); assert_eq_lines(written.to_vec(), format_expect.to_vec()); } diff --git a/src/common/datasource/src/tests.rs b/src/common/datasource/src/tests.rs index 8be9add7c88..0a63fca3da1 100644 --- a/src/common/datasource/src/tests.rs +++ b/src/common/datasource/src/tests.rs @@ -12,15 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. -use common_test_util::find_workspace_path; - use crate::test_util; #[tokio::test] async fn test_stream_to_json() { - let origin_path = &find_workspace_path("/src/common/datasource/tests/json/basic.json") - .display() - .to_string(); + let origin_path = "json/basic.json"; // A small threshold // Triggers the flush each writes @@ -33,13 +29,9 @@ async fn test_stream_to_json() { #[tokio::test] async fn test_stream_to_csv() { - let origin_path = &find_workspace_path("/src/common/datasource/tests/csv/basic.csv") - .display() - .to_string(); + let origin_path = "csv/basic.csv"; - let format_path = &find_workspace_path("/src/common/datasource/tests/csv/basic_format.csv") - .display() - .to_string(); + let format_path = "csv/basic_format.csv"; // A small threshold // Triggers the flush each writes