caching obfuscated queries from pg_stat_statements

coroot · Mar 13, 2023 · 5f6170e · 5f6170e
1 parent e08ac37
commit 5f6170e
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 9 deletions.
diff --git 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
@@ -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,17 +48,23 @@ 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
 		}
 		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