Skip to content

Commit

Permalink
Merge pull request #41136 from hashicorp/s-bedrockagent
Browse files Browse the repository at this point in the history
bedrockagent: Add sweepers
  • Loading branch information
ewbankkit authored Jan 29, 2025
2 parents b7c675d + 1ef8718 commit f95e046
Show file tree
Hide file tree
Showing 3 changed files with 106 additions and 1 deletion.
2 changes: 1 addition & 1 deletion internal/service/bedrockagent/agent.go
Original file line number Diff line number Diff line change
Expand Up @@ -389,7 +389,7 @@ func (r *agentResource) Delete(ctx context.Context, request resource.DeleteReque

agentID := data.ID.ValueString()
_, err := conn.DeleteAgent(ctx, &bedrockagent.DeleteAgentInput{
AgentId: fwflex.StringFromFramework(ctx, data.AgentID),
AgentId: aws.String(agentID),
SkipResourceInUseCheck: fwflex.BoolValueFromFramework(ctx, data.SkipResourceInUseCheck),
})

Expand Down
103 changes: 103 additions & 0 deletions internal/service/bedrockagent/sweep.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0

package bedrockagent

import (
"context"

"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/service/bedrockagent"
"github.com/hashicorp/terraform-provider-aws/internal/conns"
"github.com/hashicorp/terraform-provider-aws/internal/sweep"
"github.com/hashicorp/terraform-provider-aws/internal/sweep/awsv2"
"github.com/hashicorp/terraform-provider-aws/internal/sweep/framework"
"github.com/hashicorp/terraform-provider-aws/names"
)

func RegisterSweepers() {
awsv2.Register("aws_bedrockagent_agent", sweepAgents)
awsv2.Register("aws_bedrockagent_data_source", sweepDataSources)
awsv2.Register("aws_bedrockagent_knowledge_base", sweepKnowledgeBases, "aws_bedrockagent_agent", "aws_bedrockagent_data_source")
}

func sweepAgents(ctx context.Context, client *conns.AWSClient) ([]sweep.Sweepable, error) {
input := &bedrockagent.ListAgentsInput{}
conn := client.BedrockAgentClient(ctx)
sweepResources := make([]sweep.Sweepable, 0)

pages := bedrockagent.NewListAgentsPaginator(conn, input)
for pages.HasMorePages() {
page, err := pages.NextPage(ctx)

if err != nil {
return nil, err
}

for _, v := range page.AgentSummaries {
sweepResources = append(sweepResources, framework.NewSweepResource(newAgentResource, client,
framework.NewAttribute(names.AttrID, aws.ToString(v.AgentId)), framework.NewAttribute("skip_resource_in_use_check", true)))
}
}

return sweepResources, nil
}

func sweepDataSources(ctx context.Context, client *conns.AWSClient) ([]sweep.Sweepable, error) {
input := &bedrockagent.ListKnowledgeBasesInput{}
conn := client.BedrockAgentClient(ctx)
sweepResources := make([]sweep.Sweepable, 0)

pages := bedrockagent.NewListKnowledgeBasesPaginator(conn, input)
for pages.HasMorePages() {
page, err := pages.NextPage(ctx)

if err != nil {
return nil, err
}

for _, v := range page.KnowledgeBaseSummaries {
input := &bedrockagent.ListDataSourcesInput{
KnowledgeBaseId: v.KnowledgeBaseId,
}

pages := bedrockagent.NewListDataSourcesPaginator(conn, input)
for pages.HasMorePages() {
page, err := pages.NextPage(ctx)

if err != nil {
return nil, err
}

for _, v := range page.DataSourceSummaries {
sweepResources = append(sweepResources, framework.NewSweepResource(newDataSourceResource, client,
framework.NewAttribute("data_source_id", aws.ToString(v.DataSourceId)), framework.NewAttribute("knowledge_base_id", aws.ToString(v.KnowledgeBaseId))))
}
}
}
}

return sweepResources, nil
}

func sweepKnowledgeBases(ctx context.Context, client *conns.AWSClient) ([]sweep.Sweepable, error) {
input := &bedrockagent.ListKnowledgeBasesInput{}
conn := client.BedrockAgentClient(ctx)
sweepResources := make([]sweep.Sweepable, 0)

pages := bedrockagent.NewListKnowledgeBasesPaginator(conn, input)
for pages.HasMorePages() {
page, err := pages.NextPage(ctx)

if err != nil {
return nil, err
}

for _, v := range page.KnowledgeBaseSummaries {
sweepResources = append(sweepResources, framework.NewSweepResource(newKnowledgeBaseResource, client,
framework.NewAttribute(names.AttrID, aws.ToString(v.KnowledgeBaseId))))
}
}

return sweepResources, nil
}
2 changes: 2 additions & 0 deletions internal/sweep/register_gen_test.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit f95e046

Please sign in to comment.