Added AWS IAM authentication

go.mod

@@ -3,6 +3,8 @@ module github.com/newrelic/nri-postgresql
 go 1.24.6
 
 require (
+	github.com/aws/aws-sdk-go-v2/config v1.27.18
+	github.com/aws/aws-sdk-go-v2/feature/rds/auth v1.6.2
 	github.com/blang/semver/v4 v4.0.0
 	github.com/go-viper/mapstructure/v2 v2.4.0
 	github.com/jmoiron/sqlx v1.4.0
@@ -15,6 +17,18 @@ require (
 )
 
 require (
+	github.com/aws/aws-sdk-go-v2 v1.37.2 // indirect
+	github.com/aws/aws-sdk-go-v2/credentials v1.17.18 // indirect
+	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.5 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.9 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.9 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.2 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.11 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sso v1.20.11 // indirect
+	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.24.5 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sts v1.28.12 // indirect
+	github.com/aws/smithy-go v1.22.5 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/kr/text v0.2.0 // indirect
 	github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e // indirect

go.sum

@@ -1,5 +1,33 @@
 filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA=
 filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4=
+github.com/aws/aws-sdk-go-v2 v1.37.2 h1:xkW1iMYawzcmYFYEV0UCMxc8gSsjCGEhBXQkdQywVbo=
+github.com/aws/aws-sdk-go-v2 v1.37.2/go.mod h1:9Q0OoGQoboYIAJyslFyF1f5K1Ryddop8gqMhWx/n4Wg=
+github.com/aws/aws-sdk-go-v2/config v1.27.18 h1:wFvAnwOKKe7QAyIxziwSKjmer9JBMH1vzIL6W+fYuKk=
+github.com/aws/aws-sdk-go-v2/config v1.27.18/go.mod h1:0xz6cgdX55+kmppvPm2IaKzIXOheGJhAufacPJaXZ7c=
+github.com/aws/aws-sdk-go-v2/credentials v1.17.18 h1:D/ALDWqK4JdY3OFgA2thcPO1c9aYTT5STS/CvnkqY1c=
+github.com/aws/aws-sdk-go-v2/credentials v1.17.18/go.mod h1:JuitCWq+F5QGUrmMPsk945rop6bB57jdscu+Glozdnc=
+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.5 h1:dDgptDO9dxeFkXy+tEgVkzSClHZje/6JkPW5aZyEvrQ=
+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.5/go.mod h1:gjvE2KBUgUQhcv89jqxrIxH9GaKs1JbZzWejj/DaHGA=
+github.com/aws/aws-sdk-go-v2/feature/rds/auth v1.6.2 h1:QbFjOdplTkOgviHNKyTW/TZpvIYhD6lqEc3tkIvqMoQ=
+github.com/aws/aws-sdk-go-v2/feature/rds/auth v1.6.2/go.mod h1:d0pTYUeTv5/tPSlbPZZQSqssM158jZBs02jx2LDslM8=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.9 h1:cy8ahBJuhtM8GTTSyOkfy6WVPV1IE+SS5/wfXUYuulw=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.9/go.mod h1:CZBXGLaJnEZI6EVNcPd7a6B5IC5cA/GkRWtu9fp3S6Y=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.9 h1:A4SYk07ef04+vxZToz9LWvAXl9LW0NClpPpMsi31cz0=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.9/go.mod h1:5jJcHuwDagxN+ErjQ3PU3ocf6Ylc/p9x+BLO/+X4iXw=
+github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 h1:hT8rVHwugYE2lEfdFE0QWVo81lF7jMrYJVDWI+f+VxU=
+github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0/go.mod h1:8tu/lYfQfFe6IGnaOdrpVgEL2IrrDOf6/m9RQum4NkY=
+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.2 h1:Ji0DY1xUsUr3I8cHps0G+XM3WWU16lP6yG8qu1GAZAs=
+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.2/go.mod h1:5CsjAbs3NlGQyZNFACh+zztPDI7fU6eW9QsxjfnuBKg=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.11 h1:o4T+fKxA3gTMcluBNZZXE9DNaMkJuUL1O3mffCUjoJo=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.11/go.mod h1:84oZdJ+VjuJKs9v1UTC9NaodRZRseOXCTgku+vQJWR8=
+github.com/aws/aws-sdk-go-v2/service/sso v1.20.11 h1:gEYM2GSpr4YNWc6hCd5nod4+d4kd9vWIAWrmGuLdlMw=
+github.com/aws/aws-sdk-go-v2/service/sso v1.20.11/go.mod h1:gVvwPdPNYehHSP9Rs7q27U1EU+3Or2ZpXvzAYJNh63w=
+github.com/aws/aws-sdk-go-v2/service/ssooidc v1.24.5 h1:iXjh3uaH3vsVcnyZX7MqCoCfcyxIrVE9iOQruRaWPrQ=
+github.com/aws/aws-sdk-go-v2/service/ssooidc v1.24.5/go.mod h1:5ZXesEuy/QcO0WUnt+4sDkxhdXRHTu2yG0uCSH8B6os=
+github.com/aws/aws-sdk-go-v2/service/sts v1.28.12 h1:M/1u4HBpwLuMtjlxuI2y6HoVLzF5e2mfxHCg7ZVMYmk=
+github.com/aws/aws-sdk-go-v2/service/sts v1.28.12/go.mod h1:kcfd+eTdEi/40FIbLq4Hif3XMXnl5b/+t/KTfLt9xIk=
+github.com/aws/smithy-go v1.22.5 h1:P9ATCXPMb2mPjYBgueqJNCA5S9UfktsW0tTxi+a7eqw=
+github.com/aws/smithy-go v1.22.5/go.mod h1:t1ufH5HMublsJYulve2RKmHDC15xu1f26kHCp/HgceI=
 github.com/blang/semver/v4 v4.0.0 h1:1PFHFE6yCCTv8C1TeyNNarDzntLi7wMI5i/pzqYIsAM=
 github.com/blang/semver/v4 v4.0.0/go.mod h1:IbckMUScFkM3pff0VJDNKRiT6TG/YpiHIM2yvyW5YoQ=
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
@@ -8,8 +36,6 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/go-sql-driver/mysql v1.8.1 h1:LedoTUt/eveggdHS9qUFC1EFSa8bU2+1pZjSRpvNJ1Y=
 github.com/go-sql-driver/mysql v1.8.1/go.mod h1:wEBSXgmK//2ZFJyE+qWnIsVGmvmEKlqwuVSjsCm7DZg=
-github.com/go-viper/mapstructure/v2 v2.3.0 h1:27XbWsHIqhbdR5TIC911OfYvgSaW93HM+dX7970Q7jk=
-github.com/go-viper/mapstructure/v2 v2.3.0/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM=
 github.com/go-viper/mapstructure/v2 v2.4.0 h1:EBsztssimR/CONLSZZ04E8qAkxNYq4Qp9LvH92wZUgs=
 github.com/go-viper/mapstructure/v2 v2.4.0/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM=
 github.com/jmoiron/sqlx v1.4.0 h1:1PLqN7S1UYp5t4SrVVnt4nUVNemrDAtxlulVe+Qgm3o=
@@ -32,10 +58,6 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+
 github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
 github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
-github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
-github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
-github.com/stretchr/testify v1.11.0 h1:ib4sjIrwZKxE5u/Japgo/7SJV3PvgjGiRNAvTVGqQl8=
-github.com/stretchr/testify v1.11.0/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
 github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
 github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
 github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f h1:J9EGpcZtP0E/raorCMxlFGSTBrsSlaDGf3jU/qvAE2c=

src/args/argument_list.go

@@ -3,6 +3,7 @@ package args
 
 import (
 	"errors"
+	"os"
 	sdkArgs "github.com/newrelic/infra-integrations-sdk/v3/args"
 	"github.com/newrelic/infra-integrations-sdk/v3/log"
 )
@@ -34,12 +35,40 @@ type ArgumentList struct {
 	QueryMonitoringResponseTimeThreshold int    `default:"500" help:"Threshold in milliseconds for query response time. If response time for the individual query exceeds this threshold, the individual query is reported in metrics"`
 	QueryMonitoringCountThreshold        int    `default:"20" help:"The number of records for each query performance metrics"`
 	IsRds                                bool   `default:"false" help:"If true, the integration will support on AWS RDS. This will enable RDS-specific metrics and configurations."`
+	AWSIamAuth                           bool   `default:"false" help:"If true, use AWS IAM authentication instead of username/password. Requires AWS credentials and RDS IAM database authentication enabled."`
+	AwsRegion                            string `default:"" help:"AWS region for IAM authentication. Overrides AWS_REGION environment variable and AWS config file."`
 }
 
-// Validate validates PostgreSQl arguments
+// GetEffectiveAwsRegion returns the AWS region to use, prioritizing ArgumentList input over environment/config
+func (al ArgumentList) GetEffectiveAwsRegion() string {
+	if al.AwsRegion != "" {
+		return al.AwsRegion
+	}
+
+	if envRegion := os.Getenv("AWS_REGION"); envRegion != "" {
+		return envRegion
+	}
+
+	if defaultRegion := os.Getenv("AWS_DEFAULT_REGION"); defaultRegion != "" {
+		return defaultRegion
+	}
+
+	return ""
+}
+
+// Validate validates PostgreSQL arguments
 func (al ArgumentList) Validate() error {
-	if al.Username == "" || al.Password == "" {
-		return errors.New("invalid configuration: must specify a username and password")
+	if al.AWSIamAuth {
+		if al.Username == "" {
+			return errors.New("AWS IAM authentication requires a username. Please provide --username argument")
+		}
+		if al.GetEffectiveAwsRegion() == "" {
+			return errors.New("AWS IAM authentication requires a region. Please set --aws-region argument or AWS_REGION environment variable")
+		}
+	} else {
+		if al.Username == "" || al.Password == "" {
+			return errors.New("invalid configuration: must specify a username and password")
+		}
 	}
 	if err := al.validateSSL(); err != nil {
 		return err

src/connection/pgsql_connection.go

@@ -2,8 +2,10 @@
 package connection
 
 import (
+	"context"
 	"fmt"
 	"net/url"
+	"strconv"
 
 	"github.com/jmoiron/sqlx"
 	// pq is required for postgreSQL driver but isn't used in code
@@ -45,6 +47,8 @@ type connectionInfo struct {
 	SSLRootCertLocation    string
 	SSLKeyLocation         string
 	TrustServerCertificate bool
+	AWSIamAuth             bool
+	AwsRegion              string
 }
 
 // DefaultConnectionInfo takes an argument list and constructs a default connection out of it
@@ -61,12 +65,20 @@ func DefaultConnectionInfo(al *args.ArgumentList) Info {
 		SSLRootCertLocation:    al.SSLRootCertLocation,
 		SSLKeyLocation:         al.SSLKeyLocation,
 		TrustServerCertificate: al.TrustServerCertificate,
+		AWSIamAuth:             al.AWSIamAuth,
+		AwsRegion:              al.GetEffectiveAwsRegion(),
 	}
 }
 
 // NewConnection creates a new PGSQLConnection from args
 func (ci *connectionInfo) NewConnection(database string) (*PGSQLConnection, error) {
-	db, err := sqlx.Open("postgres", createConnectionURL(ci, database))
+	log.Debug("Creating connection to database: %s", database)
+	password, err := ci.resolvePassword()
+	if err != nil {
+		return nil, fmt.Errorf("failed to resolve password: %w", err)
+	}
+
+	db, err := sqlx.Open("postgres", createConnectionURL(ci, database, password))
 	if err != nil {
 		return nil, err
 	}
@@ -76,6 +88,28 @@ func (ci *connectionInfo) NewConnection(database string) (*PGSQLConnection, erro
 	}, nil
 }
 
+// resolvePassword returns either an AWS IAM auth token (cached) or the static password
+func (ci *connectionInfo) resolvePassword() (string, error) {
+	if !ci.AWSIamAuth {
+		return ci.Password, nil
+	}
+
+	port, err := strconv.Atoi(ci.Port)
+	if err != nil {
+		return "", fmt.Errorf("invalid port: %w", err)
+	}
+	endpoint := fmt.Sprintf("%s:%d", ci.Host, port)
+
+	// Use cached token manager
+	cache := getTokenCache()
+	token, err := cache.getCachedToken(context.TODO(), endpoint, ci.AwsRegion, ci.Username)
+	if err != nil {
+		return "", fmt.Errorf("failed to get IAM auth token: %w", err)
+	}
+
+	return token, nil
+}
+
 func (ci *connectionInfo) HostPort() (string, string) {
 	return ci.Host, ci.Port
 }
@@ -151,19 +185,19 @@ func (p PGSQLConnection) HaveExtensionInSchema(extensionName, schemaName string)
 
 // createConnectionURL creates the connection string. A list of parameters
 // can be found here https://godoc.org/github.com/lib/pq#hdr-Connection_String_Parameters
-func createConnectionURL(ci *connectionInfo, database string) string {
+func createConnectionURL(ci *connectionInfo, database string, password string) string {
 	connectionURL := &url.URL{
 		Scheme: "postgres",
-		User:   url.UserPassword(ci.Username, ci.Password),
+		User:   url.UserPassword(ci.Username, password),
 		Host:   fmt.Sprintf("%s:%s", ci.Host, ci.Port),
 		Path:   database,
 	}
 
 	query := url.Values{}
 	query.Add("connect_timeout", ci.Timeout)
 
-	// SSL settings
-	if ci.EnableSSL {
+	// SSL settings - Force SSL for IAM auth
+	if ci.EnableSSL || ci.AWSIamAuth {
 		addSSLQueries(query, ci)
 	} else {
 		query.Add("sslmode", "disable")

src/connection/token_cache.go

@@ -0,0 +1,79 @@
+package connection
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/feature/rds/auth"
+	"github.com/newrelic/infra-integrations-sdk/v3/log"
+)
+
+type tokenCacheEntry struct {
+	token     string
+	createdAt time.Time
+}
+
+type tokenCache struct {
+	cache map[string]*tokenCacheEntry
+	mutex sync.RWMutex
+}
+
+var (
+	globalTokenCache *tokenCache
+	cacheOnce        sync.Once
+)
+
+func getTokenCache() *tokenCache {
+	cacheOnce.Do(func() {
+		globalTokenCache = &tokenCache{
+			cache: make(map[string]*tokenCacheEntry),
+		}
+	})
+	return globalTokenCache
+}
+
+func (tc *tokenCache) getCachedToken(ctx context.Context, endpoint, region, username string) (string, error) {
+	cacheKey := fmt.Sprintf("%s:%s:%s", endpoint, region, username)
+
+	tc.mutex.RLock()
+	entry, exists := tc.cache[cacheKey]
+	if exists {
+		tc.mutex.RUnlock()
+		log.Debug("CACHE HIT: Reusing cached AWS IAM token for %s, created: %v", cacheKey, entry.createdAt.Format("15:04:05"))
+		return entry.token, nil
+	}
+	tc.mutex.RUnlock()
+
+	tc.mutex.Lock()
+	defer tc.mutex.Unlock()
+
+	if entry, exists := tc.cache[cacheKey]; exists {
+		log.Debug("CACHE HIT (double-check): Reusing cached AWS IAM token for %s, created: %v", cacheKey, entry.createdAt.Format("15:04:05"))
+		return entry.token, nil
+	}
+
+	log.Info("CACHE MISS: Generating new AWS IAM token for %s", cacheKey)
+
+	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
+	if err != nil {
+		return "", fmt.Errorf("failed to load AWS config: %w", err)
+	}
+
+	token, err := auth.BuildAuthToken(ctx, endpoint, region, username, cfg.Credentials)
+	if err != nil {
+		return "", fmt.Errorf("failed to generate IAM auth token: %w", err)
+	}
+
+	createdAt := time.Now()
+	tc.cache[cacheKey] = &tokenCacheEntry{
+		token:     token,
+		createdAt: createdAt,
+	}
+
+	log.Info("TOKEN CACHED: New AWS IAM token cached for %s, created at %v", cacheKey, createdAt.Format("15:04:05"))
+
+	return token, nil
+}