[processor/awsentity] Add configurable entity transform option (#1682)

Author: zhihonl

internal/entity/entity.go

@@ -0,0 +1,16 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT
+
+package entity
+
+// KeyPair represents a key-value pair for entity attributes
+type KeyPair struct {
+	Key   string `mapstructure:"key"`
+	Value string `mapstructure:"value"`
+}
+
+// Transform contains configuration for overriding entity attributes
+type Transform struct {
+	KeyAttributes []KeyPair `mapstructure:"key_attributes"`
+	Attributes    []KeyPair `mapstructure:"attributes"`
+}

plugins/processors/awsentity/config.go

@@ -4,7 +4,12 @@
 package awsentity
 
 import (
+	"errors"
+
 	"go.opentelemetry.io/collector/component"
+
+	"github.com/aws/amazon-cloudwatch-agent/internal/entity"
+	"github.com/aws/amazon-cloudwatch-agent/plugins/processors/awsentity/entityattributes"
 )
 
 type Config struct {
@@ -23,11 +28,34 @@ type Config struct {
 	// EntityType determines the type of entity processing done for
 	// telemetry. Possible values are Service and Resource
 	EntityType string `mapstructure:"entity_type,omitempty"`
+	// TransformEntity contains configuration for overriding entity attributes
+	TransformEntity *entity.Transform `mapstructure:"transform_entity,omitempty"`
 }
 
 // Verify Config implements Processor interface.
 var _ component.Config = (*Config)(nil)
 
 func (cfg *Config) Validate() error {
+	if cfg.TransformEntity != nil {
+		// Validate key attributes
+		for _, keyAttr := range cfg.TransformEntity.KeyAttributes {
+			if !entityattributes.IsAllowedKeyAttribute(keyAttr.Key) {
+				return errors.New("Invalid key attribute name for entity: " + keyAttr.Key)
+			}
+			if keyAttr.Value == "" {
+				return errors.New("empty value for entity key attribute")
+			}
+		}
+
+		// Validate regular attributes
+		for _, attr := range cfg.TransformEntity.Attributes {
+			if !entityattributes.IsAllowedAttribute(attr.Key) {
+				return errors.New("Invalid attribute name for entity: " + attr.Key)
+			}
+			if attr.Value == "" {
+				return errors.New("empty value for entity attribute")
+			}
+		}
+	}
 	return nil
 }

plugins/processors/awsentity/config_test.go

@@ -8,6 +8,9 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"go.opentelemetry.io/collector/confmap"
+
+	"github.com/aws/amazon-cloudwatch-agent/internal/entity"
+	"github.com/aws/amazon-cloudwatch-agent/plugins/processors/awsentity/entityattributes"
 )
 
 func TestUnmarshalDefaultConfig(t *testing.T) {
@@ -16,3 +19,126 @@ func TestUnmarshalDefaultConfig(t *testing.T) {
 	assert.NoError(t, confmap.New().Unmarshal(cfg))
 	assert.Equal(t, factory.CreateDefaultConfig(), cfg)
 }
+
+func TestUnmarshalConfig(t *testing.T) {
+	tests := []struct {
+		name        string
+		conf        *confmap.Conf
+		expected    *Config
+		expectError bool
+	}{
+		{
+			name: "TestValidEntityTransform",
+			conf: confmap.NewFromStringMap(map[string]interface{}{
+				"entity_type": entityattributes.Service,
+				"platform":    "ec2",
+				"transform_entity": map[string]interface{}{
+					"key_attributes": []interface{}{
+						map[string]interface{}{
+							"key":   entityattributes.ServiceName,
+							"value": "config-service-name",
+						},
+						map[string]interface{}{
+							"key":   entityattributes.DeploymentEnvironment,
+							"value": "config-environment-name",
+						},
+					},
+					"attributes": []interface{}{
+						map[string]interface{}{
+							"key":   entityattributes.ServiceNameSource,
+							"value": "UserConfiguration",
+						},
+					},
+				},
+			}),
+			expected: &Config{
+				EntityType: entityattributes.Service,
+				Platform:   "ec2",
+				TransformEntity: &entity.Transform{
+					KeyAttributes: []entity.KeyPair{
+						{
+							Key:   entityattributes.ServiceName,
+							Value: "config-service-name",
+						},
+						{
+							Key:   entityattributes.DeploymentEnvironment,
+							Value: "config-environment-name",
+						},
+					},
+					Attributes: []entity.KeyPair{
+						{
+							Key:   entityattributes.ServiceNameSource,
+							Value: "UserConfiguration",
+						},
+					},
+				},
+			},
+			expectError: false,
+		},
+		{
+			name: "TestInvalidEntityTransform",
+			conf: confmap.NewFromStringMap(map[string]interface{}{
+				"entity_type": entityattributes.Service,
+				"platform":    "ec2",
+				"transform_entity": map[string]interface{}{
+					"key_attributes": []interface{}{
+						map[string]interface{}{
+							"key":   "InvalidKey",
+							"value": "some-value",
+						},
+					},
+				},
+			}),
+			expectError: true,
+		},
+		{
+			name: "TestEmptyEntityTransform",
+			conf: confmap.NewFromStringMap(map[string]interface{}{
+				"entity_type": entityattributes.Service,
+				"platform":    "ec2",
+			}),
+			expected: &Config{
+				EntityType: entityattributes.Service,
+				Platform:   "ec2",
+			},
+			expectError: false,
+		},
+		{
+			name: "TestMissingRequiredFieldEntityTransform",
+			conf: confmap.NewFromStringMap(map[string]interface{}{
+				"transform_entity": map[string]interface{}{
+					"key_attributes": []interface{}{
+						map[string]interface{}{
+							"key":   entityattributes.ServiceName,
+							"value": "",
+						},
+					},
+				},
+			}),
+			expectError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			factory := NewFactory()
+			cfg := factory.CreateDefaultConfig()
+
+			err := tt.conf.Unmarshal(cfg)
+
+			assert.NoError(t, err)
+
+			// Validate the configuration
+			err = cfg.(*Config).Validate()
+			if tt.expectError {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+			}
+
+			if err == nil {
+				assert.Equal(t, tt.expected, cfg)
+			}
+		})
+	}
+}

plugins/processors/awsentity/entityattributes/entityattributes.go

@@ -89,6 +89,52 @@ var attributeEntityToShortNameMap = map[string]string{
 	AttributeEntityServiceNameSource: ServiceNameSource,
 }
 
+// shortNameToEntityMap is the reverse mapping of keyAttributeEntityToShortNameMap
+var keyAttributeEntityToLongNameMap = map[string]string{
+	EntityType:            AttributeEntityType,
+	ResourceType:          AttributeEntityResourceType,
+	Identifier:            AttributeEntityIdentifier,
+	AwsAccountId:          AttributeEntityAwsAccountId,
+	ServiceName:           AttributeEntityServiceName,
+	DeploymentEnvironment: AttributeEntityDeploymentEnvironment,
+}
+
+// shortNameToAttributeMap is the reverse mapping of attributeEntityToShortNameMap
+var attributeEntityToLongNameMap = map[string]string{
+	NamespaceField:    AttributeEntityNamespace,
+	Workload:          AttributeEntityWorkload,
+	Node:              AttributeEntityNode,
+	Platform:          AttributeEntityPlatformType,
+	InstanceID:        AttributeEntityInstanceID,
+	AutoscalingGroup:  AttributeEntityAutoScalingGroup,
+	ServiceNameSource: AttributeEntityServiceNameSource,
+}
+
+// GetFullAttributeName returns the full attribute name for a given short name
+func GetFullAttributeName(shortName string) (string, bool) {
+	// First check key attributes
+	if fullName, ok := keyAttributeEntityToLongNameMap[shortName]; ok {
+		return fullName, true
+	}
+	// Then check regular attributes
+	if fullName, ok := attributeEntityToLongNameMap[shortName]; ok {
+		return fullName, true
+	}
+	return "", false
+}
+
+// IsAllowedKeyAttribute checks if the given key is an allowed entity key attribute name
+func IsAllowedKeyAttribute(key string) bool {
+	_, exists := keyAttributeEntityToLongNameMap[key]
+	return exists
+}
+
+// IsAllowedAttribute checks if the given key is an allowed attribute name
+func IsAllowedAttribute(key string) bool {
+	_, exists := attributeEntityToLongNameMap[key]
+	return exists
+}
+
 func CreateCloudWatchEntityFromAttributes(resourceAttributes pcommon.Map) cloudwatch.Entity {
 	keyAttributesMap := map[string]*string{}
 	attributeMap := map[string]*string{}

plugins/processors/awsentity/internal/entitytransformer/entitytransformer.go

@@ -0,0 +1,48 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: MIT
+
+package entitytransformer
+
+import (
+	"go.opentelemetry.io/collector/pdata/pcommon"
+	"go.uber.org/zap"
+
+	"github.com/aws/amazon-cloudwatch-agent/internal/entity"
+	"github.com/aws/amazon-cloudwatch-agent/plugins/processors/awsentity/entityattributes"
+)
+
+type EntityTransformer struct {
+	transform *entity.Transform
+	logger    *zap.Logger
+}
+
+func NewEntityTransformer(transform *entity.Transform, logger *zap.Logger) *EntityTransformer {
+	return &EntityTransformer{
+		transform: transform,
+		logger:    logger,
+	}
+}
+
+func (p *EntityTransformer) ApplyTransforms(resourceAttrs pcommon.Map) {
+	if p.transform == nil {
+		return
+	}
+
+	// Apply key attributes
+	for _, keyAttr := range p.transform.KeyAttributes {
+		if fullName, ok := entityattributes.GetFullAttributeName(keyAttr.Key); ok {
+			resourceAttrs.PutStr(fullName, keyAttr.Value)
+		} else {
+			p.logger.Debug("Unrecognized key attribute", zap.String("key", keyAttr.Key))
+		}
+	}
+
+	// Apply additional attributes
+	for _, attr := range p.transform.Attributes {
+		if fullName, ok := entityattributes.GetFullAttributeName(attr.Key); ok {
+			resourceAttrs.PutStr(fullName, attr.Value)
+		} else {
+			p.logger.Debug("Unrecognized attribute", zap.String("key", attr.Key))
+		}
+	}
+}