Skip to content

Commit bb287ea

Browse files
sstoopssstoops
authored
feat(bedrock): Add support for Agent Alias and Flow Alias (ekristen#664)
* feat(bedrock-agent-alias): add BedrockAgentAlias module * feat(bedrock-flow-alias): add BedrockFlowAlias module * fix(bedrock-agent-alias): adding Resource to BedrockAgentAlias registration * fix(bedrock-agent-alias): switching pointer reference to "r" * feat(bedrock-agent-alias): adding dependency on BedrockAgent * fix(bedrock-flow-alias): adding Resource to BedrockFlowAlias registration * refactor(bedrock-flow-alias): switching pointer ref and moved utility function * refactor(bedrock-agent-alias): move utility function to end of file * fix: linting adjustments * fix: linting adjustments
1 parent 21a7e67 commit bb287ea

File tree

2 files changed

+240
-0
lines changed

2 files changed

+240
-0
lines changed

‎resources/bedrock-agent-alias.go‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
package resources
2+
3+
import (
4+
"context"
5+
6+
"github.com/aws/aws-sdk-go/aws"
7+
"github.com/aws/aws-sdk-go/service/bedrockagent"
8+
9+
"github.com/ekristen/libnuke/pkg/registry"
10+
"github.com/ekristen/libnuke/pkg/resource"
11+
"github.com/ekristen/libnuke/pkg/types"
12+
13+
"github.com/ekristen/aws-nuke/v3/pkg/nuke"
14+
)
15+
16+
const BedrockAgentAliasResource = "BedrockAgentAlias"
17+
18+
func init() {
19+
registry.Register(&registry.Registration{
20+
Name: BedrockAgentAliasResource,
21+
Scope: nuke.Account,
22+
Lister: &BedrockAgentAliasLister{},
23+
Resource: &BedrockAgentAlias{},
24+
DependsOn: []string{
25+
BedrockAgentResource,
26+
},
27+
})
28+
}
29+
30+
type BedrockAgentAliasLister struct{}
31+
32+
func (l *BedrockAgentAliasLister) List(_ context.Context, o interface{}) ([]resource.Resource, error) {
33+
opts := o.(*nuke.ListerOpts)
34+
35+
svc := bedrockagent.New(opts.Session)
36+
resources := make([]resource.Resource, 0)
37+
38+
agentIDs, err := ListBedrockAgentIds(svc)
39+
if err != nil {
40+
return nil, err
41+
}
42+
43+
for _, agentID := range agentIDs {
44+
params := &bedrockagent.ListAgentAliasesInput{
45+
MaxResults: aws.Int64(100),
46+
AgentId: aws.String(agentID),
47+
}
48+
for {
49+
output, err := svc.ListAgentAliases(params)
50+
if err != nil {
51+
return nil, err
52+
}
53+
54+
for _, agentAliasInfo := range output.AgentAliasSummaries {
55+
resources = append(resources, &BedrockAgentAlias{
56+
svc: svc,
57+
AgentID: aws.String(agentID),
58+
AgentAliasName: agentAliasInfo.AgentAliasName,
59+
AgentAliasID: agentAliasInfo.AgentAliasId,
60+
})
61+
}
62+
63+
if output.NextToken == nil {
64+
break
65+
}
66+
params.NextToken = output.NextToken
67+
}
68+
}
69+
return resources, nil
70+
}
71+
72+
type BedrockAgentAlias struct {
73+
svc *bedrockagent.BedrockAgent
74+
AgentID *string
75+
AgentAliasID *string
76+
AgentAliasName *string
77+
}
78+
79+
func (r *BedrockAgentAlias) Properties() types.Properties {
80+
return types.NewPropertiesFromStruct(r)
81+
}
82+
83+
func (r *BedrockAgentAlias) Remove(_ context.Context) error {
84+
_, err := r.svc.DeleteAgentAlias(&bedrockagent.DeleteAgentAliasInput{
85+
AgentAliasId: r.AgentAliasID,
86+
AgentId: r.AgentID,
87+
})
88+
return err
89+
}
90+
91+
func (r *BedrockAgentAlias) String() string {
92+
return *r.AgentAliasName
93+
}
94+
95+
func ListBedrockAgentIds(svc *bedrockagent.BedrockAgent) ([]string, error) {
96+
agentIds := []string{}
97+
params := &bedrockagent.ListAgentsInput{
98+
MaxResults: aws.Int64(100),
99+
}
100+
for {
101+
output, err := svc.ListAgents(params)
102+
if err != nil {
103+
return nil, err
104+
}
105+
106+
for _, agent := range output.AgentSummaries {
107+
agentIds = append(agentIds, *agent.AgentId)
108+
}
109+
110+
if output.NextToken == nil {
111+
break
112+
}
113+
params.NextToken = output.NextToken
114+
}
115+
116+
return agentIds, nil
117+
}

‎resources/bedrock-flow-alias.go‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
package resources
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"strings"
7+
8+
"github.com/aws/aws-sdk-go/aws"
9+
"github.com/aws/aws-sdk-go/service/bedrockagent"
10+
11+
"github.com/ekristen/libnuke/pkg/registry"
12+
"github.com/ekristen/libnuke/pkg/resource"
13+
"github.com/ekristen/libnuke/pkg/types"
14+
15+
"github.com/ekristen/aws-nuke/v3/pkg/nuke"
16+
)
17+
18+
const BedrockFlowAliasResource = "BedrockFlowAlias"
19+
20+
func init() {
21+
registry.Register(&registry.Registration{
22+
Name: BedrockFlowAliasResource,
23+
Scope: nuke.Account,
24+
Lister: &BedrockFlowAliasLister{},
25+
Resource: &BedrockFlowAlias{},
26+
})
27+
}
28+
29+
type BedrockFlowAliasLister struct{}
30+
31+
func (l *BedrockFlowAliasLister) List(_ context.Context, o interface{}) ([]resource.Resource, error) {
32+
opts := o.(*nuke.ListerOpts)
33+
34+
svc := bedrockagent.New(opts.Session)
35+
resources := make([]resource.Resource, 0)
36+
37+
flowIDs, err := ListBedrockFlowIds(svc)
38+
if err != nil {
39+
return nil, err
40+
}
41+
42+
for _, flowID := range flowIDs {
43+
params := &bedrockagent.ListFlowAliasesInput{
44+
MaxResults: aws.Int64(100),
45+
FlowIdentifier: aws.String(flowID),
46+
}
47+
for {
48+
output, err := svc.ListFlowAliases(params)
49+
if err != nil {
50+
return nil, err
51+
}
52+
53+
for _, flowAliasInfo := range output.FlowAliasSummaries {
54+
resources = append(resources, &BedrockFlowAlias{
55+
svc: svc,
56+
FlowID: flowAliasInfo.FlowId,
57+
FlowAliasID: flowAliasInfo.Id,
58+
FlowAliasName: flowAliasInfo.Name,
59+
})
60+
}
61+
62+
if output.NextToken == nil {
63+
break
64+
}
65+
params.NextToken = output.NextToken
66+
}
67+
}
68+
return resources, nil
69+
}
70+
71+
type BedrockFlowAlias struct {
72+
svc *bedrockagent.BedrockAgent
73+
FlowID *string
74+
FlowAliasID *string
75+
FlowAliasName *string
76+
}
77+
78+
func (r *BedrockFlowAlias) Filter() error {
79+
if strings.HasPrefix(*r.FlowAliasName, "TSTALIASID") {
80+
return fmt.Errorf("cannot delete AWS managed Flow Alias")
81+
}
82+
return nil
83+
}
84+
85+
func (r *BedrockFlowAlias) Properties() types.Properties {
86+
return types.NewPropertiesFromStruct(r)
87+
}
88+
89+
func (r *BedrockFlowAlias) Remove(_ context.Context) error {
90+
_, err := r.svc.DeleteFlowAlias(&bedrockagent.DeleteFlowAliasInput{
91+
AliasIdentifier: r.FlowAliasID,
92+
FlowIdentifier: r.FlowID,
93+
})
94+
return err
95+
}
96+
97+
func (r *BedrockFlowAlias) String() string {
98+
return *r.FlowAliasName
99+
}
100+
101+
func ListBedrockFlowIds(svc *bedrockagent.BedrockAgent) ([]string, error) {
102+
flowIds := []string{}
103+
params := &bedrockagent.ListFlowsInput{
104+
MaxResults: aws.Int64(100),
105+
}
106+
for {
107+
output, err := svc.ListFlows(params)
108+
if err != nil {
109+
return nil, err
110+
}
111+
112+
for _, flow := range output.FlowSummaries {
113+
flowIds = append(flowIds, *flow.Id)
114+
}
115+
116+
if output.NextToken == nil {
117+
break
118+
}
119+
params.NextToken = output.NextToken
120+
}
121+
122+
return flowIds, nil
123+
}

0 commit comments

Comments
 (0)