PMM-13579 Fix truncate of queries. (#3419)

* PMM-13579 Fix truncate of queries. * PMM-13579 Fix tests.
percona · Mar 3, 2025 · 232fa3c · 232fa3c
1 parent 0766be6
commit 232fa3c
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 8 deletions.
diff --git a/agent/agents/postgres/pgstatstatements/pg_stat_statements_models.go b/agent/agents/postgres/pgstatstatements/pg_stat_statements_models.go
@@ -214,6 +214,7 @@ type pgStatStatementsExtended struct {
 	Tables           []string
 	IsQueryTruncated bool
 	Comments         map[string]string
+	RealQuery        string // RealQuery is a query which is not truncated, it's not sent to API and stored only locally in memory.
 }
 
 func (e *pgStatStatementsExtended) String() string {

diff --git a/agent/agents/postgres/pgstatstatements/pgstatstatements.go b/agent/agents/postgres/pgstatstatements/pgstatstatements.go
@@ -143,7 +143,7 @@ func (m *PGStatStatementsQAN) Run(ctx context.Context) {
 	var err error
 	m.changes <- agents.Change{Status: inventoryv1.AgentStatus_AGENT_STATUS_STARTING}
 
-	if current, _, err := m.getStatStatementsExtended(ctx, m.q, m.maxQueryLength); err == nil {
+	if current, _, err := m.getStatStatementsExtended(ctx); err == nil {
 		if err = m.statementsCache.Set(current); err == nil {
 			m.l.Debugf("Got %d initial stat statements.", len(current))
 			running = true
@@ -204,8 +204,6 @@ func (m *PGStatStatementsQAN) Run(ctx context.Context) {
 // and the previous cashed state.
 func (m *PGStatStatementsQAN) getStatStatementsExtended(
 	ctx context.Context,
-	q *reform.Querier,
-	maxQueryLength int32,
 ) (statementsMap, statementsMap, error) {
 	var totalN, newN, newSharedN, oldN int
 	var err error
@@ -221,6 +219,8 @@ func (m *PGStatStatementsQAN) getStatStatementsExtended(
 		return nil, nil, err
 	}
 
+	q := m.q
+
 	// load all databases and usernames first as we can't use querier while iterating over rows below
 	databases := queryDatabases(q)
 	usernames := queryUsernames(q)
@@ -251,6 +251,7 @@ func (m *PGStatStatementsQAN) getStatStatementsExtended(
 			pgStatStatements: *row,
 			Database:         databases[row.DBID],
 			Username:         usernames[row.UserID],
+			RealQuery:        row.Query,
 		}
 
 		if p := prev[c.QueryID]; p != nil {
@@ -262,7 +263,7 @@ func (m *PGStatStatementsQAN) getStatStatementsExtended(
 		} else {
 			newN++
 
-			c.Query, c.IsQueryTruncated = truncate.Query(c.Query, maxQueryLength, truncate.GetDefaultMaxQueryLength())
+			c.Query, c.IsQueryTruncated = truncate.Query(c.Query, m.maxQueryLength, truncate.GetDefaultMaxQueryLength())
 		}
 
 		current[c.QueryID] = c
@@ -278,7 +279,7 @@ func (m *PGStatStatementsQAN) getStatStatementsExtended(
 }
 
 func (m *PGStatStatementsQAN) getNewBuckets(ctx context.Context, periodStart time.Time, periodLengthSecs uint32) ([]*agentv1.MetricsBucket, error) {
-	current, prev, err := m.getStatStatementsExtended(ctx, m.q, m.maxQueryLength)
+	current, prev, err := m.getStatStatementsExtended(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -336,7 +337,7 @@ func (m *PGStatStatementsQAN) makeBuckets(current, prev statementsMap) []*agentv
 		}
 
 		if len(currentPSS.Tables) == 0 {
-			currentPSS.Tables = extractTables(currentPSS.Query, m.maxQueryLength, l)
+			currentPSS.Tables = extractTables(currentPSS.RealQuery, m.maxQueryLength, l)
 		}
 
 		if !m.disableCommentsParsing {

diff --git a/agent/agents/postgres/pgstatstatements/pgstatstatements_test.go b/agent/agents/postgres/pgstatstatements/pgstatstatements_test.go
@@ -285,7 +285,7 @@ func TestPGStatStatementsQAN(t *testing.T) {
 			Common: &agentv1.MetricsBucket_Common{
 				Fingerprint:         selectAllCitiesLong,
 				Database:            "pmm-agent",
-				Tables:              []string{},
+				Tables:              []string{"city"},
 				Comments:            map[string]string{"controller": "test"},
 				Username:            "pmm-agent",
 				AgentId:             "agent_id",
@@ -331,7 +331,7 @@ func TestPGStatStatementsQAN(t *testing.T) {
 			Common: &agentv1.MetricsBucket_Common{
 				Fingerprint:         selectAllCitiesLong,
 				Database:            "pmm-agent",
-				Tables:              []string{},
+				Tables:              []string{"city"},
 				Comments:            map[string]string{"controller": "test"},
 				Username:            "pmm-agent",
 				AgentId:             "agent_id",