fix(collector): handle Aurora's unsupported pg_last_xact_replay_times…

collector/pg_replication.go

@@ -15,7 +15,12 @@ package collector
 
 import (
 	"context"
+	"database/sql"
+	"errors"
+	"math"
+	"strings"
 
+	"github.com/lib/pq"
 	"github.com/prometheus/client_golang/prometheus"
 )
 
@@ -72,32 +77,71 @@ var (
 		ELSE 0
 	END as is_replica,
 	GREATEST (0, EXTRACT(EPOCH FROM (now() - pg_last_xact_replay_timestamp()))) as last_replay`
+
+	// pgReplicationIsReplicaQuery is a fallback used when the full query fails
+	// on Aurora PostgreSQL, which does not support pg_last_xact_replay_timestamp().
+	pgReplicationIsReplicaQuery = `SELECT CASE WHEN pg_is_in_recovery() THEN 1 ELSE 0 END as is_replica`
 )
 
+// isAuroraUnsupportedFunction returns true when Aurora PostgreSQL rejects a
+// query because it calls a function that is not supported on Aurora (e.g.
+// pg_last_xact_replay_timestamp). Aurora surfaces this as Postgres error class
+// "0A" (feature_not_supported) with a message that identifies the word "Aurora".
+func isAuroraUnsupportedFunction(err error) bool {
+	var pqErr *pq.Error
+	if errors.As(err, &pqErr) {
+		return pqErr.Code.Class() == "0A" && strings.Contains(pqErr.Message, "Aurora")
+	}
+	return false
+}
+
 func (c *PGReplicationCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error {
 	db := instance.getDB()
 	row := db.QueryRowContext(ctx,
 		pgReplicationQuery,
 	)
 
-	var lag float64
+	var lag sql.NullFloat64
 	var isReplica int64
-	var replayAge float64
+	var replayAge sql.NullFloat64
 	err := row.Scan(&lag, &isReplica, &replayAge)
 	if err != nil {
-		return err
+		if isAuroraUnsupportedFunction(err) {
+			// Aurora PostgreSQL does not support pg_last_xact_replay_timestamp().
+			// Emit NaN for the time-based metrics and fall back to a simpler query
+			// that still reports is_replica.
+			lag = sql.NullFloat64{Valid: false}
+			replayAge = sql.NullFloat64{Valid: false}
+
+			row2 := db.QueryRowContext(ctx, pgReplicationIsReplicaQuery)
+			if err2 := row2.Scan(&isReplica); err2 != nil {
+				isReplica = 0
+			}
+		} else {
+			return err
+		}
 	}
+
+	lagValue := math.NaN()
+	if lag.Valid {
+		lagValue = lag.Float64
+	}
+	replayAgeValue := math.NaN()
+	if replayAge.Valid {
+		replayAgeValue = replayAge.Float64
+	}
+
 	ch <- prometheus.MustNewConstMetric(
 		pgReplicationLag,
-		prometheus.GaugeValue, lag,
+		prometheus.GaugeValue, lagValue,
 	)
 	ch <- prometheus.MustNewConstMetric(
 		pgReplicationIsReplica,
 		prometheus.GaugeValue, float64(isReplica),
 	)
 	ch <- prometheus.MustNewConstMetric(
 		pgReplicationLastReplay,
-		prometheus.GaugeValue, replayAge,
+		prometheus.GaugeValue, replayAgeValue,
 	)
 	return nil
 }

collector/pg_replication_test.go

@@ -14,9 +14,11 @@ package collector
 
 import (
 	"context"
+	"math"
 	"testing"
 
 	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/lib/pq"
 	"github.com/prometheus/client_golang/prometheus"
 	dto "github.com/prometheus/client_model/go"
 	"github.com/smartystreets/goconvey/convey"
@@ -62,3 +64,52 @@ func TestPgReplicationCollector(t *testing.T) {
 		t.Errorf("there were unfulfilled exceptions: %s", err)
 	}
 }
+
+func TestPgReplicationCollectorAurora(t *testing.T) {
+	db, mock, err := sqlmock.New()
+	if err != nil {
+		t.Fatalf("Error opening a stub db connection: %s", err)
+	}
+	defer db.Close()
+
+	inst := &instance{db: db}
+
+	// Aurora rejects the main query because pg_last_xact_replay_timestamp() is
+	// not supported. The collector should fall back to the simpler is_replica
+	// query and emit NaN for the time-based metrics.
+	auroraErr := &pq.Error{
+		Code:    "0A000", // feature_not_supported
+		Message: "pg_last_xact_replay_timestamp() is currently not supported for Aurora",
+	}
+	mock.ExpectQuery(sanitizeQuery(pgReplicationQuery)).WillReturnError(auroraErr)
+
+	fallbackColumns := []string{"is_replica"}
+	fallbackRows := sqlmock.NewRows(fallbackColumns).AddRow(1)
+	mock.ExpectQuery(sanitizeQuery(pgReplicationIsReplicaQuery)).WillReturnRows(fallbackRows)
+
+	ch := make(chan prometheus.Metric, 3)
+	c := PGReplicationCollector{}
+	if err := c.Update(context.Background(), inst, ch); err != nil {
+		t.Fatalf("Unexpected error from Update on Aurora: %s", err)
+	}
+	close(ch)
+
+	metrics := make([]MetricResult, 0, 3)
+	for m := range ch {
+		metrics = append(metrics, readMetric(m))
+	}
+
+	convey.Convey("Aurora fallback metrics", t, func() {
+		convey.So(len(metrics), convey.ShouldEqual, 3)
+		// lag should be NaN
+		convey.So(math.IsNaN(metrics[0].value), convey.ShouldBeTrue)
+		// is_replica should be 1
+		convey.So(metrics[1].value, convey.ShouldEqual, 1)
+		// last_replay should be NaN
+		convey.So(math.IsNaN(metrics[2].value), convey.ShouldBeTrue)
+	})
+
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("there were unfulfilled expectations: %s", err)
+	}
+}