feat: Add --path-mode and change default download behavior to only keep paths after glob patterns

quodlibetor · quodlibetor · commit 958fc6e26e44 · 2025-02-19T21:28:44.000-05:00
Instead of always reproducing the full path. Previous behavior
can be restored with the new --path-mode=abs arg.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,12 @@
 
 All notable changes to this project will be documented in this file.
 
+## v0.3.0 - 2025-02-20
+
+### 🚀 Features
+
+- Add --path-mode and change default download behavior to only keep paths after glob patterns
+
 ## v0.2.7 - 2025-02-19
 
 ### 🚀 Features
diff --git a/src/glob_matcher.rs b/src/glob_matcher.rs
@@ -1,8 +1,5 @@
-#![allow(dead_code)]
-
 //! A pattern is a glob that knows how to split itself into a prefix and join with a partial prefix
-
-const GLOB_CHARS: &[char] = &['*', '?', '[', '{'];
+#![allow(dead_code)]
 
 use std::collections::BTreeSet;
 
@@ -18,6 +15,8 @@ pub use engine::{Engine, S3Engine};
 
 mod glob;
 
+pub(crate) const GLOB_CHARS: &[char] = &['*', '?', '[', '{'];
+
 /// The maximum number of prefixes that can be generated by the glob matcher
 ///
 /// Checking that constructed prefixes exist is significantly slower than
diff --git a/src/main.rs b/src/main.rs
@@ -8,8 +8,8 @@ use anyhow::{anyhow, bail, Context as _, Result};
 use aws_config::meta::region::RegionProviderChain;
 use aws_sdk_s3::types::Object;
 use aws_sdk_s3::{config::BehaviorVersion, config::Region, Client};
-use clap::{ArgAction, Parser, Subcommand};
-use glob_matcher::{S3Engine, S3GlobMatcher};
+use clap::{ArgAction, Parser, Subcommand, ValueEnum};
+use glob_matcher::{S3Engine, S3GlobMatcher, GLOB_CHARS};
 use humansize::{FormatSizeOptions, SizeFormatter, DECIMAL};
 use num_format::{Locale, ToFormattedString};
 use regex::Regex;
@@ -77,6 +77,16 @@ enum Command {
         /// The full key name will be reproduced in the directory, so multiple
         /// folders may be created.
         dest: String,
+
+        /// Control how S3 object keys are mapped to local file paths
+        ///
+        /// - absolute | abs: the full key path will be reproduced in the
+        ///   destination
+        /// - from-first-glob | g: the key path relative to the first path part
+        ///   containing a glob in the pattern will be reproduced in the
+        ///   destination
+        #[clap(short, long, verbatim_doc_comment, default_value = "from-first-glob")]
+        path_mode: PathType,
     },
 
     /// Learn how to tune s3glob's parallelism for better performance
@@ -120,6 +130,42 @@ enum Command {
     },
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum PathType {
+    Abs,
+    Absolute,
+    G,
+    FromFirstGlob,
+}
+
+impl ValueEnum for PathType {
+    fn value_variants<'a>() -> &'a [Self] {
+        &[
+            PathType::Absolute,
+            PathType::Abs,
+            PathType::FromFirstGlob,
+            PathType::G,
+        ]
+    }
+
+    fn from_str(s: &str, _ignore_case: bool) -> Result<Self, String> {
+        match s {
+            "absolute" | "abs" => Ok(PathType::Absolute),
+            "from-first-glob" | "g" => Ok(PathType::FromFirstGlob),
+            _ => Err(format!("invalid path type: {}", s)),
+        }
+    }
+
+    fn to_possible_value(&self) -> Option<clap::builder::PossibleValue> {
+        match self {
+            PathType::Abs => Some(clap::builder::PossibleValue::new("abs")),
+            PathType::Absolute => Some(clap::builder::PossibleValue::new("absolute")),
+            PathType::FromFirstGlob => Some(clap::builder::PossibleValue::new("from-first-glob")),
+            PathType::G => Some(clap::builder::PossibleValue::new("g")),
+        }
+    }
+}
+
 #[derive(Debug, Parser)]
 #[command(version, author, about, max_term_width = 80)]
 /// A fast aws s3 ls and downloader that supports glob patterns
@@ -228,11 +274,11 @@ async fn run(opts: Opts) -> Result<()> {
     let client = create_s3_client(&opts, &bucket).await?;
 
     let prefix = raw_pattern
-        .find(['*', '?', '[', '{'])
+        .find(GLOB_CHARS)
         .map_or(raw_pattern.clone(), |i| raw_pattern[..i].to_owned());
 
     let engine = S3Engine::new(client.clone(), bucket.clone(), opts.delimiter.to_string());
-    let matcher = S3GlobMatcher::parse(raw_pattern, &opts.delimiter.to_string())?;
+    let matcher = S3GlobMatcher::parse(raw_pattern.clone(), &opts.delimiter.to_string())?;
     let mut prefixes = match matcher.find_prefixes(engine).await {
         Ok(prefixes) => prefixes,
         Err(err) => {
@@ -331,7 +377,9 @@ async fn run(opts: Opts) -> Result<()> {
                 Duration::from_millis(start.elapsed().as_millis() as u64)
             );
         }
-        Command::Download { dest, .. } => {
+        Command::Download {
+            dest, path_mode, ..
+        } => {
             let mut matching_objects = Vec::new();
             let mut match_count = 0;
             while let Some(PrefixResult {
@@ -356,9 +404,13 @@ async fn run(opts: Opts) -> Result<()> {
             let obj_count = objects.len();
             let base_path = Path::new(&dest);
             let mut total_bytes = 0_usize;
+            let prefix_to_strip = extract_prefix_to_strip(&raw_pattern, path_mode);
             for (i, obj) in objects.iter().enumerate() {
                 let key = obj.key.as_ref().unwrap();
-                let path = base_path.join(key);
+                let key_suffix = key
+                    .strip_prefix(&prefix_to_strip)
+                    .expect("all found objects will include the prefix");
+                let path = base_path.join(key_suffix);
                 let dir = path.parent().unwrap();
                 std::fs::create_dir_all(dir)
                     .with_context(|| format!("Creating directory: {}", dir.display()))?;
@@ -407,6 +459,23 @@ async fn run(opts: Opts) -> Result<()> {
     Ok(())
 }
 
+fn extract_prefix_to_strip(raw_pattern: &str, path_mode: PathType) -> String {
+    match path_mode {
+        PathType::Abs | PathType::Absolute => String::new(),
+        PathType::FromFirstGlob | PathType::G => {
+            let up_to_glob: String = raw_pattern
+                .chars()
+                .take_while(|c| !GLOB_CHARS.contains(c))
+                .collect();
+            // find the last slash in the prefix and only include that
+            match up_to_glob.rfind('/') {
+                Some(slash_idx) => up_to_glob[..slash_idx + 1].to_string(),
+                None => up_to_glob,
+            }
+        }
+    }
+}
+
 #[derive(Debug)]
 struct PrefixResult {
     #[allow(dead_code)]
@@ -607,6 +676,7 @@ pub(crate) fn setup_logging(directive: Option<&str>) {
 
 #[cfg(test)]
 mod tests {
+    #![allow(clippy::comparison_to_empty)]
     use aws_sdk_s3::types::Object;
     use rstest::rstest;
 
@@ -630,4 +700,54 @@ mod tests {
     fn test_format_invalid_variable() {
         assert!(compile_format("{invalid_var}").is_err());
     }
+
+    macro_rules! assert_extract_prefix_to_strip {
+        ($pattern:expr, $path_mode:expr, $expected:expr) => {
+            let actual = extract_prefix_to_strip($pattern, $path_mode);
+            assert2::check!(
+                actual == $expected,
+                "input: {} path_mode: {:?}",
+                $pattern,
+                $path_mode,
+            );
+        };
+    }
+
+    #[test]
+    fn test_extract_prefix_to_strip() {
+        // Test absolute path mode
+        assert_extract_prefix_to_strip!("prefix/path/to/*.txt", PathType::Absolute, "");
+        assert_extract_prefix_to_strip!("bucket/deep/path/*.txt", PathType::Abs, "");
+
+        // Test from-first-glob path mode
+        assert_extract_prefix_to_strip!(
+            "prefix/path/to/*.txt",
+            PathType::FromFirstGlob,
+            "prefix/path/to/"
+        );
+        assert_extract_prefix_to_strip!(
+            "prefix/path/*/more/*.txt",
+            PathType::FromFirstGlob,
+            "prefix/path/"
+        );
+        assert_extract_prefix_to_strip!("prefix/*.txt", PathType::FromFirstGlob, "prefix/");
+        assert_extract_prefix_to_strip!("*.txt", PathType::FromFirstGlob, "");
+        assert_extract_prefix_to_strip!("prefix/a.txt", PathType::FromFirstGlob, "prefix/");
+        // Test with different glob characters
+        assert_extract_prefix_to_strip!(
+            "prefix/path/to/[abc]/*.txt",
+            PathType::FromFirstGlob,
+            "prefix/path/to/"
+        );
+        assert_extract_prefix_to_strip!(
+            "prefix/path/to/?/*.txt",
+            PathType::FromFirstGlob,
+            "prefix/path/to/"
+        );
+        assert_extract_prefix_to_strip!(
+            "prefix/path/{a,b}/*.txt",
+            PathType::FromFirstGlob,
+            "prefix/path/"
+        );
+    }
 }
diff --git a/tests/integration.rs b/tests/integration.rs
@@ -80,7 +80,12 @@ async fn test_s3glob_pattern_matching(
         let tempdir = TempDir::new()?;
         let mut cmd = run_s3glob(
             port,
-            &[command, uri.as_str(), tempdir.path().to_str().unwrap()],
+            &[
+                command,
+                "-pabs",
+                uri.as_str(),
+                tempdir.path().to_str().unwrap(),
+            ],
         )?;
         let _ = cmd.assert().success();
         for object in &test_objects {
@@ -95,6 +100,59 @@ async fn test_s3glob_pattern_matching(
     Ok(())
 }
 
+#[rstest]
+#[case("prefix/2024-01/file1.txt", &["file1.txt"])]
+#[case("prefix/2024-01/file*.txt", &["file1.txt", "file2.txt"])]
+#[case("prefix/2024-*/file1.txt", &["2024-02/file1.txt"])]
+#[case("prefix/2024-*/nested/*3*", &["2024-02/nested/file3.txt"])]
+#[case("prefix/2024-0{1,3}/*", &["2024-01/file1.txt", "2024-03/file5.txt"])]
+#[tokio::test]
+async fn test_download_prefix_from_first_glob(
+    #[case] glob: &str,
+    #[case] expected: &[&str],
+) -> anyhow::Result<()> {
+    let (_node, port, client) = minio_and_client().await;
+
+    let bucket = "test-bucket";
+    client.create_bucket().bucket(bucket).send().await?;
+
+    let test_objects = vec![
+        "prefix/2024-01/file1.txt",
+        "prefix/2024-01/file2.txt",
+        "prefix/2024-02/nested/file3.txt",
+        "prefix/2024-02/nested/file4.txt",
+        "prefix/2024-03/file5.txt",
+    ];
+    for key in &test_objects {
+        create_object(&client, bucket, key).await?;
+    }
+
+    let tempdir = TempDir::new()?;
+
+    let mut cmd = run_s3glob(
+        port,
+        &[
+            "dl",
+            "-p",
+            "from-first-glob",
+            format!("s3://{}/{}", bucket, glob).as_str(),
+            tempdir.path().to_str().unwrap(),
+        ],
+    )?;
+
+    let _ = cmd.assert().success();
+
+    for object in test_objects {
+        if expected.contains(&object) {
+            tempdir.child(object).assert(predicate::path::exists());
+        } else {
+            tempdir.child(object).assert(predicate::path::missing());
+        }
+    }
+
+    Ok(())
+}
+
 #[rstest]
 #[case("{key}", "test/file.txt")]
 #[case("{size_bytes}", "1234")]
@@ -179,7 +237,7 @@ async fn test_patterns_in_file_not_path_component(
     } else {
         let tempdir = TempDir::new()?;
         let out_path = tempdir.path().to_str().unwrap();
-        let mut cmd = run_s3glob(port, &[command, needle.as_str(), out_path])?;
+        let mut cmd = run_s3glob(port, &[command, "-pabs", needle.as_str(), out_path])?;
         let _ = cmd.assert().success();
         for object in &test_objects {
             if expected.contains(object) {
