diff --git a/src/config/src/utils/async_file.rs b/src/config/src/utils/async_file.rs index 648e3d660e..922925ad6e 100644 --- a/src/config/src/utils/async_file.rs +++ b/src/config/src/utils/async_file.rs @@ -13,7 +13,12 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::{fs::Metadata, ops::Range, path::Path, time::SystemTime}; +use std::{ + fs::Metadata, + ops::Range, + path::{Path, PathBuf}, + time::SystemTime, +}; use async_walkdir::WalkDir; use futures::StreamExt; @@ -138,9 +143,112 @@ pub async fn clean_empty_dirs( Ok(()) } +/// Asynchronously scans a directory tree and returns a vector of canonicalized file paths +/// that match a given filter. +/// +/// This function walks the directory starting from `root`, applying a filter to each +/// entry. It uses `async_walkdir` for efficient, non-blocking directory traversal. +/// The filter logic determines whether to continue the walk, ignore a directory, +/// or ignore a file. +/// +/// For entries that pass the filter and are files, their paths are canonicalized +/// to produce absolute paths. +/// +/// # Type Parameters +/// +/// * `P` - A type that can be referenced as a `Path`, e.g., `&str` or `PathBuf`. +/// * `F` - A closure that takes a `PathBuf` and returns a boolean. +/// +/// # Arguments +/// +/// * `root` - The path to the root directory to start the scan from. +/// * `filter` - An asynchronous closure that is called for each entry in the directory. It should +/// return `true` to keep an entry or `false` to discard it. If a directory is discarded, its +/// contents will not be visited. +/// * `limit` - An optional `usize` to limit the number of file paths collected. +/// +/// # Returns +/// +/// A `Result` containing either: +/// - `Ok(Vec)`: A vector of canonicalized file path strings. +/// - `Err(std::io::Error)`: An I/O error that occurred during scanning. +pub async fn scan_files_filtered( + root: P, + filter: F, + limit: Option, +) -> Result, std::io::Error> +where + P: AsRef, + F: Fn(PathBuf) -> bool + Send + Clone + 'static, +{ + let walker = WalkDir::new(root).filter(move |entry| { + let path = entry.path(); + let filter = filter.clone(); + async move { + let is_dir = path.is_dir(); + if filter(path) { + async_walkdir::Filtering::Continue + } else if is_dir { + async_walkdir::Filtering::IgnoreDir + } else { + async_walkdir::Filtering::Ignore + } + } + }); + + let walker = walker.filter_map(|item| async { + item.ok().and_then(|dir_entry| { + let pb = dir_entry.path(); + + if !pb.is_file() { + None + } else { + pb.canonicalize() + .ok() + .and_then(|cpath| cpath.to_str().map(String::from)) + } + }) + }); + + let files = if let Some(limit_count) = limit { + walker.take(limit_count).collect().await + } else { + walker.collect().await + }; + + Ok(files) +} + +#[inline(always)] +pub async fn scan_files>( + root: P, + ext: &str, + limit: Option, +) -> Result, std::io::Error> { + let ext = ext.to_lowercase(); + scan_files_filtered( + root, + move |f| { + if f.is_dir() { + true + } else { + f.extension() + .and_then(|file_ext| file_ext.to_str()) + .map(|s| s.to_lowercase() == ext) + .unwrap_or_default() + } + }, + limit, + ) + .await +} + #[cfg(test)] mod tests { + use chrono::{TimeZone, Utc}; + use super::*; + use crate::utils::file::wal_dir_datetime_filter_builder; #[tokio::test] async fn test_get_file_contents_with_range() { @@ -163,4 +271,124 @@ mod tests { std::fs::remove_file(file_name).unwrap(); } + + #[tokio::test] + async fn test_scan_files_filtered_basic() { + let threads: Vec<_> = (1..2).collect(); + let years = vec![2025]; + let months: Vec<_> = (9..=10).collect(); + let days: Vec<_> = (8..=11).collect(); + let hours: Vec<_> = (1..=5).collect(); + let filenames: Vec<_> = vec!["a", "b", "c", "d"]; + let extensions = vec!["parquet", "blink"]; + + let wal_root = tempfile::tempdir().expect("Temp dir"); + + let inner_path = wal_root + .path() + .join("files") + .join("") + .join("") + .join(""); + + for thread in threads { + for year in &years { + for month in &months { + for day in &days { + for hour in &hours { + tokio::fs::create_dir_all(format!( + "{}/{}/{}/{}/{}/{}", + inner_path.display(), + thread, + year, + month, + day, + hour + )) + .await + .expect("Pretest setup failure"); + + for filename in &filenames { + for extension in &extensions { + tokio::fs::File::create(format!( + "{}/{}/{}/{}/{}/{}/{}.{}", + inner_path.display(), + thread, + year, + month, + day, + hour, + filename, + extension + )) + .await + .expect("Pretest setup failure"); + } + } + } + } + } + } + } + + let start_time = Utc.with_ymd_and_hms(2025, 9, 10, 2, 0, 0).single().unwrap(); + let end_time = Utc.with_ymd_and_hms(2025, 9, 10, 4, 0, 0).single().unwrap(); + // 3 hours (2, 3, 4 - inclusive matching) - 4 parquet files per hour -> 12 files in total + let filter = wal_dir_datetime_filter_builder( + start_time, + end_time, + "parquet".to_string(), + inner_path.components().count() + 1, + ); + let files = scan_files_filtered(inner_path, filter, None) + .await + .expect("Basic Test Failure"); + + assert_eq!(files.len(), 12); + } + + #[tokio::test] + async fn test_scan_files() { + let test_dir = "test_scan_files"; + let sub_dir = format!("{}/subdir", test_dir); + + // Create test directory structure + tokio::fs::create_dir_all(&sub_dir).await.unwrap(); + + // Create test files with different extensions + put_file_contents(&format!("{}/file1.txt", test_dir), b"content1") + .await + .unwrap(); + put_file_contents(&format!("{}/file2.txt", test_dir), b"content2") + .await + .unwrap(); + put_file_contents(&format!("{}/file3.log", test_dir), b"log content") + .await + .unwrap(); + put_file_contents(&format!("{}/file4.txt", sub_dir), b"sub content") + .await + .unwrap(); + + // Test scanning for txt files without limit + let txt_files = scan_files(test_dir, "txt", None).await.unwrap(); + println!("{:?}", txt_files); + assert_eq!(txt_files.len(), 3); + assert!(txt_files.iter().all(|f| f.ends_with(".txt"))); + + // Test scanning with limit + let limited_files = scan_files(test_dir, "txt", Some(2)).await.unwrap(); + assert_eq!(limited_files.len(), 2); + + // Test scanning for log files + let log_files = scan_files(test_dir, "log", None).await.unwrap(); + assert_eq!(log_files.len(), 1); + assert!(log_files[0].ends_with(".log")); + + // Test scanning for non-existent extension + let empty_files = scan_files(test_dir, "nonexistent", None).await.unwrap(); + assert!(empty_files.is_empty()); + + // Cleanup + std::fs::remove_dir_all(test_dir).unwrap(); + } } diff --git a/src/config/src/utils/file.rs b/src/config/src/utils/file.rs index 5f7643ac01..28b0ed07c5 100644 --- a/src/config/src/utils/file.rs +++ b/src/config/src/utils/file.rs @@ -21,9 +21,7 @@ use std::{ }; use async_recursion::async_recursion; -use async_walkdir::WalkDir; use chrono::{DateTime, Datelike, TimeZone, Timelike, Utc}; -use futures::StreamExt; #[inline(always)] pub fn get_file_meta(path: impl AsRef) -> Result { @@ -117,6 +115,7 @@ pub fn wal_dir_datetime_filter_builder( extension_pattern: String, skip_count: usize, ) -> impl Fn(PathBuf) -> bool + Send + Clone + 'static { + let extension_pattern = extension_pattern.to_lowercase(); move |path: PathBuf| { let mut components = path .components() @@ -178,87 +177,15 @@ pub fn wal_dir_datetime_filter_builder( && (!path.is_file() || path .extension() - .and_then(|extension| extension.to_str().map(|s| s == extension_pattern)) + .and_then(|extension| { + extension + .to_str() + .map(|s| s.to_lowercase() == extension_pattern) + }) .unwrap_or_default()) } } -/// Asynchronously scans a directory tree and returns a vector of canonicalized file paths -/// that match a given filter. -/// -/// This function walks the directory starting from `root`, applying a filter to each -/// entry. It uses `async_walkdir` for efficient, non-blocking directory traversal. -/// The filter logic determines whether to continue the walk, ignore a directory, -/// or ignore a file. -/// -/// For entries that pass the filter and are files, their paths are canonicalized -/// to produce absolute paths. -/// -/// # Type Parameters -/// -/// * `P` - A type that can be referenced as a `Path`, e.g., `&str` or `PathBuf`. -/// * `F` - A closure that takes a `PathBuf` and returns a boolean. -/// -/// # Arguments -/// -/// * `root` - The path to the root directory to start the scan from. -/// * `filter` - An asynchronous closure that is called for each entry in the directory. It should -/// return `true` to keep an entry or `false` to discard it. If a directory is discarded, its -/// contents will not be visited. -/// * `limit` - An optional `usize` to limit the number of file paths collected. -/// -/// # Returns -/// -/// A `Result` containing either: -/// - `Ok(Vec)`: A vector of canonicalized file path strings. -/// - `Err(std::io::Error)`: An I/O error that occurred during scanning. -pub async fn scan_files_filtered( - root: P, - filter: F, - limit: Option, -) -> Result, std::io::Error> -where - P: AsRef, - F: Fn(PathBuf) -> bool + Send + Clone + 'static, -{ - let walker = WalkDir::new(root).filter(move |entry| { - let path = entry.path(); - let filter = filter.clone(); - async move { - let is_dir = path.is_dir(); - if filter(path) { - async_walkdir::Filtering::Continue - } else if is_dir { - async_walkdir::Filtering::IgnoreDir - } else { - async_walkdir::Filtering::Ignore - } - } - }); - - let walker = walker.filter_map(|item| async { - item.ok().and_then(|dir_entry| { - let pb = dir_entry.path(); - - if !pb.is_file() { - None - } else { - pb.canonicalize() - .ok() - .and_then(|cpath| cpath.to_str().map(String::from)) - } - }) - }); - - let files = if let Some(limit_count) = limit { - walker.take(limit_count).collect().await - } else { - walker.collect().await - }; - - Ok(files) -} - #[inline(always)] pub fn scan_files>( root: P, @@ -632,79 +559,4 @@ mod tests { assert!(has_symlink); } } - - #[tokio::test] - async fn test_scan_files_filtered_basic() { - let threads: Vec<_> = (1..2).collect(); - let years = vec![2025]; - let months: Vec<_> = (9..=10).collect(); - let days: Vec<_> = (8..=11).collect(); - let hours: Vec<_> = (1..=5).collect(); - let filenames: Vec<_> = vec!["a", "b", "c", "d"]; - let extensions = vec!["parquet", "blink"]; - - let wal_root = tempfile::tempdir().expect("Temp dir"); - - let inner_path = wal_root - .path() - .join("files") - .join("") - .join("") - .join(""); - - for thread in threads { - for year in &years { - for month in &months { - for day in &days { - for hour in &hours { - tokio::fs::create_dir_all(format!( - "{}/{}/{}/{}/{}/{}", - inner_path.display(), - thread, - year, - month, - day, - hour - )) - .await - .expect("Pretest setup failure"); - - for filename in &filenames { - for extension in &extensions { - tokio::fs::File::create(format!( - "{}/{}/{}/{}/{}/{}/{}.{}", - inner_path.display(), - thread, - year, - month, - day, - hour, - filename, - extension - )) - .await - .expect("Pretest setup failure"); - } - } - } - } - } - } - } - - let start_time = Utc.with_ymd_and_hms(2025, 9, 10, 2, 0, 0).single().unwrap(); - let end_time = Utc.with_ymd_and_hms(2025, 9, 10, 4, 0, 0).single().unwrap(); - // 3 hours (2, 3, 4 - inclusive matching) - 4 parquet files per hour -> 12 files in total - let filter = wal_dir_datetime_filter_builder( - start_time, - end_time, - "parquet".to_string(), - inner_path.components().count() + 1, - ); - let files = scan_files_filtered(inner_path, filter, None) - .await - .expect("Basic Test Failure"); - - assert_eq!(files.len(), 12); - } } diff --git a/src/service/search/grpc/wal.rs b/src/service/search/grpc/wal.rs index 5887ee1869..342ea6970c 100644 --- a/src/service/search/grpc/wal.rs +++ b/src/service/search/grpc/wal.rs @@ -26,7 +26,8 @@ use config::{ stream::{FileKey, StreamParams, StreamPartition}, }, utils::{ - file::{is_exists, scan_files_filtered, wal_dir_datetime_filter_builder}, + async_file::scan_files_filtered, + file::{is_exists, wal_dir_datetime_filter_builder}, parquet::{parse_time_range_from_filename, read_metadata_from_file}, record_batch_ext::concat_batches, size::bytes_to_human_readable,