PMM-14527 Ability to have aurora specific queries. (#373)

* PMM-14527 Ability to have aurora queries.

* PMM-14527 Functionality to skip query for Aurora completely.

* PMM-14527 Add test case for Aurora skip.

* PMM-14527 Better check if Aurora.

* PMM-14527 Check for func of Aurora.

* PMM-14527 Different check for func exists.

* PMM-14527 Bump minimum PG to 10.

* PMM-14527 Improvements.

* PMM-14527 Fix.

* PMM-14527 Fix.

* PMM-14527 Fix.

* PMM-14527 Fix.

Author: JiriCtvrtka

.promu.yml

@@ -2,7 +2,7 @@ go:
     # This must match .circle/config.yml.
     version: 1.21
 repository:
-    path: github.com/prometheus-community/postgres_exporter
+    path: github.com/percona/postgres_exporter
 build:
     binaries:
         - name: postgres_exporter

README.md

@@ -1,6 +1,6 @@
 [![Build Status](https://circleci.com/gh/prometheus-community/postgres_exporter.svg?style=svg)](https://circleci.com/gh/prometheus-community/postgres_exporter)
 [![Coverage Status](https://coveralls.io/repos/github/prometheus-community/postgres_exporter/badge.svg?branch=master)](https://coveralls.io/github/prometheus-community/postgres_exporter?branch=master)
-[![Go Report Card](https://goreportcard.com/badge/github.com/prometheus-community/postgres_exporter)](https://goreportcard.com/report/github.com/prometheus-community/postgres_exporter)
+[![Go Report Card](https://goreportcard.com/badge/github.com/percona/postgres_exporter)](https://goreportcard.com/report/github.com/percona/postgres_exporter)
 [![Docker Pulls](https://img.shields.io/docker/pulls/prometheuscommunity/postgres-exporter.svg)](https://hub.docker.com/r/prometheuscommunity/postgres-exporter/tags)
 
 # PostgreSQL Server Exporter
@@ -73,7 +73,7 @@ auth_modules:
 
 ## Building and running
 
-    git clone https://github.com/prometheus-community/postgres_exporter.git
+    git clone https://github.com/percona/postgres_exporter.git
     cd postgres_exporter
     make build
     ./postgres_exporter <flags>

cmd/postgres_exporter/main.go

@@ -23,7 +23,7 @@ import (
 	"github.com/alecthomas/kingpin/v2"
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
-	"github.com/prometheus-community/postgres_exporter/config"
+	"github.com/percona/postgres_exporter/config"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/collectors"
 	vc "github.com/prometheus/client_golang/prometheus/collectors/version"

cmd/postgres_exporter/percona_exporter.go

@@ -16,7 +16,7 @@ import (
 	"github.com/blang/semver/v4"
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
-	"github.com/prometheus-community/postgres_exporter/collector"
+	"github.com/percona/postgres_exporter/collector"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"golang.org/x/sync/semaphore"

cmd/postgres_exporter/postgres_exporter.go

@@ -81,7 +81,7 @@ type Mapping map[string]MappingOptions
 
 // Regex used to get the "short-version" from the postgres version field.
 var versionRegex = regexp.MustCompile(`^\w+ ((\d+)(\.\d+)?(\.\d+)?)`)
-var lowestSupportedVersion = semver.MustParse("9.1.0")
+var lowestSupportedVersion = semver.MustParse("10.0.0")
 
 // Parses the version of postgres into the short version string we can use to
 // match behaviors.

cmd/postgres_exporter/postgres_exporter_test.go

@@ -23,6 +23,7 @@ import (
 	"time"
 
 	"github.com/blang/semver/v4"
+	"github.com/percona/postgres_exporter/distribution"
 	"github.com/prometheus/client_golang/prometheus"
 	. "gopkg.in/check.v1"
 )
@@ -409,7 +410,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, distribution.Standard)
 		c.Assert(err, Equals, nil)
 		c.Assert(metricMaps, NotNil)
 		c.Assert(newQueryOverrides, NotNil)

cmd/postgres_exporter/probe.go

@@ -19,8 +19,8 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
-	"github.com/prometheus-community/postgres_exporter/collector"
-	"github.com/prometheus-community/postgres_exporter/config"
+	"github.com/percona/postgres_exporter/collector"
+	"github.com/percona/postgres_exporter/config"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"golang.org/x/sync/semaphore"

cmd/postgres_exporter/queries.go

@@ -21,12 +21,20 @@ import (
 	"github.com/blang/semver/v4"
 	"github.com/go-kit/log/level"
 	"gopkg.in/yaml.v2"
+
+	"github.com/percona/postgres_exporter/distribution"
+)
+
+const (
+	// '!' 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
+// 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, dist 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 dist {
+		case distribution.Aurora:
+			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,104 @@
+// Copyright (C) 2023 Percona LLC
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+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

@@ -21,16 +21,18 @@ import (
 
 	"github.com/blang/semver/v4"
 	"github.com/go-kit/log/level"
+	"github.com/percona/postgres_exporter/distribution"
 	"github.com/prometheus/client_golang/prometheus"
 )
 
 // 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.
@@ -82,6 +84,7 @@ func NewServer(dsn string, opts ...ServerOpt) (*Server, error) {
 		},
 		metricCache: make(map[string]cachedMetrics),
 	}
+	s.distribution = distribution.Get(dsn, db)
 
 	for _, opt := range opts {
 		opt(s)