feat(nrawssdk): Accept NEW_RELIC_CLOUD_AWS_ACCOUNT_ID as env var

Author: cade-conklin

v3/integrations/nrawssdk-v2/example/main.go

@@ -38,18 +38,14 @@ func main() {
 	// Start recording a New Relic transaction
 	txn := app.StartTransaction("My sample transaction")
 
-	ctx := context.Background()
+	ctx := newrelic.NewContext(context.Background(), txn)
 
 	awsConfig, err := config.LoadDefaultConfig(ctx, func(awsConfig *config.LoadOptions) error {
 		// Instrument all new AWS clients with New Relic
-
 		return nil
 	})
-	creds, err := awsConfig.Credentials.Retrieve(ctx)
-	if err != nil {
-		log.Println("Warning couldn't get flags")
-	}
-	nrawssdk.AppendMiddlewares(&awsConfig.APIOptions, nil, creds)
+
+	nrawssdk.NRAppendMiddlewares(&awsConfig.APIOptions, ctx, awsConfig)
 	if err != nil {
 		log.Fatal(err)
 	}

v3/integrations/nrawssdk-v2/nrawssdk.go

@@ -46,8 +46,8 @@ import (
 )
 
 type nrMiddleware struct {
-	txn   *newrelic.Transaction
-	creds aws.Credentials
+	txn       *newrelic.Transaction
+	accountID string
 }
 
 type contextKey string
@@ -77,14 +77,7 @@ func (m nrMiddleware) deserializeMiddleware(stack *smithymiddle.Stack) error {
 		serviceName := awsmiddle.GetServiceID(ctx)
 		operation := awsmiddle.GetOperationName(ctx)
 		region := awsmiddle.GetRegion(ctx)
-
-		creds := awsmiddle.GetSigningCredentials(ctx)
-		accountID, err := awssupport.AWSAccountIdFromAWSAccessKey(creds)
-		if err != nil {
-			accountID = ""
-			fmt.Println(err.Error())
-		}
-
+		accountID := m.accountID
 		var segment endable
 
 		if serviceName == "dynamodb" || serviceName == "DynamoDB" {
@@ -138,7 +131,9 @@ func (m nrMiddleware) deserializeMiddleware(stack *smithymiddle.Stack) error {
 				integrationsupport.AddAgentSpanAttribute(txn, newrelic.AttributeAWSElastSearchDomainEndpoint, httpRequest.URL.String()) // this way I don't have to pull it out of context
 			}
 			// Set additional span attributes
+
 			integrationsupport.AddAgentSpanAttribute(txn, newrelic.AttributeCloudAccountID, accountID) // setting account ID here, why do we only do this if it is an SQS service?
+
 			integrationsupport.AddAgentSpanAttribute(txn,
 				newrelic.AttributeResponseCode, strconv.Itoa(response.StatusCode))
 			integrationsupport.AddAgentSpanAttribute(txn,
@@ -161,7 +156,6 @@ func (m nrMiddleware) serializeMiddleware(stack *middleware.Stack) error {
 		ctx context.Context, in middleware.InitializeInput, next middleware.InitializeHandler) (
 		out middleware.InitializeOutput, metadata middleware.Metadata, err error) {
 		serviceName := awsmiddle.GetServiceID(ctx)
-		ctx = awsmiddle.SetSigningCredentials(ctx, m.creds)
 		switch serviceName {
 		case "dynamodb", "DynamoDB":
 			ctx = context.WithValue(ctx, dynamodbInputKey, dynamoDBInputFromMiddlewareInput(in))
@@ -229,8 +223,31 @@ func (m nrMiddleware) serializeMiddleware(stack *middleware.Stack) error {
 //	if err != nil {
 //		log.Fatal(err)
 //	}
-func AppendMiddlewares(apiOptions *[]func(*smithymiddle.Stack) error, txn *newrelic.Transaction, creds aws.Credentials) {
-	m := nrMiddleware{txn: txn, creds: creds}
+func AppendMiddlewares(apiOptions *[]func(*smithymiddle.Stack) error, txn *newrelic.Transaction) {
+	m := nrMiddleware{txn: txn}
+	*apiOptions = append(*apiOptions, m.deserializeMiddleware)
+	*apiOptions = append(*apiOptions, m.serializeMiddleware)
+}
+
+func NRAppendMiddlewares(apiOptions *[]func(*smithymiddle.Stack) error, ctx context.Context, awsConfig aws.Config) {
+	txn := newrelic.FromContext(ctx)
+
+	creds, err := awsConfig.Credentials.Retrieve(ctx)
+	if err != nil {
+		fmt.Println("error: Couldn't get AWS Credentials")
+	}
+
+	cfg, ok := txn.Application().Config()
+	m := nrMiddleware{txn: txn}
+
+	if ok {
+		accountID, err := ResolveAWSCredentials(cfg, creds)
+		if err != nil {
+			fmt.Println("error: Couldn't resolve AWS credentials")
+		}
+		m.accountID = accountID
+	}
+
 	*apiOptions = append(*apiOptions, m.deserializeMiddleware)
 	*apiOptions = append(*apiOptions, m.serializeMiddleware)
 }
@@ -291,3 +308,18 @@ func dynamoDBInputFromMiddlewareInput(in middleware.InitializeInput) dynamodbInp
 		return dynamodbInput{}
 	}
 }
+
+func ResolveAWSCredentials(cfg newrelic.Config, creds aws.Credentials) (string, error) {
+
+	accountID, err := awssupport.AWSAccountIdFromAWSAccessKey(creds)
+	if err != nil {
+		return "", err
+	}
+
+	// if the accountID that comes back is different than what is set in the env
+	// then we should use that accountID since it is accurate
+	if accountID != "" && accountID != cfg.CloudAWS.AccountID {
+		return accountID, nil
+	}
+	return cfg.CloudAWS.AccountID, nil
+}

v3/integrations/nrawssdk-v2/nrawssdk_test.go

@@ -64,7 +64,7 @@ var fakeCreds = func() interface{} {
 
 func newConfig(ctx context.Context, txn *newrelic.Transaction) aws.Config {
 	cfg, _ := config.LoadDefaultConfig(ctx, func(o *config.LoadOptions) error {
-		AppendMiddlewares(&o.APIOptions, txn, aws.Credentials{})
+		AppendMiddlewares(&o.APIOptions, txn)
 		return nil
 	})
 	cfg.Credentials = fakeCreds.(aws.CredentialsProvider)

v3/newrelic/config.go

@@ -93,6 +93,10 @@ type Config struct {
 		MaxSamplesStored int
 	}
 
+	CloudAWS struct {
+		AccountID string
+	}
+
 	// ErrorCollector controls the capture of errors.
 	ErrorCollector struct {
 		// Enabled controls whether errors are captured.  This setting

v3/newrelic/config_options.go

@@ -624,6 +624,9 @@ func configFromEnvironment(getenv func(string) string) ConfigOption {
 		// Error Collector Env Variables
 		assignInt(&cfg.ErrorCollector.MaxSamplesStored, "NEW_RELIC_ERROR_COLLECTOR_MAX_EVENT_SAMPLES_STORED", maxErrorEvents)
 
+		// AWS Env Variables
+		assignString(&cfg.CloudAWS.AccountID, "NEW_RELIC_CLOUD_AWS_ACCOUNT_ID")
+
 		if env := getenv("NEW_RELIC_LABELS"); env != "" {
 			labels, err := getLabels(getenv("NEW_RELIC_LABELS"))
 			if err != nil {