From acc21a71f155473e275b2a7b6123c302c207ec08 Mon Sep 17 00:00:00 2001
From: Nevermind <79126473+NevermindZ4@users.noreply.github.com>
Date: Mon, 13 Jan 2025 22:15:03 +0100
Subject: [PATCH] pg_stat_statements PG17

---
 collector/pg_stat_statements.go      | 30 +++++++++++++++++--
 collector/pg_stat_statements_test.go | 43 ++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/collector/pg_stat_statements.go b/collector/pg_stat_statements.go
index 7926f533e..cabdc3619 100644
--- a/collector/pg_stat_statements.go
+++ b/collector/pg_stat_statements.go
@@ -112,12 +112,38 @@ var (
 		)
 	ORDER BY seconds_total DESC
 	LIMIT 100;`
+
+	pgStatStatementsQuery_PG17 = `SELECT
+		pg_get_userbyid(userid) as user,
+		pg_database.datname,
+		pg_stat_statements.queryid,
+		pg_stat_statements.calls as calls_total,
+		pg_stat_statements.total_exec_time / 1000.0 as seconds_total,
+		pg_stat_statements.rows as rows_total,
+		pg_stat_statements.shared_blk_read_time / 1000.0 as block_read_seconds_total,
+		pg_stat_statements.shared_blk_write_time / 1000.0 as block_write_seconds_total
+		FROM pg_stat_statements
+	JOIN pg_database
+		ON pg_database.oid = pg_stat_statements.dbid
+	WHERE
+		total_exec_time > (
+		SELECT percentile_cont(0.1)
+			WITHIN GROUP (ORDER BY total_exec_time)
+			FROM pg_stat_statements
+		)
+	ORDER BY seconds_total DESC
+	LIMIT 100;`
 )
 
 func (PGStatStatementsCollector) Update(ctx context.Context, instance *instance, ch chan<- prometheus.Metric) error {
-	query := pgStatStatementsQuery
-	if instance.version.GE(semver.MustParse("13.0.0")) {
+	var query string
+	switch {
+	case instance.version.GE(semver.MustParse("17.0.0")):
+		query = pgStatStatementsQuery_PG17
+	case instance.version.GE(semver.MustParse("13.0.0")):
 		query = pgStatStatementsNewQuery
+	default:
+		query = pgStatStatementsQuery
 	}
 
 	db := instance.getDB()
diff --git a/collector/pg_stat_statements_test.go b/collector/pg_stat_statements_test.go
index 08aba34c2..ea035f557 100644
--- a/collector/pg_stat_statements_test.go
+++ b/collector/pg_stat_statements_test.go
@@ -151,3 +151,46 @@ func TestPGStateStatementsCollectorNewPG(t *testing.T) {
 		t.Errorf("there were unfulfilled exceptions: %s", err)
 	}
 }
+
+func TestPGStateStatementsCollector_PG17(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, version: semver.MustParse("17.0.0")}
+
+	columns := []string{"user", "datname", "queryid", "calls_total", "seconds_total", "rows_total", "block_read_seconds_total", "block_write_seconds_total"}
+	rows := sqlmock.NewRows(columns).
+		AddRow("postgres", "postgres", 1500, 5, 0.4, 100, 0.1, 0.2)
+	mock.ExpectQuery(sanitizeQuery(pgStatStatementsQuery_PG17)).WillReturnRows(rows)
+
+	ch := make(chan prometheus.Metric)
+	go func() {
+		defer close(ch)
+		c := PGStatStatementsCollector{}
+
+		if err := c.Update(context.Background(), inst, ch); err != nil {
+			t.Errorf("Error calling PGStatStatementsCollector.Update: %s", err)
+		}
+	}()
+
+	expected := []MetricResult{
+		{labels: labelMap{"user": "postgres", "datname": "postgres", "queryid": "1500"}, metricType: dto.MetricType_COUNTER, value: 5},
+		{labels: labelMap{"user": "postgres", "datname": "postgres", "queryid": "1500"}, metricType: dto.MetricType_COUNTER, value: 0.4},
+		{labels: labelMap{"user": "postgres", "datname": "postgres", "queryid": "1500"}, metricType: dto.MetricType_COUNTER, value: 100},
+		{labels: labelMap{"user": "postgres", "datname": "postgres", "queryid": "1500"}, metricType: dto.MetricType_COUNTER, value: 0.1},
+		{labels: labelMap{"user": "postgres", "datname": "postgres", "queryid": "1500"}, metricType: dto.MetricType_COUNTER, value: 0.2},
+	}
+
+	convey.Convey("Metrics comparison", t, func() {
+		for _, expect := range expected {
+			m := readMetric(<-ch)
+			convey.So(expect, convey.ShouldResemble, m)
+		}
+	})
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Errorf("there were unfulfilled exceptions: %s", err)
+	}
+}