Skip to content

Commit

Permalink
Add account alias support
Browse files Browse the repository at this point in the history
  • Loading branch information
@kgeckhart
kgeckhart committed Jul 1, 2024
1 parent be5080d commit b79c043
Show file tree
Hide file tree
Showing 2 changed files with 150 additions and 41 deletions.
44 changes: 30 additions & 14 deletions pkg/job/scraper.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,19 +51,34 @@ var (
CloudWatchCollectionErr ErrorType = "Failed to gather cloudwatch metrics for job"
)

type Account struct {
ID string
Alias string
}

func (s Scraper) Scrape(ctx context.Context) ([]model.TaggedResourceResult, []model.CloudwatchMetricResult, []Error) {
// Setup so we only do one GetAccount call per region + role combo when running jobs
roleRegionToAccount := map[model.Role]map[string]func() (string, error){}
roleRegionToAccount := map[model.Role]map[string]func() (Account, error){}
jobConfigVisitor(s.jobsCfg, func(_ any, role model.Role, region string) {
if _, exists := roleRegionToAccount[role]; !exists {
roleRegionToAccount[role] = map[string]func() (string, error){}
roleRegionToAccount[role] = map[string]func() (Account, error){}
}
roleRegionToAccount[role][region] = sync.OnceValues[string, error](func() (string, error) {
accountID, err := s.runnerFactory.GetAccountClient(region, role).GetAccount(ctx)
roleRegionToAccount[role][region] = sync.OnceValues[Account, error](func() (Account, error) {
client := s.runnerFactory.GetAccountClient(region, role)
accountID, err := client.GetAccount(ctx)
if err != nil {
return Account{}, fmt.Errorf("failed to get Account: %w", err)
}
a := Account{
ID: accountID,
}
accountAlias, err := client.GetAccountAlias(ctx)
if err != nil {
return "", fmt.Errorf("failed to get Account: %w", err)
s.logger.Warn("Failed to get optional account alias from account", "err", err, "account_id", accountID)
} else {
a.Alias = accountAlias
}
return accountID, nil
return a, nil
})
})

Expand Down Expand Up @@ -92,16 +107,16 @@ func (s Scraper) Scrape(ctx context.Context) ([]model.TaggedResourceResult, []mo
}
jobLogger := s.logger.With("namespace", jobContext.Namespace, "region", jobContext.Region, "arn", jobContext.RoleARN)

accountID, err := roleRegionToAccount[role][region]()
account, err := roleRegionToAccount[role][region]()
if err != nil {
jobError := NewError(jobContext, AccountErr, err)
mux.Lock()
jobErrors = append(jobErrors, jobError)
mux.Unlock()
return
}
jobContext.AccountID = accountID
jobLogger = jobLogger.With("account", jobContext.AccountID)
jobContext.Account = account
jobLogger = jobLogger.With("account_id", jobContext.Account.ID)

var jobToRun cloudwatchrunner.Job
jobAction(jobLogger, job,
Expand Down Expand Up @@ -211,17 +226,18 @@ func jobAction(logger logging.Logger, job any, discovery func(job model.Discover
// This makes it easier to track the data additively and morph it to the final shape necessary be it a model.ScrapeContext
// or an Error. It's an exported type for tests but is not part of the public interface
type JobContext struct { //nolint:revive
AccountID string
Account Account
Namespace string
Region string
RoleARN string
}

func (jc JobContext) ToScrapeContext(customTags []model.Tag) *model.ScrapeContext {
return &model.ScrapeContext{
AccountID: jc.AccountID,
Region: jc.Region,
CustomTags: customTags,
AccountID: jc.Account.ID,
Region: jc.Region,
CustomTags: customTags,
AccountAlias: jc.Account.Alias,
}
}

Expand All @@ -241,7 +257,7 @@ func NewError(context JobContext, errorType ErrorType, err error) Error {

func (e Error) ToLoggerKeyVals() []interface{} {
return []interface{}{
"account_id", e.AccountID,
"account_id", e.Account.ID,
"namespace", e.Namespace,
"region", e.Region,
"role_arn", e.RoleARN,
Expand Down
147 changes: 120 additions & 27 deletions pkg/job/scraper_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,14 @@ import (
)

type testRunnerFactory struct {
GetAccountFunc func() (string, error)
MetadataRunFunc func(ctx context.Context, region string, job model.DiscoveryJob) ([]*model.TaggedResource, error)
CloudwatchRunFunc func(ctx context.Context, job cloudwatchrunner.Job) ([]*model.CloudwatchData, error)
GetAccountAliasFunc func() (string, error)
GetAccountFunc func() (string, error)
MetadataRunFunc func(ctx context.Context, region string, job model.DiscoveryJob) ([]*model.TaggedResource, error)
CloudwatchRunFunc func(ctx context.Context, job cloudwatchrunner.Job) ([]*model.CloudwatchData, error)
}

func (t *testRunnerFactory) GetAccountAlias(context.Context) (string, error) {
return t.GetAccountAliasFunc()
}

func (t *testRunnerFactory) GetAccount(context.Context) (string, error) {
Expand Down Expand Up @@ -63,14 +68,15 @@ func (t testCloudwatchRunner) Run(ctx context.Context) ([]*model.CloudwatchData,

func TestScrapeRunner_Run(t *testing.T) {
tests := []struct {
name string
jobsCfg model.JobsConfig
getAccountFunc func() (string, error)
metadataRunFunc func(ctx context.Context, region string, job model.DiscoveryJob) ([]*model.TaggedResource, error)
cloudwatchRunFunc func(ctx context.Context, job cloudwatchrunner.Job) ([]*model.CloudwatchData, error)
expectedResources []model.TaggedResourceResult
expectedMetrics []model.CloudwatchMetricResult
expectedErrs []job.Error
name string
jobsCfg model.JobsConfig
getAccountFunc func() (string, error)
getAccountAliasFunc func() (string, error)
metadataRunFunc func(ctx context.Context, region string, job model.DiscoveryJob) ([]*model.TaggedResource, error)
cloudwatchRunFunc func(ctx context.Context, job cloudwatchrunner.Job) ([]*model.CloudwatchData, error)
expectedResources []model.TaggedResourceResult
expectedMetrics []model.CloudwatchMetricResult
expectedErrs []job.Error
}{
{
name: "can run a discovery job",
Expand All @@ -88,6 +94,9 @@ func TestScrapeRunner_Run(t *testing.T) {
getAccountFunc: func() (string, error) {
return "aws-account-1", nil
},
getAccountAliasFunc: func() (string, error) {
return "my-aws-account", nil
},
metadataRunFunc: func(_ context.Context, _ string, _ model.DiscoveryJob) ([]*model.TaggedResource, error) {
return []*model.TaggedResource{{
ARN: "resource-1", Namespace: "aws-namespace", Region: "us-east-1", Tags: []model.Tag{{Key: "tag1", Value: "value1"}},
Expand All @@ -107,15 +116,15 @@ func TestScrapeRunner_Run(t *testing.T) {
},
expectedResources: []model.TaggedResourceResult{
{
Context: &model.ScrapeContext{Region: "us-east-1", AccountID: "aws-account-1"},
Context: &model.ScrapeContext{Region: "us-east-1", AccountID: "aws-account-1", AccountAlias: "my-aws-account"},
Data: []*model.TaggedResource{
{ARN: "resource-1", Namespace: "aws-namespace", Region: "us-east-1", Tags: []model.Tag{{Key: "tag1", Value: "value1"}}},
},
},
},
expectedMetrics: []model.CloudwatchMetricResult{
{
Context: &model.ScrapeContext{Region: "us-east-1", AccountID: "aws-account-1"},
Context: &model.ScrapeContext{Region: "us-east-1", AccountID: "aws-account-1", AccountAlias: "my-aws-account"},
Data: []*model.CloudwatchData{
{
MetricName: "metric-1",
Expand Down Expand Up @@ -146,6 +155,9 @@ func TestScrapeRunner_Run(t *testing.T) {
getAccountFunc: func() (string, error) {
return "aws-account-1", nil
},
getAccountAliasFunc: func() (string, error) {
return "my-aws-account", nil
},
cloudwatchRunFunc: func(_ context.Context, _ cloudwatchrunner.Job) ([]*model.CloudwatchData, error) {
return []*model.CloudwatchData{
{
Expand All @@ -159,7 +171,7 @@ func TestScrapeRunner_Run(t *testing.T) {
},
expectedMetrics: []model.CloudwatchMetricResult{
{
Context: &model.ScrapeContext{Region: "us-east-2", AccountID: "aws-account-1"},
Context: &model.ScrapeContext{Region: "us-east-2", AccountID: "aws-account-1", AccountAlias: "my-aws-account"},
Data: []*model.CloudwatchData{
{
MetricName: "metric-2",
Expand Down Expand Up @@ -198,6 +210,9 @@ func TestScrapeRunner_Run(t *testing.T) {
getAccountFunc: func() (string, error) {
return "aws-account-1", nil
},
getAccountAliasFunc: func() (string, error) {
return "my-aws-account", nil
},
metadataRunFunc: func(_ context.Context, _ string, _ model.DiscoveryJob) ([]*model.TaggedResource, error) {
return []*model.TaggedResource{{
ARN: "resource-1", Namespace: "aws-namespace", Region: "us-east-1", Tags: []model.Tag{{Key: "tag1", Value: "value1"}},
Expand Down Expand Up @@ -228,15 +243,15 @@ func TestScrapeRunner_Run(t *testing.T) {
},
expectedResources: []model.TaggedResourceResult{
{
Context: &model.ScrapeContext{Region: "us-east-1", AccountID: "aws-account-1"},
Context: &model.ScrapeContext{Region: "us-east-1", AccountID: "aws-account-1", AccountAlias: "my-aws-account"},
Data: []*model.TaggedResource{
{ARN: "resource-1", Namespace: "aws-namespace", Region: "us-east-1", Tags: []model.Tag{{Key: "tag1", Value: "value1"}}},
},
},
},
expectedMetrics: []model.CloudwatchMetricResult{
{
Context: &model.ScrapeContext{Region: "us-east-1", AccountID: "aws-account-1"},
Context: &model.ScrapeContext{Region: "us-east-1", AccountID: "aws-account-1", AccountAlias: "my-aws-account"},
Data: []*model.CloudwatchData{
{
MetricName: "metric-1",
Expand All @@ -249,7 +264,7 @@ func TestScrapeRunner_Run(t *testing.T) {
},
},
{
Context: &model.ScrapeContext{Region: "us-east-2", AccountID: "aws-account-1"},
Context: &model.ScrapeContext{Region: "us-east-2", AccountID: "aws-account-1", AccountAlias: "my-aws-account"},
Data: []*model.CloudwatchData{
{
MetricName: "metric-2",
Expand Down Expand Up @@ -289,8 +304,66 @@ func TestScrapeRunner_Run(t *testing.T) {
return "", errors.New("failed to get account")
},
expectedErrs: []job.Error{
{JobContext: job.JobContext{AccountID: "", Namespace: "aws-namespace", Region: "us-east-1", RoleARN: "aws-arn-1"}, ErrorType: job.AccountErr},
{JobContext: job.JobContext{AccountID: "", Namespace: "custom-namespace", Region: "us-east-2", RoleARN: "aws-arn-2"}, ErrorType: job.AccountErr},
{JobContext: job.JobContext{Account: job.Account{}, Namespace: "aws-namespace", Region: "us-east-1", RoleARN: "aws-arn-1"}, ErrorType: job.AccountErr},
{JobContext: job.JobContext{Account: job.Account{}, Namespace: "custom-namespace", Region: "us-east-2", RoleARN: "aws-arn-2"}, ErrorType: job.AccountErr},
},
},
{
name: "ignores errors from GetAccountAlias",
jobsCfg: model.JobsConfig{
DiscoveryJobs: []model.DiscoveryJob{
{
Regions: []string{"us-east-1"},
Type: "aws-namespace",
Roles: []model.Role{
{RoleArn: "aws-arn-1", ExternalID: "external-id-1"},
},
},
},
},
getAccountFunc: func() (string, error) {
return "aws-account-1", nil
},
getAccountAliasFunc: func() (string, error) { return "", errors.New("No alias here") },
metadataRunFunc: func(_ context.Context, _ string, _ model.DiscoveryJob) ([]*model.TaggedResource, error) {
return []*model.TaggedResource{{
ARN: "resource-1", Namespace: "aws-namespace", Region: "us-east-1", Tags: []model.Tag{{Key: "tag1", Value: "value1"}},
}}, nil
},
cloudwatchRunFunc: func(_ context.Context, _ cloudwatchrunner.Job) ([]*model.CloudwatchData, error) {
return []*model.CloudwatchData{
{
MetricName: "metric-1",
ResourceName: "resource-1",
Namespace: "aws-namespace",
Tags: []model.Tag{{Key: "tag1", Value: "value1"}},
Dimensions: []model.Dimension{{Name: "dimension1", Value: "value1"}},
GetMetricDataResult: &model.GetMetricDataResult{Statistic: "Maximum", Datapoint: aws.Float64(1.0), Timestamp: time.Time{}},
},
}, nil
},
expectedResources: []model.TaggedResourceResult{
{
Context: &model.ScrapeContext{Region: "us-east-1", AccountID: "aws-account-1", AccountAlias: ""},
Data: []*model.TaggedResource{
{ARN: "resource-1", Namespace: "aws-namespace", Region: "us-east-1", Tags: []model.Tag{{Key: "tag1", Value: "value1"}}},
},
},
},
expectedMetrics: []model.CloudwatchMetricResult{
{
Context: &model.ScrapeContext{Region: "us-east-1", AccountID: "aws-account-1", AccountAlias: ""},
Data: []*model.CloudwatchData{
{
MetricName: "metric-1",
ResourceName: "resource-1",
Namespace: "aws-namespace",
Tags: []model.Tag{{Key: "tag1", Value: "value1"}},
Dimensions: []model.Dimension{{Name: "dimension1", Value: "value1"}},
GetMetricDataResult: &model.GetMetricDataResult{Statistic: "Maximum", Datapoint: aws.Float64(1.0), Timestamp: time.Time{}},
},
},
},
},
},
{
Expand Down Expand Up @@ -319,6 +392,9 @@ func TestScrapeRunner_Run(t *testing.T) {
getAccountFunc: func() (string, error) {
return "aws-account-1", nil
},
getAccountAliasFunc: func() (string, error) {
return "my-aws-account", nil
},
metadataRunFunc: func(_ context.Context, _ string, _ model.DiscoveryJob) ([]*model.TaggedResource, error) {
return nil, errors.New("I failed you")
},
Expand All @@ -335,7 +411,7 @@ func TestScrapeRunner_Run(t *testing.T) {
},
expectedMetrics: []model.CloudwatchMetricResult{
{
Context: &model.ScrapeContext{Region: "us-east-2", AccountID: "aws-account-1"},
Context: &model.ScrapeContext{Region: "us-east-2", AccountID: "aws-account-1", AccountAlias: "my-aws-account"},
Data: []*model.CloudwatchData{
{
MetricName: "metric-2",
Expand All @@ -348,7 +424,14 @@ func TestScrapeRunner_Run(t *testing.T) {
},
},
expectedErrs: []job.Error{
{JobContext: job.JobContext{AccountID: "aws-account-1", Namespace: "aws-namespace", Region: "us-east-1", RoleARN: "aws-arn-1"}, ErrorType: job.ResourceMetadataErr},
{
JobContext: job.JobContext{
Account: job.Account{ID: "aws-account-1", Alias: "my-aws-account"},
Namespace: "aws-namespace",
Region: "us-east-1",
RoleARN: "aws-arn-1"},
ErrorType: job.ResourceMetadataErr,
},
},
},
{
Expand Down Expand Up @@ -377,6 +460,9 @@ func TestScrapeRunner_Run(t *testing.T) {
getAccountFunc: func() (string, error) {
return "aws-account-1", nil
},
getAccountAliasFunc: func() (string, error) {
return "my-aws-account", nil
},
metadataRunFunc: func(_ context.Context, _ string, _ model.DiscoveryJob) ([]*model.TaggedResource, error) {
return []*model.TaggedResource{{
ARN: "resource-1", Namespace: "aws-namespace", Region: "us-east-1", Tags: []model.Tag{{Key: "tag1", Value: "value1"}},
Expand All @@ -399,15 +485,15 @@ func TestScrapeRunner_Run(t *testing.T) {
},
expectedResources: []model.TaggedResourceResult{
{
Context: &model.ScrapeContext{Region: "us-east-1", AccountID: "aws-account-1"},
Context: &model.ScrapeContext{Region: "us-east-1", AccountID: "aws-account-1", AccountAlias: "my-aws-account"},
Data: []*model.TaggedResource{
{ARN: "resource-1", Namespace: "aws-namespace", Region: "us-east-1", Tags: []model.Tag{{Key: "tag1", Value: "value1"}}},
},
},
},
expectedMetrics: []model.CloudwatchMetricResult{
{
Context: &model.ScrapeContext{Region: "us-east-1", AccountID: "aws-account-1"},
Context: &model.ScrapeContext{Region: "us-east-1", AccountID: "aws-account-1", AccountAlias: "my-aws-account"},
Data: []*model.CloudwatchData{
{
MetricName: "metric-1",
Expand All @@ -421,16 +507,23 @@ func TestScrapeRunner_Run(t *testing.T) {
},
},
expectedErrs: []job.Error{
{JobContext: job.JobContext{AccountID: "aws-account-1", Namespace: "custom-namespace", Region: "us-east-2", RoleARN: "aws-arn-2"}, ErrorType: job.CloudWatchCollectionErr},
{
JobContext: job.JobContext{
Account: job.Account{ID: "aws-account-1", Alias: "my-aws-account"},
Namespace: "custom-namespace",
Region: "us-east-2",
RoleARN: "aws-arn-2"},
ErrorType: job.CloudWatchCollectionErr},
},
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
rf := testRunnerFactory{
GetAccountFunc: tc.getAccountFunc,
MetadataRunFunc: tc.metadataRunFunc,
CloudwatchRunFunc: tc.cloudwatchRunFunc,
GetAccountFunc: tc.getAccountFunc,
GetAccountAliasFunc: tc.getAccountAliasFunc,
MetadataRunFunc: tc.metadataRunFunc,
CloudwatchRunFunc: tc.cloudwatchRunFunc,
}
sr := job.NewScraper(logging.NewLogger("", true), tc.jobsCfg, &rf)
resources, metrics, errs := sr.Scrape(context.Background())
Expand Down

0 comments on commit b79c043

Please sign in to comment.