PMM-14527 Ability to have aurora specific queries.

cmd/postgres_exporter/postgres_exporter_test.go

@@ -409,7 +409,7 @@ func (s *FunctionalSuite) TestBooleanConversionToValueAndString(c *C) {
 func (s *FunctionalSuite) TestParseUserQueries(c *C) {
 	userQueriesData, err := os.ReadFile("./tests/user_queries_ok.yaml")
 	if err == nil {
-		metricMaps, newQueryOverrides, err := parseUserQueries(userQueriesData)
+		metricMaps, newQueryOverrides, err := parseUserQueries(userQueriesData, distributionStandard)
 		c.Assert(err, Equals, nil)
 		c.Assert(metricMaps, NotNil)
 		c.Assert(newQueryOverrides, NotNil)

cmd/postgres_exporter/queries.go

@@ -23,10 +23,18 @@ import (
 	"gopkg.in/yaml.v2"
 )
 
-// UserQuery represents a user defined query
+const (
+	distributionStandard = "standard"
+	distributionAurora   = "aurora"
+	// '!' is a reserved character indicating that the query is not supported on Aurora and should be skipped for Aurora instances.
+	notSupportedByAurora = "!"
+)
+
+// UserQuery represents a user defined query, including support for Aurora, if needed
 type UserQuery struct {
-	Query        string    `yaml:"query"`
-	Metrics      []Mapping `yaml:"metrics"`
+	Query        string    `yaml:"query"`         // Standard query
+	QueryAurora  string    `yaml:"query_aurora"`  // Aurora specific query
+	Metrics      []Mapping `yaml:"metrics"`       // Metrics to be collected
 	Master       bool      `yaml:"master"`        // Querying only for master database
 	CacheSeconds uint64    `yaml:"cache_seconds"` // Number of seconds to cache the namespace result metrics for.
 	RunOnServer  string    `yaml:"runonserver"`   // Querying to run on which server version
@@ -197,7 +205,7 @@ func makeQueryOverrideMap(pgVersion semver.Version, queryOverrides map[string][]
 	return resultMap
 }
 
-func parseUserQueries(content []byte) (map[string]intermediateMetricMap, map[string]string, error) {
+func parseUserQueries(content []byte, distribution string) (map[string]intermediateMetricMap, map[string]string, error) {
 	var userQueries UserQueries
 
 	err := yaml.Unmarshal(content, &userQueries)
@@ -211,7 +219,25 @@ func parseUserQueries(content []byte) (map[string]intermediateMetricMap, map[str
 
 	for metric, specs := range userQueries {
 		level.Debug(logger).Log("msg", "New user metric namespace from YAML metric", "metric", metric, "cache_seconds", specs.CacheSeconds)
-		newQueryOverrides[metric] = specs.Query
+
+		// Query selection logic:
+		// For Aurora: use query_aurora if defined and not empty, otherwise use query if defined and not empty.
+		// If query_aurora is set to '!', skip this query for Aurora (not supported).
+		// For standard (non-Aurora): always use query.
+		switch distribution {
+		case distributionAurora:
+			if specs.QueryAurora != "" {
+				if specs.QueryAurora == notSupportedByAurora {
+					continue
+				}
+				newQueryOverrides[metric] = specs.QueryAurora
+			} else {
+				newQueryOverrides[metric] = specs.Query
+			}
+		default:
+			newQueryOverrides[metric] = specs.Query
+		}
+
 		metricMap, ok := metricMaps[metric]
 		if !ok {
 			// Namespace for metric not found - add it.
@@ -251,7 +277,7 @@ func parseUserQueries(content []byte) (map[string]intermediateMetricMap, map[str
 // TODO: test code for all cu.
 // TODO: the YAML this supports is "non-standard" - we should move away from it.
 func addQueries(content []byte, pgVersion semver.Version, server *Server) error {
-	metricMaps, newQueryOverrides, err := parseUserQueries(content)
+	metricMaps, newQueryOverrides, err := parseUserQueries(content, server.distribution)
 	if err != nil {
 		return err
 	}

cmd/postgres_exporter/queries_test.go

@@ -0,0 +1,102 @@
+// Copyright 2021 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestParseUserQueries_DistributionSelection(t *testing.T) {
+	cases := []struct {
+		name         string
+		yamlInput    string
+		distribution string
+		wantQuery    string
+	}{
+		{
+			name: "Standard uses query",
+			yamlInput: `
+pg_replication:
+  query: "standard"
+  query_aurora: "aurora"
+`,
+			distribution: "standard",
+			wantQuery:    "standard",
+		},
+		{
+			name: "Aurora uses query_aurora",
+			yamlInput: `
+pg_replication:
+  query: "standard"
+  query_aurora: "aurora"
+`,
+			distribution: "aurora",
+			wantQuery:    "aurora",
+		},
+		{
+			name: "Aurora falls back to query",
+			yamlInput: `
+pg_replication:
+  query: "standard"
+`,
+			distribution: "aurora",
+			wantQuery:    "standard",
+		},
+		{
+			name: "Aurora skips if neither",
+			yamlInput: `
+pg_replication:
+`,
+			distribution: "aurora",
+			wantQuery:    "",
+		},
+		{
+			name: "Standard query only",
+			yamlInput: `
+pg_replication:
+  query: "standard"
+`,
+			distribution: "standard",
+			wantQuery:    "standard",
+		},
+		{
+			name: "Aurora query only",
+			yamlInput: `
+pg_replication:
+  query_aurora: "aurora"
+`,
+			distribution: "aurora",
+			wantQuery:    "aurora",
+		},
+		{
+			name: "Not supported by Aurora",
+			yamlInput: `
+pg_replication:
+  query: "standard"
+  query_aurora: "!"
+`,
+			distribution: "aurora",
+			wantQuery:    "",
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			_, metricsQueries, err := parseUserQueries([]byte(tc.yamlInput), tc.distribution)
+			require.NoError(t, err)
+			require.Equal(t, tc.wantQuery, metricsQueries["pg_replication"])
+		})
+	}
+}

cmd/postgres_exporter/server.go

@@ -27,10 +27,11 @@ import (
 // Server describes a connection to Postgres.
 // Also it contains metrics map and query overrides.
 type Server struct {
-	db          *sql.DB
-	labels      prometheus.Labels
-	master      bool
-	runonserver string
+	db           *sql.DB
+	distribution string
+	labels       prometheus.Labels
+	master       bool
+	runonserver  string
 
 	// Last version used to calculate metric map. If mismatch on scrape,
 	// then maps are recalculated.
@@ -80,7 +81,15 @@ func NewServer(dsn string, opts ...ServerOpt) (*Server, error) {
 		labels: prometheus.Labels{
 			serverLabelName: fingerprint,
 		},
-		metricCache: make(map[string]cachedMetrics),
+		metricCache:  make(map[string]cachedMetrics),
+		distribution: distributionStandard,
+	}
+
+	// Detect Aurora by checking if aurora_version function exists
+	row := db.QueryRow("SELECT to_regproc('aurora_version') IS NOT NULL;")
+	var isAurora bool
+	if err := row.Scan(&isAurora); err == nil && isAurora {
+		s.distribution = distributionAurora
 	}
 
 	for _, opt := range opts {

percona_tests/custom-queries/medium-resolution/queries.yaml

@@ -1,6 +1,9 @@
 #### Queries are commented due to PMM-8859
 pg_replication:
-  query: "SELECT CASE WHEN NOT pg_is_in_recovery() THEN 0 ELSE GREATEST (0, EXTRACT(EPOCH FROM (now() - pg_last_xact_replay_timestamp()))) END AS lag"
+# Standard query
+- query: "SELECT CASE WHEN NOT pg_is_in_recovery() THEN 0 ELSE GREATEST (0, EXTRACT(EPOCH FROM (now() - pg_last_xact_replay_timestamp()))) END AS lag"
+# Aurora query
+- query_aurora: "SELECT CASE WHEN pg_is_in_recovery() THEN (SELECT COALESCE(MAX(replica_lag_in_msec), 0) / 1000.0 FROM aurora_replica_status()) ELSE 0 END AS lag"
   master: true
   metrics:
     - lag:

queries-hr.yml

@@ -31,6 +31,7 @@ pg_custom_replication_slots:
 
 pg_custom_stat_wal_receiver:
   master: true
+  # standard query
   query: |
     SELECT
       status,
@@ -43,6 +44,8 @@ pg_custom_stat_wal_receiver:
       EXTRACT(EPOCH FROM last_msg_send_time) AS last_msg_send_time_seconds,
       EXTRACT(EPOCH FROM last_msg_receipt_time) AS last_msg_receipt_time_seconds
     FROM pg_stat_wal_receiver;
+  # skipped for Aurora by query_aurora == '!'
+  query_aurora: "!"
   metrics:
     - status:
         usage: "LABEL"

queries-mr.yaml

@@ -1,6 +1,9 @@
 #### Queries are commented due to PMM-8859
 pg_replication:
+  # Standard query
   query: "SELECT CASE WHEN NOT pg_is_in_recovery() THEN 0 ELSE GREATEST (0, EXTRACT(EPOCH FROM (now() - pg_last_xact_replay_timestamp()))) END AS lag"
+  # Aurora query
+  query_aurora: "SELECT CASE WHEN pg_is_in_recovery() THEN (SELECT COALESCE(MAX(replica_lag_in_msec), 0) / 1000.0 FROM aurora_replica_status()) ELSE 0 END AS lag"
   master: true
   metrics:
     - lag: