Feat performance monitoring

Makefile

@@ -32,9 +32,19 @@ test:
 
 integration-test:
 	@echo "=== $(INTEGRATION) === [ test ]: running integration tests..."
-	@docker compose -f tests/docker-compose.yml pull
-	@go test -v -tags=integration -count 1 ./tests/. || (ret=$$?; docker compose -f tests/docker-compose.yml down && exit $$ret)
-	@docker compose -f tests/docker-compose.yml down
+	@docker compose -f tests/docker-compose.yml up -d
+	# Sleep added to allow postgres with test data and extensions to start up
+	@sleep 10
+	@go test -v -tags=integration -count 1 ./tests/postgresql_test.go -timeout 300s || (ret=$$?; docker compose -f tests/docker-compose.yml down -v && exit $$ret)
+	@docker compose -f tests/docker-compose.yml down -v
+	@echo "=== $(INTEGRATION) === [ test ]: running integration tests for query performance monitoring..."
+	@echo "Starting containers for performance tests..."
+	@docker compose -f tests/docker-compose-performance.yml up -d
+	# Sleep added to allow postgres with test data and extensions to start up
+	@sleep 30
+	@go test -v -tags=query_performance ./tests/postgresqlperf_test.go -timeout 600s || (ret=$$?; docker compose -f tests/docker-compose-performance.yml down -v && exit $$ret)
+	@echo "Stopping performance test containers..."
+	@docker compose -f tests/docker-compose-performance.yml down -v
 
 install: compile
 	@echo "=== $(INTEGRATION) === [ install ]: installing bin/$(BINARY_NAME)..."

go.mod

@@ -11,6 +11,7 @@ require (
 	github.com/xeipuuv/gojsonschema v1.2.0
 	gopkg.in/DATA-DOG/go-sqlmock.v1 v1.3.0
 	gopkg.in/yaml.v3 v3.0.1
+	github.com/go-viper/mapstructure/v2 v2.2.1
 )
 
 require (

go.sum

@@ -8,6 +8,8 @@ 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.2.1 h1:ZAaOCxANMuZx5RCeg0mBdEZk7DZasvvZIxtHqx8aGss=
+github.com/go-viper/mapstructure/v2 v2.2.1/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM=
 github.com/jmoiron/sqlx v1.4.0 h1:1PLqN7S1UYp5t4SrVVnt4nUVNemrDAtxlulVe+Qgm3o=
 github.com/jmoiron/sqlx v1.4.0/go.mod h1:ZrZ7UsYB/weZdl2Bxg6jCRO9c3YHl8r3ahlKmRT4JLY=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=

postgresql-config.yml.sample

@@ -50,6 +50,15 @@ integrations:
     # True if SSL is to be used. Defaults to false.
     ENABLE_SSL: "false"
 
+    # Enable query performance monitoring - Defaults to false
+    # ENABLE_QUERY_MONITORING : "false"
+
+    # Threshold in milliseconds for query response time to fetch individual query performance metrics - Defaults to 500
+    # QUERY_MONITORING_RESPONSE_TIME_THRESHOLD : "500"
+
+    # The number of records for each query performance metrics - Defaults to 20
+    # QUERY_MONITORING_COUNT_THRESHOLD : "20"
+
     # True if the SSL certificate should be trusted without validating.
     # Setting this to true may open up the monitoring service to MITM attacks.
     # Defaults to false.

src/args/argument_list.go

@@ -3,46 +3,46 @@ package args
 
 import (
 	"errors"
-
 	sdkArgs "github.com/newrelic/infra-integrations-sdk/v3/args"
 	"github.com/newrelic/infra-integrations-sdk/v3/log"
 )
 
 // ArgumentList struct that holds all PostgreSQL arguments
 type ArgumentList struct {
 	sdkArgs.DefaultArgumentList
-	Username                     string `default:"" help:"The username for the PostgreSQL database"`
-	Password                     string `default:"" help:"The password for the specified username"`
-	Hostname                     string `default:"localhost" help:"The PostgreSQL hostname to connect to"`
-	Database                     string `default:"postgres" help:"The PostgreSQL database name to connect to"`
-	Port                         string `default:"5432" help:"The port to connect to the PostgreSQL database"`
-	CollectionList               string `default:"{}" help:"A JSON object which defines the databases, schemas, tables, and indexes to collect. Can also be a JSON array that list databases to be collected. Can also be the string literal 'ALL' to collect everything. Collects nothing by default."`
-	CollectionIgnoreDatabaseList string `default:"[]" help:"A JSON array that list databases that will be excluded from collection. Nothing is excluded by default."`
-	CollectionIgnoreTableList    string `default:"[]" help:"A JSON array that list tables that will be excluded from collection. Nothing is excluded by default."`
-	SSLRootCertLocation          string `default:"" help:"Absolute path to PEM encoded root certificate file"`
-	SSLCertLocation              string `default:"" help:"Absolute path to PEM encoded client cert file"`
-	SSLKeyLocation               string `default:"" help:"Absolute path to PEM encoded client key file"`
-	Timeout                      string `default:"10" help:"Maximum wait for connection, in seconds. Set 0 for no timeout"`
-	CustomMetricsQuery           string `default:"" help:"A SQL query to collect custom metrics. Must have the columns metric_name, metric_type, and metric_value. Additional columns are added as attributes"`
-	CustomMetricsConfig          string `default:"" help:"YAML configuration with one or more custom SQL queries to collect"`
-	EnableSSL                    bool   `default:"false" help:"If true will use SSL encryption, false will not use encryption"`
-	TrustServerCertificate       bool   `default:"false" help:"If true server certificate is not verified for SSL. If false certificate will be verified against supplied certificate"`
-	Pgbouncer                    bool   `default:"false" help:"Collects metrics from PgBouncer instance. Assumes connection is through PgBouncer."`
-	CollectDbLockMetrics         bool   `default:"false" help:"If true, enables collection of lock metrics for the specified database. (Note: requires that the 'tablefunc' extension is installed)"` //nolint: stylecheck
-	CollectBloatMetrics          bool   `default:"true" help:"Enable collecting bloat metrics which can be performance intensive"`
-	ShowVersion                  bool   `default:"false" help:"Print build information and exit"`
+	Username                             string `default:"" help:"The username for the PostgreSQL database"`
+	Password                             string `default:"" help:"The password for the specified username"`
+	Hostname                             string `default:"localhost" help:"The PostgreSQL hostname to connect to"`
+	Database                             string `default:"postgres" help:"The PostgreSQL database name to connect to"`
+	Port                                 string `default:"5432" help:"The port to connect to the PostgreSQL database"`
+	CollectionList                       string `default:"{}" help:"A JSON object which defines the databases, schemas, tables, and indexes to collect. Can also be a JSON array that list databases to be collected. Can also be the string literal 'ALL' to collect everything. Collects nothing by default."`
+	CollectionIgnoreDatabaseList         string `default:"[]" help:"A JSON array that list databases that will be excluded from collection. Nothing is excluded by default."`
+	CollectionIgnoreTableList            string `default:"[]" help:"A JSON array that list tables that will be excluded from collection. Nothing is excluded by default."`
+	SSLRootCertLocation                  string `default:"" help:"Absolute path to PEM encoded root certificate file"`
+	SSLCertLocation                      string `default:"" help:"Absolute path to PEM encoded client cert file"`
+	SSLKeyLocation                       string `default:"" help:"Absolute path to PEM encoded client key file"`
+	Timeout                              string `default:"10" help:"Maximum wait for connection, in seconds. Set 0 for no timeout"`
+	CustomMetricsQuery                   string `default:"" help:"A SQL query to collect custom metrics. Must have the columns metric_name, metric_type, and metric_value. Additional columns are added as attributes"`
+	CustomMetricsConfig                  string `default:"" help:"YAML configuration with one or more custom SQL queries to collect"`
+	EnableSSL                            bool   `default:"false" help:"If true will use SSL encryption, false will not use encryption"`
+	TrustServerCertificate               bool   `default:"false" help:"If true server certificate is not verified for SSL. If false certificate will be verified against supplied certificate"`
+	Pgbouncer                            bool   `default:"false" help:"Collects metrics from PgBouncer instance. Assumes connection is through PgBouncer."`
+	CollectDbLockMetrics                 bool   `default:"false" help:"If true, enables collection of lock metrics for the specified database. (Note: requires that the 'tablefunc' extension is installed)"` //nolint: stylecheck
+	CollectBloatMetrics                  bool   `default:"true" help:"Enable collecting bloat metrics which can be performance intensive"`
+	ShowVersion                          bool   `default:"false" help:"Print build information and exit"`
+	EnableQueryMonitoring                bool   `default:"false" help:"Enable collection of detailed query performance metrics."`
+	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"`
 }
 
 // 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 err := al.validateSSL(); err != nil {
 		return err
 	}
-
 	return nil
 }
 

src/connection/pgsql_connection.go

@@ -149,7 +149,7 @@ func (p PGSQLConnection) HaveExtensionInSchema(extensionName, schemaName string)
 	return true
 }
 
-// createConnectionURL creates the connection string. A list of paramters
+// 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 {
 	connectionURL := &url.URL{
@@ -170,7 +170,6 @@ func createConnectionURL(ci *connectionInfo, database string) string {
 	}
 
 	connectionURL.RawQuery = query.Encode()
-
 	return connectionURL.String()
 }
 

src/main.go

@@ -7,6 +7,8 @@ import (
 	"runtime"
 	"strings"
 
+	queryperformancemonitoring "github.com/newrelic/nri-postgresql/src/query-performance-monitoring"
+
 	"github.com/newrelic/infra-integrations-sdk/v3/integration"
 	"github.com/newrelic/infra-integrations-sdk/v3/log"
 	"github.com/newrelic/nri-postgresql/src/args"
@@ -27,6 +29,7 @@ var (
 )
 
 func main() {
+
 	var args args.ArgumentList
 	// Create Integration
 	pgIntegration, err := integration.New(integrationName, integrationVersion, integration.Args(&args))
@@ -62,7 +65,6 @@ func main() {
 		log.Error("Error creating list of entities to collect: %s", err)
 		os.Exit(1)
 	}
-
 	instance, err := pgIntegration.Entity(fmt.Sprintf("%s:%s", args.Hostname, args.Port), "pg-instance")
 	if err != nil {
 		log.Error("Error creating instance entity: %s", err.Error())
@@ -89,4 +91,9 @@ func main() {
 	if err = pgIntegration.Publish(); err != nil {
 		log.Error(err.Error())
 	}
+
+	if args.EnableQueryMonitoring {
+		queryperformancemonitoring.QueryPerformanceMain(args, pgIntegration, collectionList)
+	}
+
 }

src/metrics/metrics.go

@@ -37,7 +37,7 @@ func PopulateMetrics(
 	}
 	defer con.Close()
 
-	version, err := collectVersion(con)
+	version, err := CollectVersion(con)
 	if err != nil {
 		log.Error("Metrics collection failed: error collecting version number: %s", err.Error())
 		return
@@ -223,15 +223,14 @@ type serverVersionRow struct {
 	Version string `db:"server_version"`
 }
 
-func collectVersion(connection *connection.PGSQLConnection) (*semver.Version, error) {
+func CollectVersion(connection *connection.PGSQLConnection) (*semver.Version, error) {
 	var versionRows []*serverVersionRow
 	if err := connection.Query(&versionRows, versionQuery); err != nil {
 		return nil, err
 	}
 
 	re := regexp.MustCompile(`[0-9]+\.[0-9]+(\.[0-9])?`)
 	version := re.FindString(versionRows[0].Version)
-
 	// special cases for ubuntu/debian parsing
 	//version := versionRows[0].Version
 	//if strings.Contains(version, "Ubuntu") {

src/metrics/version_test.go

@@ -22,7 +22,7 @@ func Test_collectVersion(t *testing.T) {
 		Minor: 3,
 	}
 
-	version, err := collectVersion(testConnection)
+	version, err := CollectVersion(testConnection)
 
 	assert.Nil(t, err)
 	assert.Equal(t, expected, version)
@@ -42,7 +42,7 @@ func Test_collectVersion_EnterpriseDB(t *testing.T) {
 		Patch: 7,
 	}
 
-	version, err := collectVersion(testConnection)
+	version, err := CollectVersion(testConnection)
 
 	assert.Nil(t, err)
 	assert.Equal(t, expected, version)
@@ -61,7 +61,7 @@ func Test_collectVersion_Ubuntu(t *testing.T) {
 		Minor: 4,
 	}
 
-	version, err := collectVersion(testConnection)
+	version, err := CollectVersion(testConnection)
 
 	assert.Nil(t, err)
 	assert.Equal(t, expected, version)
@@ -80,7 +80,7 @@ func Test_collectVersion_Debian(t *testing.T) {
 		Minor: 4,
 	}
 
-	version, err := collectVersion(testConnection)
+	version, err := CollectVersion(testConnection)
 
 	assert.Nil(t, err)
 	assert.Equal(t, expected, version)
@@ -94,7 +94,7 @@ func Test_collectVersion_Err(t *testing.T) {
 
 	mock.ExpectQuery(versionQuery).WillReturnRows(versionRows)
 
-	_, err := collectVersion(testConnection)
+	_, err := CollectVersion(testConnection)
 
 	assert.NotNil(t, err)
 }

src/query-performance-monitoring/common-parameters/common_parameters.go

@@ -0,0 +1,55 @@
+package commonparameters
+
+import (
+	"github.com/newrelic/infra-integrations-sdk/v3/log"
+	"github.com/newrelic/nri-postgresql/src/args"
+)
+
+// The maximum number records that can be fetched in a single metrics
+const MaxQueryCountThreshold = 30
+
+// DefaultQueryMonitoringCountThreshold is the default threshold for the number of queries to monitor.
+const DefaultQueryMonitoringCountThreshold = 20
+
+// DefaultQueryResponseTimeThreshold is the default threshold for the response time of a query.
+const DefaultQueryResponseTimeThreshold = 500
+
+type CommonParameters struct {
+	Version                              uint64
+	Databases                            string
+	QueryMonitoringCountThreshold        int
+	QueryMonitoringResponseTimeThreshold int
+	Host                                 string
+	Port                                 string
+}
+
+func SetCommonParameters(args args.ArgumentList, version uint64, databases string) *CommonParameters {
+	return &CommonParameters{
+		Version:                              version,
+		Databases:                            databases, // comma separated database names
+		QueryMonitoringCountThreshold:        validateAndGetQueryMonitoringCountThreshold(args),
+		QueryMonitoringResponseTimeThreshold: validateAndGetQueryMonitoringResponseTimeThreshold(args),
+		Host:                                 args.Hostname,
+		Port:                                 args.Port,
+	}
+}
+
+func validateAndGetQueryMonitoringResponseTimeThreshold(args args.ArgumentList) int {
+	if args.QueryMonitoringResponseTimeThreshold < 0 {
+		log.Warn("QueryResponseTimeThreshold should be greater than or equal to 0 but the input is %d, setting value to default which is %d", args.QueryMonitoringResponseTimeThreshold, DefaultQueryResponseTimeThreshold)
+		return DefaultQueryResponseTimeThreshold
+	}
+	return args.QueryMonitoringResponseTimeThreshold
+}
+
+func validateAndGetQueryMonitoringCountThreshold(args args.ArgumentList) int {
+	if args.QueryMonitoringCountThreshold < 0 {
+		log.Warn("QueryCountThreshold should be greater than 0 but the input is %d, setting value to default which is %d", args.QueryMonitoringCountThreshold, DefaultQueryMonitoringCountThreshold)
+		return DefaultQueryMonitoringCountThreshold
+	}
+	if args.QueryMonitoringCountThreshold > MaxQueryCountThreshold {
+		log.Warn("QueryCountThreshold should be less than or equal to max limit but the input is %d, setting value to max limit which is %d", args.QueryMonitoringCountThreshold, MaxQueryCountThreshold)
+		return MaxQueryCountThreshold
+	}
+	return args.QueryMonitoringCountThreshold
+}

src/query-performance-monitoring/common-utils/common_helpers.go

@@ -0,0 +1,42 @@
+package commonutils
+
+import (
+	"crypto/rand"
+	"fmt"
+	"math/big"
+	"regexp"
+	"strings"
+	"time"
+
+	"github.com/newrelic/nri-postgresql/src/collection"
+)
+
+// re is a regular expression that matches single-quoted strings, numbers, or double-quoted strings
+var re = regexp.MustCompile(`'[^']*'|\d+|".*?"`)
+
+func GetDatabaseListInString(dbMap collection.DatabaseList) string {
+	if len(dbMap) == 0 {
+		return ""
+	}
+	var quotedNames = make([]string, 0)
+	for dbName := range dbMap {
+		quotedNames = append(quotedNames, fmt.Sprintf("'%s'", dbName))
+	}
+	return strings.Join(quotedNames, ",")
+}
+
+func AnonymizeQueryText(query string) string {
+	anonymizedQuery := re.ReplaceAllString(query, "?")
+	return anonymizedQuery
+}
+
+// This function is used to generate a unique plan ID for a query
+func GeneratePlanID() (string, error) {
+	randomInt, err := rand.Int(rand.Reader, big.NewInt(RandomIntRange))
+	if err != nil {
+		return "", ErrUnExpectedError
+	}
+	currentTime := time.Now().Format(TimeFormat)
+	result := fmt.Sprintf("%d-%s", randomInt.Int64(), currentTime)
+	return result, nil
+}

src/query-performance-monitoring/common-utils/common_helpers_test.go

@@ -0,0 +1,38 @@
+package commonutils
+
+import (
+	"sort"
+	"testing"
+
+	"github.com/newrelic/nri-postgresql/src/collection"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetDatabaseListInString(t *testing.T) {
+	dbListKeys := []string{"db1"}
+	sort.Strings(dbListKeys) // Sort the keys to ensure consistent order
+	dbList := collection.DatabaseList{}
+	for _, key := range dbListKeys {
+		dbList[key] = collection.SchemaList{}
+	}
+	expected := "'db1'"
+	result := GetDatabaseListInString(dbList)
+	assert.Equal(t, expected, result)
+
+	// Test with empty database list
+	dbList = collection.DatabaseList{}
+	expected = ""
+	result = GetDatabaseListInString(dbList)
+	assert.Equal(t, expected, result)
+}
+
+func TestAnonymizeQueryText(t *testing.T) {
+	query := "SELECT * FROM users WHERE id = 1 AND name = 'John'"
+	expected := "SELECT * FROM users WHERE id = ? AND name = ?"
+	result := AnonymizeQueryText(query)
+	assert.Equal(t, expected, result)
+	query = "SELECT * FROM employees WHERE id = 10 OR name <> 'John Doe'   OR name != 'John Doe'   OR age < 30 OR age <= 30   OR salary > 50000OR salary >= 50000  OR department LIKE 'Sales%' OR department ILIKE 'sales%'OR join_date BETWEEN '2023-01-01' AND '2023-12-31' OR department IN ('HR', 'Engineering', 'Marketing') OR department IS NOT NULL OR department IS NULL;"
+	expected = "SELECT * FROM employees WHERE id = ? OR name <> ?   OR name != ?   OR age < ? OR age <= ?   OR salary > ?OR salary >= ?  OR department LIKE ? OR department ILIKE ?OR join_date BETWEEN ? AND ? OR department IN (?, ?, ?) OR department IS NOT NULL OR department IS NULL;"
+	result = AnonymizeQueryText(query)
+	assert.Equal(t, expected, result)
+}