diff --git a/go.mod b/go.mod index 54cf7b2a..a851827e 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 99c68dfc..b3faff08 100644 --- a/go.sum +++ b/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= diff --git a/src/args/argument_list.go b/src/args/argument_list.go index f41e242a..80d4825b 100644 --- a/src/args/argument_list.go +++ b/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 diff --git a/src/connection/pgsql_connection.go b/src/connection/pgsql_connection.go index fd762f2a..b983ba26 100644 --- a/src/connection/pgsql_connection.go +++ b/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,10 +185,10 @@ 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, } @@ -162,8 +196,8 @@ func createConnectionURL(ci *connectionInfo, database string) string { 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") diff --git a/src/connection/token_cache.go b/src/connection/token_cache.go new file mode 100644 index 00000000..b105f09f --- /dev/null +++ b/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 +}