Allow to filter buckets

Author: vprus

tools/storage-advisor/src/main.go

@@ -28,6 +28,8 @@ var mode = flag.String("mode", "cli", "Use as a CLI tool")
 
 var region = flag.String("region", "", "AWS region")
 
+var bucketsFilterS = flag.String("buckets", "", "Comma-separated list of buckets to check")
+
 var apiToken = flag.String("api-token", "", "API token for Joom Cloud")
 
 var apiEndpoint = flag.String("api-endpoint", "https://api.cloud.joom.ai/v1", "API endpoint URL")
@@ -46,6 +48,14 @@ func main() {
 	client := s3.NewFromConfig(cfg)
 	api := NewApi(*apiEndpoint, apiToken)
 
+	bucketsFilter := map[string]struct{}{}
+	if bucketsFilterS != nil && *bucketsFilterS != "" {
+		for _, bucket := range strings.Split(*bucketsFilterS, ",") {
+			bucketsFilter[bucket] = struct{}{}
+		}
+		fmt.Printf("Filtering buckets: %v\n", len(bucketsFilter))
+	}
+
 	switch *mode {
 	case "cli":
 		// For CLI, use a simplified logger
@@ -89,7 +99,7 @@ func main() {
 		identity := *stsOutput.Arn
 		fmt.Printf("Info: your AWS identity is %s\n", identity)
 
-		buckets, issues, err := runS3Checks(rootCtx, client)
+		buckets, issues, err := runS3Checks(rootCtx, client, bucketsFilter)
 		if err != nil {
 			var ae smithy.APIError
 			if errors.As(err, &ae) {
@@ -125,7 +135,7 @@ func main() {
 				log.Info().Msgf("running with AWS identity %s", *stsOutput.Arn)
 			}
 
-			buckets, _, err := runS3Checks(rootCtx, client)
+			buckets, _, err := runS3Checks(rootCtx, client, bucketsFilter)
 			if err != nil {
 				log.Fatal().Err(err).Msg("could not list buckets")
 			}

tools/storage-advisor/src/s3Issues.go

@@ -508,14 +508,25 @@ func summarizeIssues(results []Finding, checkName string) {
 // as separate global list of all issues.
 // The function tries to run and return something in face of all errors, and will only
 // return error is we can't do anything -- like we can't list buckets.
-func runS3Checks(ctx context.Context, s3client *s3.Client) ([]Bucket, []FindingWithBucket, error) {
+func runS3Checks(ctx context.Context, s3client *s3.Client, bucketsFilter map[string]struct{}) ([]Bucket, []FindingWithBucket, error) {
 	fmt.Printf(green("Running S3 checks\n\n"))
 
 	buckets, err := listBuckets(s3client)
 	if err != nil {
 		return nil, nil, err
 	}
 
+	if len(bucketsFilter) > 0 {
+		filtered := make([]Bucket, 0, len(bucketsFilter))
+		for _, b := range buckets {
+			if _, found := bucketsFilter[b.Name]; found {
+				filtered = append(filtered, b)
+			}
+		}
+		fmt.Printf("Filtered %d buckets using %d filters into %d buckets", len(buckets), len(bucketsFilter), len(filtered))
+		buckets = filtered
+	}
+
 	for i := range buckets {
 		// Get bucket region
 		r, err := s3client.GetBucketLocation(ctx, &s3.GetBucketLocationInput{