From 279c846f0614a3d75dc5a0dd02a30af300617143 Mon Sep 17 00:00:00 2001
From: Nate Anderson <nate@momentohq.com>
Date: Fri, 2 Aug 2024 16:06:55 -0700
Subject: [PATCH] chore: make s3 and API gateway cloud linting optional

Make s3 and API gateway cloud linting optional.
---
 momento-cli-opts/src/lib.rs                   | 14 ++++++-
 .../src/commands/cloud_linter/linter_cli.rs   | 39 ++++++++++++-------
 momento/src/main.rs                           |  4 ++
 3 files changed, 41 insertions(+), 16 deletions(-)

diff --git a/momento-cli-opts/src/lib.rs b/momento-cli-opts/src/lib.rs
index 8c42b3a..11b8008 100644
--- a/momento-cli-opts/src/lib.rs
+++ b/momento-cli-opts/src/lib.rs
@@ -226,14 +226,24 @@ to help find opportunities for optimizations with Momento.
         region: String,
         #[arg(
             long = "enable-ddb-ttl-check",
-            help = "Opt in to check whether ddb tables have ttl enabled. If there are lots of tables, could slow down data collection"
+            help = "Opt in to check whether ddb tables have ttl enabled. If there are lots of tables, this could slow down data collection"
         )]
         enable_ddb_ttl_check: bool,
         #[arg(
             long = "enable-gsi",
-            help = "Opt in to check metrics on dynamodb gsi's. If there are lots of tables with gsi's, could slow down data collection"
+            help = "Opt in to check metrics on dynamodb gsi's. If there are lots of tables with gsi's, this could slow down data collection"
         )]
         enable_gsi: bool,
+        #[arg(
+            long = "enable-s3",
+            help = "Opt in to check metrics on s3. If there are lots of s3 buckets, this could slow down data collection"
+        )]
+        enable_s3: bool,
+        #[arg(
+            long = "enable-apt-gateway",
+            help = "Opt in to check metrics on API Gateway"
+        )]
+        enable_api_gateway: bool,
         #[arg(
             value_enum,
             long = "resource",
diff --git a/momento/src/commands/cloud_linter/linter_cli.rs b/momento/src/commands/cloud_linter/linter_cli.rs
index 86776f4..0bc2a74 100644
--- a/momento/src/commands/cloud_linter/linter_cli.rs
+++ b/momento/src/commands/cloud_linter/linter_cli.rs
@@ -28,6 +28,8 @@ pub async fn run_cloud_linter(
     region: String,
     enable_ddb_ttl_check: bool,
     enable_gsi: bool,
+    enable_s3: bool,
+    enable_api_gateway: bool,
     only_collect_for_resource: Option<CloudLinterResources>,
     metric_collection_rate: u32,
 ) -> Result<(), CliError> {
@@ -49,6 +51,8 @@ pub async fn run_cloud_linter(
             tx,
             enable_ddb_ttl_check,
             enable_gsi,
+            enable_s3,
+            enable_api_gateway,
             only_collect_for_resource,
             metric_collection_rate,
         )
@@ -80,11 +84,14 @@ pub async fn run_cloud_linter(
     Ok(())
 }
 
+#[allow(clippy::too_many_arguments)]
 async fn process_data(
     region: String,
     sender: Sender<Resource>,
     enable_ddb_ttl_check: bool,
     enable_gsi: bool,
+    enable_s3: bool,
+    enable_api_gateway: bool,
     only_collect_for_resource: Option<CloudLinterResources>,
     metric_collection_rate: u32,
 ) -> Result<(), CliError> {
@@ -170,21 +177,25 @@ async fn process_data(
         };
     };
 
-    process_s3_resources(
-        &config,
-        Arc::clone(&control_plane_limiter),
-        Arc::clone(&metrics_limiter),
-        sender.clone(),
-    )
-    .await?;
+    if enable_s3 {
+        process_s3_resources(
+            &config,
+            Arc::clone(&control_plane_limiter),
+            Arc::clone(&metrics_limiter),
+            sender.clone(),
+        )
+        .await?;
+    }
 
-    process_api_gateway_resources(
-        &config,
-        Arc::clone(&control_plane_limiter),
-        Arc::clone(&metrics_limiter),
-        sender.clone(),
-    )
-    .await?;
+    if enable_api_gateway {
+        process_api_gateway_resources(
+            &config,
+            Arc::clone(&control_plane_limiter),
+            Arc::clone(&metrics_limiter),
+            sender.clone(),
+        )
+        .await?;
+    }
 
     process_ddb_resources(
         &config,
diff --git a/momento/src/main.rs b/momento/src/main.rs
index 3b20cea..635a66f 100644
--- a/momento/src/main.rs
+++ b/momento/src/main.rs
@@ -269,11 +269,15 @@ async fn run_momento_command(args: momento_cli_opts::Momento) -> Result<(), CliE
                 resource,
                 metric_collection_rate,
                 enable_gsi,
+                enable_s3,
+                enable_api_gateway,
             } => {
                 commands::cloud_linter::linter_cli::run_cloud_linter(
                     region,
                     enable_ddb_ttl_check,
                     enable_gsi,
+                    enable_s3,
+                    enable_api_gateway,
                     resource,
                     metric_collection_rate,
                 )