Merge pull request #39686 from hashicorp/s-athena

Additional Athena sweepers
hashicorp · Oct 14, 2024 · 8d411b1 · 8d411b1
2 parents a5a8c95 + 2a24a11
commit 8d411b1
Show file tree

Hide file tree

Showing 4 changed files with 142 additions and 10 deletions.
diff --git a/internal/service/athena/data_catalog.go b/internal/service/athena/data_catalog.go
@@ -192,7 +192,7 @@ func resourceDataCatalogDelete(ctx context.Context, d *schema.ResourceData, meta
 
 	conn := meta.(*conns.AWSClient).AthenaClient(ctx)
 
-	log.Printf("[DEBUG] Deleting Athena Data Catalog: (%s)", d.Id())
+	log.Printf("[DEBUG] Deleting Athena Data Catalog (%s)", d.Id())
 	_, err := conn.DeleteDataCatalog(ctx, &athena.DeleteDataCatalogInput{
 		Name: aws.String(d.Id()),
 	})

diff --git a/internal/service/athena/database.go b/internal/service/athena/database.go
@@ -197,6 +197,7 @@ func resourceDatabaseDelete(ctx context.Context, d *schema.ResourceData, meta in
 		ResultConfiguration: expandResultConfiguration(d),
 	}
 
+	log.Printf("[DEBUG] Deleting Athena Database (%s)", d.Id())
 	output, err := conn.StartQueryExecution(ctx, input)
 
 	if err != nil {

diff --git a/internal/service/athena/sweep.go b/internal/service/athena/sweep.go
@@ -16,25 +16,87 @@ import (
 )
 
 func RegisterSweepers() {
+	resource.AddTestSweepers("aws_athena_data_catalog", &resource.Sweeper{
+		Name: "aws_athena_data_catalog",
+		F:    sweepDataCatalogs,
+		Dependencies: []string{
+			"aws_athena_database",
+		},
+	})
+
 	resource.AddTestSweepers("aws_athena_database", &resource.Sweeper{
 		Name: "aws_athena_database",
 		F:    sweepDatabases,
 	})
+
+	resource.AddTestSweepers("aws_athena_workgroup", &resource.Sweeper{
+		Name: "aws_athena_workgroup",
+		F:    sweepWorkGroups,
+		Dependencies: []string{
+			"aws_athena_data_catalog",
+		},
+	})
 }
 
-func sweepDatabases(region string) error {
+func sweepDataCatalogs(region string) error {
 	ctx := sweep.Context(region)
 	client, err := sweep.SharedRegionalSweepClient(ctx, region)
 	if err != nil {
 		return fmt.Errorf("error getting client: %s", err)
 	}
 	conn := client.AthenaClient(ctx)
-	input := &athena.ListDatabasesInput{
-		CatalogName: aws.String("AwsDataCatalog"),
+	input := &athena.ListDataCatalogsInput{}
+	sweepResources := make([]sweep.Sweepable, 0)
+
+	pages := athena.NewListDataCatalogsPaginator(conn, input)
+	for pages.HasMorePages() {
+		page, err := pages.NextPage(ctx)
+
+		if awsv2.SkipSweepError(err) {
+			log.Printf("[WARN] Skipping Athena Data Catalog sweep for %s: %s", region, err)
+			return nil
+		}
+
+		if err != nil {
+			return fmt.Errorf("error listing Athena Data Catalogs (%s): %w", region, err)
+		}
+
+		for _, v := range page.DataCatalogsSummary {
+			name := aws.ToString(v.CatalogName)
+
+			if name == "AwsDataCatalog" {
+				log.Printf("[INFO] Skipping Athena Data Catalog %s", name)
+				continue
+			}
+
+			r := resourceDataCatalog()
+			d := r.Data(nil)
+			d.SetId(name)
+
+			sweepResources = append(sweepResources, sweep.NewSweepResource(r, d, client))
+		}
 	}
+
+	err = sweep.SweepOrchestrator(ctx, sweepResources)
+
+	if err != nil {
+		return fmt.Errorf("error sweeping Athena Data Catalogs (%s): %w", region, err)
+	}
+
+	return nil
+}
+
+func sweepDatabases(region string) error {
+	ctx := sweep.Context(region)
+	client, err := sweep.SharedRegionalSweepClient(ctx, region)
+	if err != nil {
+		return fmt.Errorf("error getting client: %s", err)
+	}
+	conn := client.AthenaClient(ctx)
+	input := &athena.ListDataCatalogsInput{}
 	sweepResources := make([]sweep.Sweepable, 0)
 
-	pages := athena.NewListDatabasesPaginator(conn, input)
+	pages := athena.NewListDataCatalogsPaginator(conn, input)
 	for pages.HasMorePages() {
 		page, err := pages.NextPage(ctx)
 
@@ -44,15 +106,83 @@ func sweepDatabases(region string) error {
 		}
 
 		if err != nil {
-			return fmt.Errorf("error listing Athena Databases (%s): %w", region, err)
+			return fmt.Errorf("error listing Athena Data Catalogs (%s): %w", region, err)
 		}
 
-		for _, v := range page.DatabaseList {
+		for _, v := range page.DataCatalogsSummary {
+			catalogName := aws.ToString(v.CatalogName)
+			input := &athena.ListDatabasesInput{
+				CatalogName: aws.String(catalogName),
+			}
+
+			pages := athena.NewListDatabasesPaginator(conn, input)
+			for pages.HasMorePages() {
+				page, err := pages.NextPage(ctx)
+
+				if err != nil {
+					continue
+				}
+
+				for _, v := range page.DatabaseList {
+					name := aws.ToString(v.Name)
+
+					if name == "default" {
+						log.Printf("[INFO] Skipping Athena Database %s", name)
+						continue
+					}
+
+					r := resourceDatabase()
+					d := r.Data(nil)
+					d.SetId(name)
+					d.Set(names.AttrForceDestroy, true)
+
+					sweepResources = append(sweepResources, sweep.NewSweepResource(r, d, client))
+				}
+			}
+		}
+	}
+
+	err = sweep.SweepOrchestrator(ctx, sweepResources)
+
+	if err != nil {
+		return fmt.Errorf("error sweeping Athena Databases (%s): %w", region, err)
+	}
+
+	return nil
+}
+
+func sweepWorkGroups(region string) error {
+	ctx := sweep.Context(region)
+	client, err := sweep.SharedRegionalSweepClient(ctx, region)
+	if err != nil {
+		return fmt.Errorf("error getting client: %s", err)
+	}
+	conn := client.AthenaClient(ctx)
+	input := &athena.ListWorkGroupsInput{}
+	sweepResources := make([]sweep.Sweepable, 0)
+
+	pages := athena.NewListWorkGroupsPaginator(conn, input)
+	for pages.HasMorePages() {
+		page, err := pages.NextPage(ctx)
+
+		if awsv2.SkipSweepError(err) {
+			log.Printf("[WARN] Skipping Athena WorkGroup sweep for %s: %s", region, err)
+			return nil
+		}
+
+		if err != nil {
+			return fmt.Errorf("error listing Athena WorkGroups (%s): %w", region, err)
+		}
+
+		for _, v := range page.WorkGroups {
 			name := aws.ToString(v.Name)
-			if name == "default" {
+
+			if name == "primary" {
+				log.Printf("[INFO] Skipping Athena WorkGroup %s", name)
 				continue
 			}
-			r := resourceDatabase()
+
+			r := resourceWorkGroup()
 			d := r.Data(nil)
 			d.SetId(name)
 			d.Set(names.AttrForceDestroy, true)
@@ -64,7 +194,7 @@ func sweepDatabases(region string) error {
 	err = sweep.SweepOrchestrator(ctx, sweepResources)
 
 	if err != nil {
-		return fmt.Errorf("error sweeping Athena Databases (%s): %w", region, err)
+		return fmt.Errorf("error sweeping Athena WorkGroups (%s): %w", region, err)
 	}
 
 	return nil

diff --git a/internal/service/athena/workgroup.go b/internal/service/athena/workgroup.go
@@ -301,6 +301,7 @@ func resourceWorkGroupDelete(ctx context.Context, d *schema.ResourceData, meta i
 		input.RecursiveDeleteOption = aws.Bool(v.(bool))
 	}
 
+	log.Printf("[DEBUG] Deleting Athena WorkGroup (%s)", d.Id())
 	_, err := conn.DeleteWorkGroup(ctx, input)
 
 	if err != nil {