From 5f6170e2006e2fed5d90dddc981959eea128ae27 Mon Sep 17 00:00:00 2001
From: Nikolay Sivko <n.sivko@gmail.com>
Date: Mon, 13 Mar 2023 11:43:41 +0300
Subject: [PATCH] caching obfuscated queries from `pg_stat_statements`

---
 collector/collector.go          |  6 +++++-
 collector/pg_stat_statements.go | 22 ++++++++++++++--------
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/collector/collector.go b/collector/collector.go
index 4a4bbfe..ccda568 100644
--- a/collector/collector.go
+++ b/collector/collector.go
@@ -155,7 +155,11 @@ func (c *Collector) snapshot() {
 
 	c.ssPrev = c.ssCurr
 	c.saPrev = c.saCurr
-	c.ssCurr, err = c.getStatStatements(version, querySizeLimit)
+	prevStatements := map[statementId]ssRow{}
+	if c.ssPrev != nil {
+		prevStatements = c.ssPrev.rows
+	}
+	c.ssCurr, err = c.getStatStatements(version, querySizeLimit, prevStatements)
 	if err != nil {
 		c.logger.Warning(err)
 		return
diff --git a/collector/pg_stat_statements.go b/collector/pg_stat_statements.go
index 38597e3..7d13529 100644
--- a/collector/pg_stat_statements.go
+++ b/collector/pg_stat_statements.go
@@ -9,15 +9,15 @@ import (
 )
 
 type ssRow struct {
-	queryText sql.NullString
-	calls     sql.NullInt64
-	rows      sql.NullInt64
-	totalTime sql.NullFloat64
-	ioTime    sql.NullFloat64
+	obfuscatedQueryText string
+	calls               sql.NullInt64
+	rows                sql.NullInt64
+	totalTime           sql.NullFloat64
+	ioTime              sql.NullFloat64
 }
 
 func (r ssRow) QueryKey(id statementId) QueryKey {
-	return QueryKey{Query: obfuscate.Sql(r.queryText.String), User: id.user.String, DB: id.db.String}
+	return QueryKey{Query: r.obfuscatedQueryText, User: id.user.String, DB: id.db.String}
 }
 
 type statementId struct {
@@ -31,7 +31,7 @@ type ssSnapshot struct {
 	rows map[statementId]ssRow
 }
 
-func (c *Collector) getStatStatements(version semver.Version, querySizeLimit int) (*ssSnapshot, error) {
+func (c *Collector) getStatStatements(version semver.Version, querySizeLimit int, prev map[statementId]ssRow) (*ssSnapshot, error) {
 	snapshot := &ssSnapshot{ts: time.Now(), rows: map[statementId]ssRow{}}
 	var query string
 	switch {
@@ -48,10 +48,11 @@ func (c *Collector) getStatStatements(version semver.Version, querySizeLimit int
 		return nil, err
 	}
 	defer rows.Close()
+	var queryText sql.NullString
 	for rows.Next() {
 		var id statementId
 		r := ssRow{}
-		err := rows.Scan(&id.db, &id.user, &r.queryText, &id.id, &r.calls, &r.totalTime, &r.ioTime)
+		err := rows.Scan(&id.db, &id.user, &queryText, &id.id, &r.calls, &r.totalTime, &r.ioTime)
 		if err != nil {
 			c.logger.Warning("failed to scan pg_stat_statements row:", err)
 			continue
@@ -59,6 +60,11 @@ func (c *Collector) getStatStatements(version semver.Version, querySizeLimit int
 		if id.user.String == "" || id.db.String == "" || !id.id.Valid {
 			continue
 		}
+		if p, ok := prev[id]; ok {
+			r.obfuscatedQueryText = p.obfuscatedQueryText
+		} else {
+			r.obfuscatedQueryText = obfuscate.Sql(queryText.String)
+		}
 		snapshot.rows[id] = r
 	}
 	return snapshot, nil