aws: Add regex pattern matching to cleanup-ssm tool (#6870)

I had noticed that some of the things we were deleting might be useful
like:

```
Would delete: pytorch-labs-cloudwatch_agent_config_runner_linux_nvidia
```

So I added pattern matching to only delete things related to runner
spin-up.

---------

Signed-off-by: Eli Uriegas <[email protected]>

Author: seemethere

aws/tools/cleanup-ssm/Cargo.toml

@@ -12,6 +12,7 @@ chrono = "0.4"
 aws-smithy-types = "1.3"
 humantime = "2.1"
 indicatif = "0.17"
+regex = "1.0"
 
 [dev-dependencies]
 mockall = "0.12"

aws/tools/cleanup-ssm/src/filter.rs

@@ -1,13 +1,16 @@
 use crate::TimeProvider;
 use aws_sdk_ssm::types::ParameterMetadata;
 use chrono::{DateTime, Duration};
+use regex::Regex;
 
 pub fn filter_old_parameters<T: TimeProvider>(
     parameters: &[ParameterMetadata],
     time_provider: &T,
     older_than_seconds: f64,
-) -> Vec<String> {
+    pattern: &str,
+) -> Result<Vec<String>, regex::Error> {
     let threshold = time_provider.now() - Duration::seconds(older_than_seconds as i64);
+    let regex = Regex::new(pattern)?;
     let mut parameters_to_delete = Vec::new();
 
     for parameter in parameters {
@@ -17,13 +20,15 @@ pub fn filter_old_parameters<T: TimeProvider>(
 
             if last_modified_time < threshold {
                 if let Some(name) = parameter.name() {
-                    parameters_to_delete.push(name.to_string());
+                    if regex.is_match(name) {
+                        parameters_to_delete.push(name.to_string());
+                    }
                 }
             }
         }
     }
 
-    parameters_to_delete
+    Ok(parameters_to_delete)
 }
 
 #[cfg(test)]
@@ -54,7 +59,7 @@ mod tests {
         let parameters = vec![];
         let time_provider = MockTimeProvider::new(Utc::now());
 
-        let result = filter_old_parameters(&parameters, &time_provider, 86400.0); // 1 day in seconds
+        let result = filter_old_parameters(&parameters, &time_provider, 86400.0, ".*").unwrap(); // 1 day in seconds
 
         assert_eq!(result.len(), 0);
     }
@@ -72,7 +77,7 @@ mod tests {
         let time_provider = MockTimeProvider::new(now);
         let parameters = vec![parameter];
 
-        let result = filter_old_parameters(&parameters, &time_provider, 86400.0); // 1 day in seconds
+        let result = filter_old_parameters(&parameters, &time_provider, 86400.0, ".*").unwrap(); // 1 day in seconds
 
         assert_eq!(result.len(), 0);
     }
@@ -90,7 +95,7 @@ mod tests {
         let time_provider = MockTimeProvider::new(now);
         let parameters = vec![parameter];
 
-        let result = filter_old_parameters(&parameters, &time_provider, 86400.0); // 1 day in seconds
+        let result = filter_old_parameters(&parameters, &time_provider, 86400.0, ".*").unwrap(); // 1 day in seconds
 
         assert_eq!(result.len(), 1);
         assert_eq!(result[0], "old-param");
@@ -115,7 +120,7 @@ mod tests {
         let time_provider = MockTimeProvider::new(now);
         let parameters = vec![old_parameter, recent_parameter];
 
-        let result = filter_old_parameters(&parameters, &time_provider, 172800.0); // 2 days in seconds
+        let result = filter_old_parameters(&parameters, &time_provider, 172800.0, ".*").unwrap(); // 2 days in seconds
 
         assert_eq!(result.len(), 1);
         assert_eq!(result[0], "old-param");
@@ -130,7 +135,7 @@ mod tests {
         let time_provider = MockTimeProvider::new(Utc::now());
         let parameters = vec![parameter];
 
-        let result = filter_old_parameters(&parameters, &time_provider, 86400.0); // 1 day in seconds
+        let result = filter_old_parameters(&parameters, &time_provider, 86400.0, ".*").unwrap(); // 1 day in seconds
 
         assert_eq!(result.len(), 0);
     }
@@ -147,7 +152,51 @@ mod tests {
         let time_provider = MockTimeProvider::new(now);
         let parameters = vec![parameter];
 
-        let result = filter_old_parameters(&parameters, &time_provider, 86400.0); // 1 day in seconds
+        let result = filter_old_parameters(&parameters, &time_provider, 86400.0, ".*").unwrap(); // 1 day in seconds
+
+        assert_eq!(result.len(), 0);
+    }
+
+    #[test]
+    fn test_filter_old_parameters_with_pattern_match() {
+        let now = Utc::now();
+        let old_time = now - Duration::days(5);
+
+        let matching_parameter = ParameterMetadata::builder()
+            .name("gh-ci-i-test-param")
+            .last_modified_date(AwsDateTime::from_secs(old_time.timestamp()))
+            .build();
+
+        let non_matching_parameter = ParameterMetadata::builder()
+            .name("other-param")
+            .last_modified_date(AwsDateTime::from_secs(old_time.timestamp()))
+            .build();
+
+        let time_provider = MockTimeProvider::new(now);
+        let parameters = vec![matching_parameter, non_matching_parameter];
+
+        let result =
+            filter_old_parameters(&parameters, &time_provider, 86400.0, "gh-ci-i-.*").unwrap(); // 1 day in seconds
+
+        assert_eq!(result.len(), 1);
+        assert_eq!(result[0], "gh-ci-i-test-param");
+    }
+
+    #[test]
+    fn test_filter_old_parameters_with_pattern_no_match() {
+        let now = Utc::now();
+        let old_time = now - Duration::days(5);
+
+        let parameter = ParameterMetadata::builder()
+            .name("other-param")
+            .last_modified_date(AwsDateTime::from_secs(old_time.timestamp()))
+            .build();
+
+        let time_provider = MockTimeProvider::new(now);
+        let parameters = vec![parameter];
+
+        let result =
+            filter_old_parameters(&parameters, &time_provider, 86400.0, "gh-ci-i-.*").unwrap(); // 1 day in seconds
 
         assert_eq!(result.len(), 0);
     }

aws/tools/cleanup-ssm/src/lib.rs

@@ -1,6 +1,42 @@
 use aws_sdk_ssm::types::ParameterMetadata;
 use chrono::{DateTime, Utc};
 
+#[derive(Debug)]
+pub enum CleanupError {
+    Aws(aws_sdk_ssm::Error),
+    Regex(regex::Error),
+}
+
+impl From<aws_sdk_ssm::Error> for CleanupError {
+    fn from(err: aws_sdk_ssm::Error) -> Self {
+        CleanupError::Aws(err)
+    }
+}
+
+impl From<regex::Error> for CleanupError {
+    fn from(err: regex::Error) -> Self {
+        CleanupError::Regex(err)
+    }
+}
+
+impl std::fmt::Display for CleanupError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            CleanupError::Aws(err) => write!(f, "AWS error: {}", err),
+            CleanupError::Regex(err) => write!(f, "Regex error: {}", err),
+        }
+    }
+}
+
+impl std::error::Error for CleanupError {
+    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        match self {
+            CleanupError::Aws(err) => Some(err),
+            CleanupError::Regex(err) => Some(err),
+        }
+    }
+}
+
 #[cfg(test)]
 use mockall::automock;
 
@@ -13,6 +49,7 @@ pub struct CleanupConfig {
     pub region: String,
     pub dry_run: bool,
     pub older_than_seconds: f64,
+    pub pattern: String,
 }
 
 #[derive(Debug)]
@@ -49,11 +86,15 @@ pub async fn cleanup_ssm_parameters<C: SsmClient, T: TimeProvider>(
     client: &C,
     time_provider: &T,
     config: &CleanupConfig,
-) -> Result<CleanupResult, aws_sdk_ssm::Error> {
+) -> Result<CleanupResult, CleanupError> {
     let parameters = client.describe_parameters().await?;
 
-    let parameters_to_delete =
-        filter::filter_old_parameters(&parameters, time_provider, config.older_than_seconds);
+    let parameters_to_delete = filter::filter_old_parameters(
+        &parameters,
+        time_provider,
+        config.older_than_seconds,
+        &config.pattern,
+    )?;
 
     println!("Found {} parameters to delete", parameters_to_delete.len());
     let parameters_found = parameters_to_delete.len();
@@ -126,6 +167,7 @@ mod tests {
             region: "us-east-1".to_string(),
             dry_run: true,
             older_than_seconds: 86400.0, // 1 day in seconds
+            pattern: ".*".to_string(),
         };
 
         let result = cleanup_ssm_parameters(&mock_client, &time_provider, &config)

aws/tools/cleanup-ssm/src/main.rs

@@ -13,10 +13,13 @@ struct Args {
     // time duration older than which to delete parameters (e.g., "1d", "2h", "30m")
     #[clap(long, default_value = "1d")]
     older_than: String,
+    // regex pattern to match parameter names for deletion
+    #[clap(long, default_value = "gh-ci-i-.*")]
+    pattern: String,
 }
 
 #[tokio::main]
-async fn main() -> Result<(), Box<aws_sdk_ssm::Error>> {
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let args = Args::parse();
 
     // Parse the human-readable time string into a Duration
@@ -36,6 +39,7 @@ async fn main() -> Result<(), Box<aws_sdk_ssm::Error>> {
         region: args.region,
         dry_run: args.dry_run,
         older_than_seconds,
+        pattern: args.pattern,
     };
 
     let result = cleanup_ssm_parameters(&client, &time_provider, &config).await?;