Merge pull request #599 from Meroje/nrawssdk-v2

nrawssdk-v2: fix usage examples of nrawssdk.AppendMiddlewares
newrelic · Oct 10, 2024 · fe29069 · fe29069
2 parents 1c35a75 + 441b3a9
commit fe29069
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 21 deletions.
diff --git a/v3/integrations/nrawssdk-v2/example/main.go b/v3/integrations/nrawssdk-v2/example/main.go
@@ -10,7 +10,7 @@ import (
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/service/s3"
-	nraws "github.com/newrelic/go-agent/v3/integrations/nrawssdk-v2"
+	"github.com/newrelic/go-agent/v3/integrations/nrawssdk-v2"
 	"github.com/newrelic/go-agent/v3/newrelic"
 )
 
@@ -39,14 +39,15 @@ func main() {
 	txn := app.StartTransaction("My sample transaction")
 
 	ctx := context.Background()
-	awsConfig, err := config.LoadDefaultConfig(ctx)
+	awsConfig, err := config.LoadDefaultConfig(ctx, func(awsConfig *config.LoadOptions) error {
+		// Instrument all new AWS clients with New Relic
+		nrawssdk.AppendMiddlewares(&awsConfig.APIOptions, nil)
+		return nil
+	})
 	if err != nil {
 		log.Fatal(err)
 	}
 
-	// Instrument all new AWS clients with New Relic
-	nraws.AppendMiddlewares(&awsConfig.APIOptions, nil)
-
 	s3Client := s3.NewFromConfig(awsConfig)
 	output, err := s3Client.ListBuckets(ctx, nil)
 	if err != nil {

diff --git a/v3/integrations/nrawssdk-v2/nrawssdk.go b/v3/integrations/nrawssdk-v2/nrawssdk.go
@@ -120,24 +120,50 @@ func (m nrMiddleware) deserializeMiddleware(stack *smithymiddle.Stack) error {
 // To see segments and spans for all AWS invocations, call AppendMiddlewares
 // with the AWS Config `apiOptions` and provide nil for `txn`. For example:
 //
-//  awsConfig, err := config.LoadDefaultConfig(ctx)
-//  if err != nil {
-//      log.Fatal(err)
-//  }
-//  nraws.AppendMiddlewares(&awsConfig.APIOptions, nil)
+//	awsConfig, err := config.LoadDefaultConfig(ctx, func(o *config.LoadOptions) error {
+//		// Instrument all new AWS clients with New Relic
+//		nrawssdk.AppendMiddlewares(&o.APIOptions, nil)
+//		return nil
+//	})
+//	if err != nil {
+//		log.Fatal(err)
+//	}
 //
-// If do not want the transaction to be retrived from the context, you can
+// If do not want the transaction to be retrieved from the context, you can
 // explicitly set `txn`. For example:
 //
-//  awsConfig, err := config.LoadDefaultConfig(ctx)
-//  if err != nil {
-//      log.Fatal(err)
-//  }
+//	txn := loadNewRelicTransaction()
+//	awsConfig, err := config.LoadDefaultConfig(ctx, func(o *config.LoadOptions) error {
+//		// Instrument all new AWS clients with New Relic
+//		nrawssdk.AppendMiddlewares(&o.APIOptions, txn)
+//		return nil
+//	})
+//	if err != nil {
+//		log.Fatal(err)
+//	}
 //
-//  ...
+// The middleware can also be added later, per AWS service call using
+// the `optFns` parameter. For example:
 //
-//  txn := loadNewRelicTransaction()
-//  nraws.AppendMiddlewares(&awsConfig.APIOptions, txn)
+//	awsConfig, err := config.LoadDefaultConfig(ctx)
+//	if err != nil {
+//		log.Fatal(err)
+//	}
+//
+//	...
+//
+//	s3Client := s3.NewFromConfig(awsConfig)
+//
+//	...
+//
+//	txn := loadNewRelicTransaction()
+//	output, err := s3Client.ListBuckets(ctx, nil, func(o *s3.Options) error {
+//		nrawssdk.AppendMiddlewares(&o.APIOptions, txn)
+//		return nil
+//	})
+//	if err != nil {
+//		log.Fatal(err)
+//	}
 func AppendMiddlewares(apiOptions *[]func(*smithymiddle.Stack) error, txn *newrelic.Transaction) {
 	m := nrMiddleware{txn: txn}
 	*apiOptions = append(*apiOptions, m.deserializeMiddleware)

diff --git a/v3/integrations/nrawssdk-v2/nrawssdk_test.go b/v3/integrations/nrawssdk-v2/nrawssdk_test.go
@@ -60,15 +60,16 @@ var fakeCreds = func() interface{} {
 }()
 
 func newConfig(ctx context.Context, txn *newrelic.Transaction) aws.Config {
-	cfg, _ := config.LoadDefaultConfig(ctx)
+	cfg, _ := config.LoadDefaultConfig(ctx, func(o *config.LoadOptions) error {
+		AppendMiddlewares(&o.APIOptions, txn)
+		return nil
+	})
 	cfg.Credentials = fakeCreds.(aws.CredentialsProvider)
 	cfg.Region = awsRegion
 	cfg.HTTPClient = &http.Client{
 		Transport: &fakeTransport{},
 	}
 
-	AppendMiddlewares(&cfg.APIOptions, txn)
-
 	return cfg
 }