From 67304c8d775034b13997c4fa90fc3d51498a12af Mon Sep 17 00:00:00 2001 From: Ryan Geyer Date: Mon, 2 Feb 2026 10:59:13 -0800 Subject: [PATCH] fix: Improvements to SET search_path for postgres explain plans (#5422) ### Brief description of Pull Request During dogfooding we discovered some cases where explain plans were failing to run due to missing or ambiguous schema names. This PR addresses several small issues with the `SET search_path` command which is meant to address these issues. ### Pull Request Details - Sets the search path first, before any other requests are issued, including creating the prepared statement. - Places the datname in double quotes to prevent syntax errors - Explicitly sets the search path for SESSION rather than depending upon the default. ### PR Checklist - [x] Tests updated (cherry picked from commit 0044e0365710c7394c994c57690a585caab22440) --- .../postgres/collector/explain_plans.go | 10 +++++----- .../postgres/collector/explain_plans_test.go | 6 +++--- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/internal/component/database_observability/postgres/collector/explain_plans.go b/internal/component/database_observability/postgres/collector/explain_plans.go index 9cf9f4d9fa2..50a2f5b0802 100644 --- a/internal/component/database_observability/postgres/collector/explain_plans.go +++ b/internal/component/database_observability/postgres/collector/explain_plans.go @@ -568,6 +568,11 @@ func (c *ExplainPlans) fetchExplainPlanJSON(ctx context.Context, qi queryInfo) ( } defer conn.Close() + setSearchPathStatement := fmt.Sprintf("SET SESSION search_path TO \"%s\", public", qi.datname) + if _, err := conn.ExecContext(ctx, setSearchPathStatement); err != nil { + return nil, fmt.Errorf("failed to set search path: %w", err) + } + preparedStatementName := strings.ReplaceAll(fmt.Sprintf("explain_plan_%s", qi.queryId), "-", "_") preparedStatementText := fmt.Sprintf("PREPARE %s AS %s", preparedStatementName, qi.queryText) logger := log.With(c.logger, "query_id", qi.queryId, "datname", qi.datname, "preparedStatementName", preparedStatementName, "preparedStatementText", preparedStatementText) @@ -581,11 +586,6 @@ func (c *ExplainPlans) fetchExplainPlanJSON(ctx context.Context, qi queryInfo) ( } }() - setSearchPathStatement := fmt.Sprintf("SET search_path TO %s, public", qi.datname) - if _, err := conn.ExecContext(ctx, setSearchPathStatement); err != nil { - return nil, fmt.Errorf("failed to set search path: %w", err) - } - if _, err := conn.ExecContext(ctx, "SET plan_cache_mode = force_generic_plan"); err != nil { return nil, fmt.Errorf("failed to set plan cache mode: %w", err) } 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 f29a916ef45..c2daa6b5dc8 100644 --- a/internal/component/database_observability/postgres/collector/explain_plans_test.go +++ b/internal/component/database_observability/postgres/collector/explain_plans_test.go @@ -2915,8 +2915,8 @@ func TestExplainPlanFetchExplainPlans(t *testing.T) { require.Equal(t, "complex_aggregation_with_case.json", jsonFile.Name) jsonData := jsonFile.Data + mock.ExpectExec("SET SESSION search_path TO \"testdb\", public").WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectExec("PREPARE explain_plan_123456 AS select * from some_table where id = $1").WillReturnResult(sqlmock.NewResult(0, 1)) - mock.ExpectExec("SET search_path TO testdb, public").WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectExec("SET plan_cache_mode = force_generic_plan").WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectQuery("EXPLAIN (FORMAT JSON) EXECUTE explain_plan_123456(null)").WillReturnRows(sqlmock.NewRows([]string{"json"}).AddRow(jsonData)) mock.ExpectExec("DEALLOCATE explain_plan_123456").WillReturnResult(sqlmock.NewResult(0, 1)) @@ -2982,8 +2982,8 @@ func TestExplainPlanFetchExplainPlans(t *testing.T) { require.Equal(t, "complex_aggregation_with_case.json", jsonFile.Name) jsonData := jsonFile.Data + mock.ExpectExec("SET SESSION search_path TO \"testdb\", public").WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectExec("PREPARE explain_plan_123456 AS with cte as (select * from some_table where id = $1) select * from cte").WillReturnResult(sqlmock.NewResult(0, 1)) - mock.ExpectExec("SET search_path TO testdb, public").WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectExec("SET plan_cache_mode = force_generic_plan").WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectQuery("EXPLAIN (FORMAT JSON) EXECUTE explain_plan_123456(null)").WillReturnRows(sqlmock.NewRows([]string{"json"}).AddRow(jsonData)) mock.ExpectExec("DEALLOCATE explain_plan_123456").WillReturnResult(sqlmock.NewResult(0, 1)) @@ -3049,8 +3049,8 @@ func TestExplainPlanFetchExplainPlans(t *testing.T) { require.Equal(t, "complex_aggregation_with_case.json", jsonFile.Name) jsonData := jsonFile.Data + mock.ExpectExec("SET SESSION search_path TO \"testdb\", public").WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectExec("PREPARE explain_plan_123456 AS select * from some_table").WillReturnResult(sqlmock.NewResult(0, 1)) - mock.ExpectExec("SET search_path TO testdb, public").WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectExec("SET plan_cache_mode = force_generic_plan").WillReturnResult(sqlmock.NewResult(0, 1)) mock.ExpectQuery("EXPLAIN (FORMAT JSON) EXECUTE explain_plan_123456").WillReturnRows(sqlmock.NewRows([]string{"json"}).AddRow(jsonData)) mock.ExpectExec("DEALLOCATE explain_plan_123456").WillReturnResult(sqlmock.NewResult(0, 1))