From a4dd6495278428eace54657ce98559bc9d8607a6 Mon Sep 17 00:00:00 2001 From: Cristian Greco Date: Mon, 3 Aug 2026 15:19:36 +0200 Subject: [PATCH 1/2] fix(database_observability.postgres): explain_plans collector for postgres Fix `explain_plans` collector so that a cached query is finalised immediately instead of being deferred (which prevented correct update of processedCount and thus made batch-size limit ineffective). --- .../postgres/collector/explain_plans.go | 203 +++++++++--------- 1 file changed, 101 insertions(+), 102 deletions(-) diff --git a/internal/component/database_observability/postgres/collector/explain_plans.go b/internal/component/database_observability/postgres/collector/explain_plans.go index 18707bd0602..9c9179bc963 100644 --- a/internal/component/database_observability/postgres/collector/explain_plans.go +++ b/internal/component/database_observability/postgres/collector/explain_plans.go @@ -497,8 +497,6 @@ func (c *ExplainPlans) fetchExplainPlans(ctx context.Context) error { processedCount := 0 for _, qi := range c.queryCache { - generatedAt := time.Now().Format(time.RFC3339) - nonRecoverableFailureOccurred := false if c.isThrottled(qi.uniqueKey) { c.markFinished(qi) delete(c.queryCache, qi.uniqueKey) @@ -507,126 +505,127 @@ func (c *ExplainPlans) fetchExplainPlans(ctx context.Context) error { if processedCount >= c.currentBatchSize { break } - logger := c.logger.With("query_id", qi.queryId) - defer func(nonRecoverableFailureOccurred *bool) { - if *nonRecoverableFailureOccurred { - c.queryDenylist[qi.uniqueKey] = struct{}{} - } else { - c.markFinished(qi) - } - delete(c.queryCache, qi.uniqueKey) - processedCount++ - }(&nonRecoverableFailureOccurred) - - if strings.HasSuffix(qi.queryText, "...") { - err := c.sendExplainPlansOutput( - qi.datname, - qi.queryId, - generatedAt, - database_observability.ExplainProcessingResultSkipped, - "query is truncated", - nil, - ) - if err != nil { - c.logger.Error("failed to send truncated query skip explain plan output", "err", err) - } - continue + if c.processExplainPlan(ctx, qi) { + c.queryDenylist[qi.uniqueKey] = struct{}{} + } else { + c.markFinished(qi) } + delete(c.queryCache, qi.uniqueKey) + processedCount++ + } - containsReservedWord, err := database_observability.ContainsReservedKeywords(qi.queryText, database_observability.ExplainReservedWordDenyList, sqllexer.DBMSPostgres) + return nil +} + +// processExplainPlan processes a single query from the cache. +// Returns true if the query encountered a non-recoverable failure and should be denylisted. +func (c *ExplainPlans) processExplainPlan(ctx context.Context, qi *queryInfo) bool { + generatedAt := time.Now().Format(time.RFC3339) + logger := c.logger.With("query_id", qi.queryId) + + if strings.HasSuffix(qi.queryText, "...") { + err := c.sendExplainPlansOutput( + qi.datname, + qi.queryId, + generatedAt, + database_observability.ExplainProcessingResultSkipped, + "query is truncated", + nil, + ) if err != nil { - logger.Error("failed to check for reserved keywords", "err", err) - err := c.sendExplainPlansOutput( - qi.datname, - qi.queryId, - generatedAt, - database_observability.ExplainProcessingResultError, - fmt.Sprintf("failed to check for reserved keywords: %s", err.Error()), - nil, - ) - if err != nil { - c.logger.Error("failed to send reserved keyword check error explain plan output", "err", err) - } - continue + c.logger.Error("failed to send truncated query skip explain plan output", "err", err) } + return false + } - if containsReservedWord { - err := c.sendExplainPlansOutput( - qi.datname, - qi.queryId, - generatedAt, - database_observability.ExplainProcessingResultSkipped, - "query contains reserved word", - nil, - ) - if err != nil { - c.logger.Error("failed to send reserved keyword check error explain plan output", "err", err) - } - continue + containsReservedWord, err := database_observability.ContainsReservedKeywords(qi.queryText, database_observability.ExplainReservedWordDenyList, sqllexer.DBMSPostgres) + if err != nil { + logger.Error("failed to check for reserved keywords", "err", err) + err := c.sendExplainPlansOutput( + qi.datname, + qi.queryId, + generatedAt, + database_observability.ExplainProcessingResultError, + fmt.Sprintf("failed to check for reserved keywords: %s", err.Error()), + nil, + ) + if err != nil { + c.logger.Error("failed to send reserved keyword check error explain plan output", "err", err) } + return false + } - logger = logger.With("datname", qi.datname) - - byteExplainPlanJSON, err := c.fetchExplainPlanJSON(ctx, *qi) + if containsReservedWord { + err := c.sendExplainPlansOutput( + qi.datname, + qi.queryId, + generatedAt, + database_observability.ExplainProcessingResultSkipped, + "query contains reserved word", + nil, + ) if err != nil { - logger.Debug("failed to fetch explain plan json bytes", "err", err) - for _, code := range unrecoverablePostgresSQLErrors { - if strings.Contains(err.Error(), code) { - nonRecoverableFailureOccurred = true - break - } - } - continue + c.logger.Error("failed to send reserved keyword check error explain plan output", "err", err) } + return false + } - if len(byteExplainPlanJSON) == 0 { - logger.Error("explain plan json bytes is empty") - nonRecoverableFailureOccurred = true - continue - } + logger = logger.With("datname", qi.datname) - if !utf8.Valid(byteExplainPlanJSON) { - logger.Error("explain plan json bytes is not valid UTF-8") - nonRecoverableFailureOccurred = true - continue + byteExplainPlanJSON, err := c.fetchExplainPlanJSON(ctx, *qi) + if err != nil { + logger.Debug("failed to fetch explain plan json bytes", "err", err) + for _, code := range unrecoverablePostgresSQLErrors { + if strings.Contains(err.Error(), code) { + return true + } } + return false + } - redactedByteExplainPlanJSON := database_observability.RedactSql(string(byteExplainPlanJSON)) + if len(byteExplainPlanJSON) == 0 { + logger.Error("explain plan json bytes is empty") + return true + } - logger.Debug("db native explain plan", "db_native_explain_plan", base64.StdEncoding.EncodeToString([]byte(redactedByteExplainPlanJSON))) + if !utf8.Valid(byteExplainPlanJSON) { + logger.Error("explain plan json bytes is not valid UTF-8") + return true + } - explainPlanOutput, genErr := newExplainPlanOutput(byteExplainPlanJSON) - explainPlanOutputJSON, err := json.Marshal(explainPlanOutput) - if err != nil { - logger.Error("failed to marshal explain plan output", "err", err) - nonRecoverableFailureOccurred = true - continue - } + redactedByteExplainPlanJSON := database_observability.RedactSql(string(byteExplainPlanJSON)) - if genErr != nil { - logger.Error( - "failed to create explain plan output", - "incomplete_explain_plan", base64.StdEncoding.EncodeToString(explainPlanOutputJSON), - "err", genErr, - ) - nonRecoverableFailureOccurred = true - continue - } + logger.Debug("db native explain plan", "db_native_explain_plan", base64.StdEncoding.EncodeToString([]byte(redactedByteExplainPlanJSON))) - if err := c.sendExplainPlansOutput( - qi.datname, - qi.queryId, - generatedAt, - database_observability.ExplainProcessingResultSuccess, - "", - explainPlanOutput, - ); err != nil { - c.logger.Error("failed to send explain plan output", "err", err) - } + explainPlanOutput, genErr := newExplainPlanOutput(byteExplainPlanJSON) + explainPlanOutputJSON, err := json.Marshal(explainPlanOutput) + if err != nil { + logger.Error("failed to marshal explain plan output", "err", err) + return true } - return nil + if genErr != nil { + logger.Error( + "failed to create explain plan output", + "incomplete_explain_plan", base64.StdEncoding.EncodeToString(explainPlanOutputJSON), + "err", genErr, + ) + return true + } + + if err := c.sendExplainPlansOutput( + qi.datname, + qi.queryId, + generatedAt, + database_observability.ExplainProcessingResultSuccess, + "", + explainPlanOutput, + ); err != nil { + c.logger.Error("failed to send explain plan output", "err", err) + } + + return false } // postgresPreparedStatementParamCount returns N for EXECUTE, where N is the highest From c764b14b333b5e9980c3a1069cbd64100ade90ae Mon Sep 17 00:00:00 2001 From: Cristian Greco Date: Tue, 4 Aug 2026 16:52:58 +0200 Subject: [PATCH 2/2] review feedback --- .../postgres/collector/explain_plans.go | 12 +++---- .../postgres/collector/explain_plans_test.go | 35 +++++++++++++++++++ 2 files changed, 41 insertions(+), 6 deletions(-) diff --git a/internal/component/database_observability/postgres/collector/explain_plans.go b/internal/component/database_observability/postgres/collector/explain_plans.go index 9c9179bc963..9f0f0e192ef 100644 --- a/internal/component/database_observability/postgres/collector/explain_plans.go +++ b/internal/component/database_observability/postgres/collector/explain_plans.go @@ -542,7 +542,7 @@ func (c *ExplainPlans) processExplainPlan(ctx context.Context, qi *queryInfo) bo containsReservedWord, err := database_observability.ContainsReservedKeywords(qi.queryText, database_observability.ExplainReservedWordDenyList, sqllexer.DBMSPostgres) if err != nil { logger.Error("failed to check for reserved keywords", "err", err) - err := c.sendExplainPlansOutput( + sendErr := c.sendExplainPlansOutput( qi.datname, qi.queryId, generatedAt, @@ -550,14 +550,14 @@ func (c *ExplainPlans) processExplainPlan(ctx context.Context, qi *queryInfo) bo fmt.Sprintf("failed to check for reserved keywords: %s", err.Error()), nil, ) - if err != nil { - c.logger.Error("failed to send reserved keyword check error explain plan output", "err", err) + if sendErr != nil { + logger.Error("failed to send reserved keyword check error explain plan output", "err", sendErr) } return false } if containsReservedWord { - err := c.sendExplainPlansOutput( + sendErr := c.sendExplainPlansOutput( qi.datname, qi.queryId, generatedAt, @@ -565,8 +565,8 @@ func (c *ExplainPlans) processExplainPlan(ctx context.Context, qi *queryInfo) bo "query contains reserved word", nil, ) - if err != nil { - c.logger.Error("failed to send reserved keyword check error explain plan output", "err", err) + if sendErr != nil { + logger.Error("failed to send reserved keyword skip explain plan output", "err", sendErr) } return false } diff --git a/internal/component/database_observability/postgres/collector/explain_plans_test.go b/internal/component/database_observability/postgres/collector/explain_plans_test.go index b07dcd02363..e28cd60a7f0 100644 --- a/internal/component/database_observability/postgres/collector/explain_plans_test.go +++ b/internal/component/database_observability/postgres/collector/explain_plans_test.go @@ -2914,6 +2914,41 @@ func TestPlanNode_ToExplainPlanOutputNode(t *testing.T) { assert.Equal(t, database_observability.ExplainPlanJoinAlgorithmHash, *result.Details.JoinAlgorithm) } +func TestExplainPlanBatchSizeLimitsProcessing(t *testing.T) { + lokiClient := loki.NewCollectingHandler() + defer lokiClient.Stop() + + c, err := NewExplainPlan(ExplainPlansArguments{ + Logger: logging.NewSlogNop(), + ScrapeInterval: time.Second, + PerScrapeRatio: 1, + EntryHandler: lokiClient, + DBVersion: "17.0", + }) + require.NoError(t, err) + + c.queryCache = map[string]*queryInfo{ + explainPlanQueryKey("db", "1"): newQueryInfo("db", "1", "select * from table_1 where ...", 1, time.Now()), + explainPlanQueryKey("db", "2"): newQueryInfo("db", "2", "select * from table_2 where ...", 1, time.Now()), + explainPlanQueryKey("db", "3"): newQueryInfo("db", "3", "select * from table_3 where ...", 1, time.Now()), + explainPlanQueryKey("db", "4"): newQueryInfo("db", "4", "select * from table_4 where ...", 1, time.Now()), + } + c.currentBatchSize = 2 + + require.NoError(t, c.fetchExplainPlans(t.Context())) + require.Len(t, c.queryCache, 2, "batch size limit should leave unprocessed items in cache") + require.Len(t, c.finishedQueryCache, 2) + require.Empty(t, c.queryDenylist) + require.Eventually( + t, + func() bool { return len(lokiClient.Received()) == 2 }, + 5*time.Second, + 10*time.Millisecond, + "expected exactly 2 Loki entries, got %d", + len(lokiClient.Received()), + ) +} + func TestExplainPlanFetchExplainPlans(t *testing.T) { db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual)) require.NoError(t, err)