Add AttributeBuilder, which will extend the AttributeSetter contract …

…with output value. Thanks to that, users can instrument attributes based on the AWS SDK responses.

instrumentation/github.com/aws/aws-sdk-go-v2/otelaws/attributes.go

@@ -23,10 +23,10 @@ const (
 	AWSSystemVal string        = "aws-api"
 )
 
-var servicemap = map[string]AttributeSetter{
-	dynamodb.ServiceID: DynamoDBAttributeSetter,
-	sqs.ServiceID:      SQSAttributeSetter,
-	sns.ServiceID:      SNSAttributeSetter,
+var servicemap = map[string]AttributeBuilder{
+	dynamodb.ServiceID: DynamoDBAttributeBuilder,
+	sqs.ServiceID:      SQSAttributeBuilder,
+	sns.ServiceID:      SNSAttributeBuilder,
 }
 
 // SystemAttr return the AWS RPC system attribute.
@@ -56,11 +56,18 @@ func RequestIDAttr(requestID string) attribute.KeyValue {
 
 // DefaultAttributeSetter checks to see if there are service specific attributes available to set for the AWS service.
 // If there are service specific attributes available then they will be included.
+// Deprecated: Kept for backward compatibility, use DefaultAttributeBuilder instead. This will be removed in a future release.
 func DefaultAttributeSetter(ctx context.Context, in middleware.InitializeInput) []attribute.KeyValue {
+	return DefaultAttributeBuilder(ctx, in, middleware.InitializeOutput{})
+}
+
+// DefaultAttributeBuilder checks to see if there are service specific attributes available to set for the AWS service.
+// If there are service specific attributes available then they will be included.
+func DefaultAttributeBuilder(ctx context.Context, in middleware.InitializeInput, out middleware.InitializeOutput) []attribute.KeyValue {
 	serviceID := v2Middleware.GetServiceID(ctx)
 
 	if fn, ok := servicemap[serviceID]; ok {
-		return fn(ctx, in)
+		return fn(ctx, in, out)
 	}
 
 	return []attribute.KeyValue{}

instrumentation/github.com/aws/aws-sdk-go-v2/otelaws/aws.go

@@ -27,12 +27,16 @@ const (
 type spanTimestampKey struct{}
 
 // AttributeSetter returns an array of KeyValue pairs, it can be used to set custom attributes.
+// Deprecated: Kept for backward compatibility, use AttributeBuilder instead. This will be removed in a future release.
 type AttributeSetter func(context.Context, middleware.InitializeInput) []attribute.KeyValue
 
+// AttributeBuilder returns an array of KeyValue pairs, it can be used to set custom attributes.
+type AttributeBuilder func(ctx context.Context, in middleware.InitializeInput, out middleware.InitializeOutput) []attribute.KeyValue
+
 type otelMiddlewares struct {
-	tracer          trace.Tracer
-	propagator      propagation.TextMapPropagator
-	attributeSetter []AttributeSetter
+	tracer            trace.Tracer
+	propagator        propagation.TextMapPropagator
+	attributeBuilders []AttributeBuilder
 }
 
 func (m otelMiddlewares) initializeMiddlewareBefore(stack *middleware.Stack) error {
@@ -61,9 +65,6 @@ func (m otelMiddlewares) initializeMiddlewareAfter(stack *middleware.Stack) erro
 			RegionAttr(region),
 			OperationAttr(operation),
 		}
-		for _, setter := range m.attributeSetter {
-			attributes = append(attributes, setter(ctx, in)...)
-		}
 
 		ctx, span := m.tracer.Start(ctx, spanName(serviceID, operation),
 			trace.WithTimestamp(ctx.Value(spanTimestampKey{}).(time.Time)),
@@ -73,6 +74,7 @@ func (m otelMiddlewares) initializeMiddlewareAfter(stack *middleware.Stack) erro
 		defer span.End()
 
 		out, metadata, err = next.HandleInitialize(ctx, in)
+		span.SetAttributes(m.buildAttributes(ctx, in, out)...)
 		if err != nil {
 			span.RecordError(err)
 			span.SetStatus(codes.Error, err.Error())
@@ -125,6 +127,14 @@ func (m otelMiddlewares) deserializeMiddleware(stack *middleware.Stack) error {
 		middleware.Before)
 }
 
+func (m otelMiddlewares) buildAttributes(ctx context.Context, in middleware.InitializeInput, out middleware.InitializeOutput) (attributes []attribute.KeyValue) {
+	for _, builder := range m.attributeBuilders {
+		attributes = append(attributes, builder(ctx, in, out)...)
+	}
+
+	return attributes
+}
+
 func spanName(serviceID, operation string) string {
 	spanName := serviceID
 	if operation != "" {
@@ -145,15 +155,15 @@ func AppendMiddlewares(apiOptions *[]func(*middleware.Stack) error, opts ...Opti
 		opt.apply(&cfg)
 	}
 
-	if cfg.AttributeSetter == nil {
-		cfg.AttributeSetter = []AttributeSetter{DefaultAttributeSetter}
+	if cfg.AttributeBuilders == nil {
+		cfg.AttributeBuilders = []AttributeBuilder{DefaultAttributeBuilder}
 	}
 
 	m := otelMiddlewares{
 		tracer: cfg.TracerProvider.Tracer(ScopeName,
 			trace.WithInstrumentationVersion(Version())),
-		propagator:      cfg.TextMapPropagator,
-		attributeSetter: cfg.AttributeSetter,
+		propagator:        cfg.TextMapPropagator,
+		attributeBuilders: cfg.AttributeBuilders,
 	}
 	*apiOptions = append(*apiOptions, m.initializeMiddlewareBefore, m.initializeMiddlewareAfter, m.finalizeMiddlewareAfter, m.deserializeMiddleware)
 }

instrumentation/github.com/aws/aws-sdk-go-v2/otelaws/config.go

@@ -4,14 +4,17 @@
 package otelaws // import "go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-sdk-go-v2/otelaws"
 
 import (
+	"context"
+	"github.com/aws/smithy-go/middleware"
+	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/propagation"
 	"go.opentelemetry.io/otel/trace"
 )
 
 type config struct {
 	TracerProvider    trace.TracerProvider
 	TextMapPropagator propagation.TextMapPropagator
-	AttributeSetter   []AttributeSetter
+	AttributeBuilders []AttributeBuilder
 }
 
 // Option applies an option value.
@@ -48,9 +51,22 @@ func WithTextMapPropagator(propagator propagation.TextMapPropagator) Option {
 }
 
 // WithAttributeSetter specifies an attribute setter function for setting service specific attributes.
-// If none is specified, the service will be determined by the DefaultAttributeSetter function and the corresponding attributes will be included.
+// If none is specified, the service will be determined by the DefaultAttributeBuilder function and the corresponding attributes will be included.
 func WithAttributeSetter(attributesetters ...AttributeSetter) Option {
+	var attributeBuilders []AttributeBuilder
+	for _, setter := range attributesetters {
+		attributeBuilders = append(attributeBuilders, func(ctx context.Context, in middleware.InitializeInput, out middleware.InitializeOutput) []attribute.KeyValue {
+			return setter(ctx, in)
+		})
+	}
+
+	return WithAttributeBuilder(attributeBuilders...)
+}
+
+// WithAttributeBuilder specifies an attribute setter function for setting service specific attributes.
+// If none is specified, the service will be determined by the DefaultAttributeBuilder function and the corresponding attributes will be included.
+func WithAttributeBuilder(attributeBuilders ...AttributeBuilder) Option {
 	return optionFunc(func(cfg *config) {
-		cfg.AttributeSetter = append(cfg.AttributeSetter, attributesetters...)
+		cfg.AttributeBuilders = append(cfg.AttributeBuilders, attributeBuilders...)
 	})
 }

instrumentation/github.com/aws/aws-sdk-go-v2/otelaws/dynamodbattributes.go

@@ -14,8 +14,8 @@ import (
 	semconv "go.opentelemetry.io/otel/semconv/v1.21.0"
 )
 
-// DynamoDBAttributeSetter sets DynamoDB specific attributes depending on the DynamoDB operation being performed.
-func DynamoDBAttributeSetter(ctx context.Context, in middleware.InitializeInput) []attribute.KeyValue {
+// DynamoDBAttributeBuilder sets DynamoDB specific attributes depending on the DynamoDB operation being performed.
+func DynamoDBAttributeBuilder(ctx context.Context, in middleware.InitializeInput, out middleware.InitializeOutput) []attribute.KeyValue {
 	dynamodbAttributes := []attribute.KeyValue{semconv.DBSystemDynamoDB}
 
 	switch v := in.Parameters.(type) {

instrumentation/github.com/aws/aws-sdk-go-v2/otelaws/dynamodbattributes_test.go

@@ -31,7 +31,7 @@ func TestDynamodbTagsBatchGetItemInput(t *testing.T) {
 		},
 	}
 
-	attributes := DynamoDBAttributeSetter(context.TODO(), input)
+	attributes := DynamoDBAttributeBuilder(context.TODO(), input, middleware.InitializeOutput{})
 
 	assert.Contains(t, attributes, attribute.StringSlice("aws.dynamodb.table_names", []string{"table1"}))
 }
@@ -60,7 +60,7 @@ func TestDynamodbTagsBatchWriteItemInput(t *testing.T) {
 		},
 	}
 
-	attributes := DynamoDBAttributeSetter(context.TODO(), input)
+	attributes := DynamoDBAttributeBuilder(context.TODO(), input, middleware.InitializeOutput{})
 
 	assert.Contains(t, attributes, attribute.StringSlice("aws.dynamodb.table_names", []string{"table1"}))
 }
@@ -114,7 +114,7 @@ func TestDynamodbTagsCreateTableInput(t *testing.T) {
 		},
 	}
 
-	attributes := DynamoDBAttributeSetter(context.TODO(), input)
+	attributes := DynamoDBAttributeBuilder(context.TODO(), input, middleware.InitializeOutput{})
 
 	assert.Contains(t, attributes, attribute.StringSlice(
 		"aws.dynamodb.table_names", []string{"table1"},
@@ -144,7 +144,7 @@ func TestDynamodbTagsDeleteItemInput(t *testing.T) {
 			TableName: aws.String("table1"),
 		},
 	}
-	attributes := DynamoDBAttributeSetter(context.TODO(), input)
+	attributes := DynamoDBAttributeBuilder(context.TODO(), input, middleware.InitializeOutput{})
 
 	assert.Contains(t, attributes, attribute.StringSlice(
 		"aws.dynamodb.table_names", []string{"table1"},
@@ -157,7 +157,7 @@ func TestDynamodbTagsDeleteTableInput(t *testing.T) {
 			TableName: aws.String("table1"),
 		},
 	}
-	attributes := DynamoDBAttributeSetter(context.TODO(), input)
+	attributes := DynamoDBAttributeBuilder(context.TODO(), input, middleware.InitializeOutput{})
 
 	assert.Contains(t, attributes, attribute.StringSlice(
 		"aws.dynamodb.table_names", []string{"table1"},
@@ -170,7 +170,7 @@ func TestDynamodbTagsDescribeTableInput(t *testing.T) {
 			TableName: aws.String("table1"),
 		},
 	}
-	attributes := DynamoDBAttributeSetter(context.TODO(), input)
+	attributes := DynamoDBAttributeBuilder(context.TODO(), input, middleware.InitializeOutput{})
 
 	assert.Contains(t, attributes, attribute.StringSlice(
 		"aws.dynamodb.table_names", []string{"table1"},
@@ -184,7 +184,7 @@ func TestDynamodbTagsListTablesInput(t *testing.T) {
 			Limit:                   aws.Int32(10),
 		},
 	}
-	attributes := DynamoDBAttributeSetter(context.TODO(), input)
+	attributes := DynamoDBAttributeBuilder(context.TODO(), input, middleware.InitializeOutput{})
 
 	assert.Contains(t, attributes, attribute.String("aws.dynamodb.exclusive_start_table", "table1"))
 	assert.Contains(t, attributes, attribute.Int("aws.dynamodb.limit", 10))
@@ -202,7 +202,7 @@ func TestDynamodbTagsPutItemInput(t *testing.T) {
 		},
 	}
 
-	attributes := DynamoDBAttributeSetter(context.TODO(), input)
+	attributes := DynamoDBAttributeBuilder(context.TODO(), input, middleware.InitializeOutput{})
 
 	assert.Contains(t, attributes, attribute.StringSlice(
 		"aws.dynamodb.table_names", []string{"table1"},
@@ -230,7 +230,7 @@ func TestDynamodbTagsQueryInput(t *testing.T) {
 		},
 	}
 
-	attributes := DynamoDBAttributeSetter(context.TODO(), input)
+	attributes := DynamoDBAttributeBuilder(context.TODO(), input, middleware.InitializeOutput{})
 
 	assert.Contains(t, attributes, attribute.StringSlice(
 		"aws.dynamodb.table_names", []string{"table1"},
@@ -257,7 +257,7 @@ func TestDynamodbTagsScanInput(t *testing.T) {
 		},
 	}
 
-	attributes := DynamoDBAttributeSetter(context.TODO(), input)
+	attributes := DynamoDBAttributeBuilder(context.TODO(), input, middleware.InitializeOutput{})
 
 	assert.Contains(t, attributes, attribute.StringSlice(
 		"aws.dynamodb.table_names", []string{"my-table"},
@@ -285,7 +285,7 @@ func TestDynamodbTagsUpdateItemInput(t *testing.T) {
 		},
 	}
 
-	attributes := DynamoDBAttributeSetter(context.TODO(), input)
+	attributes := DynamoDBAttributeBuilder(context.TODO(), input, middleware.InitializeOutput{})
 
 	assert.Contains(t, attributes, attribute.StringSlice(
 		"aws.dynamodb.table_names", []string{"my-table"},
@@ -326,7 +326,7 @@ func TestDynamodbTagsUpdateTableInput(t *testing.T) {
 		},
 	}
 
-	attributes := DynamoDBAttributeSetter(context.TODO(), input)
+	attributes := DynamoDBAttributeBuilder(context.TODO(), input, middleware.InitializeOutput{})
 
 	assert.Contains(t, attributes, attribute.StringSlice(
 		"aws.dynamodb.table_names", []string{"my-table"},

instrumentation/github.com/aws/aws-sdk-go-v2/otelaws/snsattributes.go

@@ -14,8 +14,8 @@ import (
 	semconv "go.opentelemetry.io/otel/semconv/v1.27.0"
 )
 
-// SNSAttributeSetter sets SNS specific attributes depending on the SNS operation is being performed.
-func SNSAttributeSetter(ctx context.Context, in middleware.InitializeInput) []attribute.KeyValue {
+// SNSAttributeBuilder sets SNS specific attributes depending on the SNS operation is being performed.
+func SNSAttributeBuilder(ctx context.Context, in middleware.InitializeInput, out middleware.InitializeOutput) []attribute.KeyValue {
 	snsAttributes := []attribute.KeyValue{semconv.MessagingSystemKey.String("aws_sns")}
 
 	switch v := in.Parameters.(type) {

instrumentation/github.com/aws/aws-sdk-go-v2/otelaws/snsattributes_test.go

@@ -23,7 +23,7 @@ func TestPublishInput(t *testing.T) {
 		},
 	}
 
-	attributes := SNSAttributeSetter(context.Background(), input)
+	attributes := SNSAttributeBuilder(context.Background(), input, middleware.InitializeOutput{})
 
 	assert.Contains(t, attributes, semconv.MessagingSystemKey.String("aws_sns"))
 	assert.Contains(t, attributes, semconv.MessagingDestinationName("my-topic"))
@@ -36,7 +36,7 @@ func TestPublishInputWithNoDestination(t *testing.T) {
 		Parameters: &sns.PublishInput{},
 	}
 
-	attributes := SNSAttributeSetter(context.Background(), input)
+	attributes := SNSAttributeBuilder(context.Background(), input, middleware.InitializeOutput{})
 
 	assert.Contains(t, attributes, semconv.MessagingSystemKey.String("aws_sns"))
 	assert.Contains(t, attributes, semconv.MessagingDestinationName(""))
@@ -52,7 +52,7 @@ func TestPublishBatchInput(t *testing.T) {
 		},
 	}
 
-	attributes := SNSAttributeSetter(context.Background(), input)
+	attributes := SNSAttributeBuilder(context.Background(), input, middleware.InitializeOutput{})
 
 	assert.Contains(t, attributes, semconv.MessagingSystemKey.String("aws_sns"))
 	assert.Contains(t, attributes, semconv.MessagingDestinationName("my-topic-batch"))

instrumentation/github.com/aws/aws-sdk-go-v2/otelaws/sqsattributes.go

@@ -13,8 +13,8 @@ import (
 	semconv "go.opentelemetry.io/otel/semconv/v1.21.0"
 )
 
-// SQSAttributeSetter sets SQS specific attributes depending on the SQS operation being performed.
-func SQSAttributeSetter(ctx context.Context, in middleware.InitializeInput) []attribute.KeyValue {
+// SQSAttributeBuilder sets SQS specific attributes depending on the SQS operation being performed.
+func SQSAttributeBuilder(ctx context.Context, in middleware.InitializeInput, out middleware.InitializeOutput) []attribute.KeyValue {
 	sqsAttributes := []attribute.KeyValue{semconv.MessagingSystem("AmazonSQS")}
 
 	key := semconv.NetPeerNameKey