From 47fc20ea6c83a1acd3302c907c95ce09c33c90a0 Mon Sep 17 00:00:00 2001 From: Lukas Fittl Date: Tue, 10 Dec 2024 15:31:12 -0800 Subject: [PATCH] WIP --- input/postgres/explain.go | 174 +++++++++++++++++++++++++------------- runner/logs.go | 2 +- runner/query_run.go | 20 +---- 3 files changed, 120 insertions(+), 76 deletions(-) diff --git a/input/postgres/explain.go b/input/postgres/explain.go index 4ea3ee7fc..daa3b7775 100644 --- a/input/postgres/explain.go +++ b/input/postgres/explain.go @@ -14,7 +14,7 @@ import ( "github.com/pganalyze/collector/util" ) -func RunExplain(ctx context.Context, server *state.Server, inputs []state.PostgresQuerySample, collectionOpts state.CollectionOpts, logger *util.Logger) (outputs []state.PostgresQuerySample) { +func RunExplainForSamples(ctx context.Context, server *state.Server, inputs []state.PostgresQuerySample, collectionOpts state.CollectionOpts, logger *util.Logger) (outputs []state.PostgresQuerySample) { var samplesByDb = make(map[string]([]state.PostgresQuerySample)) skip := func(sample state.PostgresQuerySample) bool { @@ -84,68 +84,19 @@ func RunExplain(ctx context.Context, server *state.Server, inputs []state.Postgr func runDbExplain(ctx context.Context, db *sql.DB, inputs []state.PostgresQuerySample, useHelper bool) ([]state.PostgresQuerySample, error) { var outputs []state.PostgresQuerySample for _, sample := range inputs { - var isUtil []bool - // To be on the safe side never EXPLAIN a statement that can't be parsed, - // or multiple statements in one (leading to accidental execution) - isUtil, err := util.IsUtilityStmt(sample.Query) - if err == nil && len(isUtil) == 1 && !isUtil[0] { - var explainOutput []byte - + isExplainable, explainJSON, explainError, err := runExplainOrExplainAnalyze(ctx, db, sample.Query, sample.Parameters, useHelper) + if err != nil { + return nil, err + } + if isExplainable { sample.HasExplain = true sample.ExplainSource = pganalyze_collector.QuerySample_STATEMENT_LOG_EXPLAIN_SOURCE sample.ExplainFormat = pganalyze_collector.QuerySample_JSON_EXPLAIN_FORMAT - - if useHelper { - err = db.QueryRowContext(ctx, QueryMarkerSQL+"SELECT pganalyze.explain($1, $2)", sample.Query, pq.Array(sample.Parameters)).Scan(&explainOutput) - if err != nil { - if ctx.Err() != nil { - return nil, err - } - sample.ExplainError = fmt.Sprintf("%s", err) - } - } else { - if len(sample.Parameters) > 0 { - _, err = db.ExecContext(ctx, QueryMarkerSQL+"PREPARE pganalyze_explain AS "+sample.Query) - if err != nil { - if ctx.Err() != nil { - return nil, err - } - sample.ExplainError = fmt.Sprintf("%s", err) - } else { - paramStr := getQuotedParamsStr(sample.Parameters) - err = db.QueryRowContext(ctx, QueryMarkerSQL+"EXPLAIN (VERBOSE, FORMAT JSON) EXECUTE pganalyze_explain("+paramStr+")").Scan(&explainOutput) - if err != nil { - if ctx.Err() != nil { - return nil, err - } - sample.ExplainError = fmt.Sprintf("%s", err) - } - - _, err = db.ExecContext(ctx, QueryMarkerSQL+"DEALLOCATE pganalyze_explain") - if err != nil { - return nil, err - } - } - } else { - err = db.QueryRowContext(ctx, QueryMarkerSQL+"EXPLAIN (VERBOSE, FORMAT JSON) "+sample.Query).Scan(&explainOutput) - if err != nil { - if ctx.Err() != nil { - return nil, err - } - sample.ExplainError = fmt.Sprintf("%s", err) - } - } + if explainJSON != nil { + sample.ExplainOutputJSON = explainJSON } - - if len(explainOutput) > 0 { - var explainOutputJSON []state.ExplainPlanContainer - if err := json.Unmarshal(explainOutput, &explainOutputJSON); err != nil { - sample.ExplainError = fmt.Sprintf("%s", err) - } else if len(explainOutputJSON) != 1 { - sample.ExplainError = fmt.Sprintf("Unexpected plan size: %d (expected 1)", len(explainOutputJSON)) - } else { - sample.ExplainOutputJSON = &explainOutputJSON[0] - } + if explainError != "" { + sample.ExplainError = explainError } } @@ -164,6 +115,111 @@ func contains(strs []string, val string) bool { return false } +func RunExplainAnalyzeForQueryRun(ctx context.Context, db *sql.DB, query string, parameters []null.String, useHelper bool) (string, error, error) { + var err error + var firstErr error + var result string + + // check for helper here (before calling runExplainOrExplainAnalyze) + + sql := "BEGIN; EXPLAIN (ANALYZE, VERBOSE, BUFFERS, FORMAT JSON) " + query + "; ROLLBACK" + err = db.QueryRowContext(ctx, sql).Scan(&result) + firstErr = err + + // Run EXPLAIN ANALYZE a second time to get a warm cache result + err = db.QueryRowContext(ctx, sql).Scan(&result) + + // If the first run failed and the second run succeeded, run once more to get a warm cache result + if err == nil && firstErr != nil { + err = db.QueryRowContext(ctx, sql).Scan(&result) + } + + // If it timed out, capture a non-ANALYZE EXPLAIN instead + if err != nil && strings.Contains(err.Error(), "statement timeout") { + sql = "BEGIN; EXPLAIN (VERBOSE, FORMAT JSON) " + query + "; ROLLBACK" + err = db.QueryRowContext(ctx, sql).Scan(&result) + } + + // Run one time with EXPLAIN ANALYZE (could do this one without buffers and timing, since we won't use them) + // Run again with EXPLAIN ANALYZE + buffers and timing + // If either run times out we should instead capture a regular EXPLAIN plan +} + +// The helper variants are: +// - have no helper +// - have explain helper +// - have explain analyze helper +// and the actions are: +// - do EXPLAIN +// - do EXPLAIN ANALYZE with timings and buffers +// - do EXPLAIN ANALYZE without timings and buffers + +func runExplainOrExplainAnalyze(ctx context.Context, db *sql.DB, query string, parameters []null.String, analyze bool, helperFunction string) (isExplainable bool, explainJSON *state.ExplainPlanContainer, explainError string, _err error) { + var isUtil []bool + // To be on the safe side never EXPLAIN a statement that can't be parsed, + // or multiple statements in one (leading to accidental execution) + isUtil, err := util.IsUtilityStmt(query) + if err == nil && len(isUtil) == 1 && !isUtil[0] { + var explainOutput []byte + isExplainable = true + + if useHelper { + err = db.QueryRowContext(ctx, QueryMarkerSQL+"SELECT pganalyze.explain($1, $2)", query, pq.Array(parameters)).Scan(&explainOutput) + if err != nil { + if ctx.Err() != nil { + return false, nil, "", err + } + explainError = fmt.Sprintf("%s", err) + } + } else { + if len(parameters) > 0 { + _, err = db.ExecContext(ctx, QueryMarkerSQL+"PREPARE pganalyze_explain AS "+query) + if err != nil { + if ctx.Err() != nil { + return false, nil, "", err + } + explainError = fmt.Sprintf("%s", err) + } else { + paramStr := getQuotedParamsStr(parameters) + err = db.QueryRowContext(ctx, QueryMarkerSQL+"EXPLAIN (VERBOSE, FORMAT JSON) EXECUTE pganalyze_explain("+paramStr+")").Scan(&explainOutput) + if err != nil { + if ctx.Err() != nil { + return false, nil, "", err + } + explainError = fmt.Sprintf("%s", err) + } + + _, err = db.ExecContext(ctx, QueryMarkerSQL+"DEALLOCATE pganalyze_explain") + if err != nil { + return false, nil, "", err + } + } + } else { + err = db.QueryRowContext(ctx, QueryMarkerSQL+"EXPLAIN (VERBOSE, FORMAT JSON) "+query).Scan(&explainOutput) + if err != nil { + if ctx.Err() != nil { + return false, nil, "", err + } + explainError = fmt.Sprintf("%s", err) + } + } + } + + if len(explainOutput) > 0 { + var explainOutputJSON []state.ExplainPlanContainer + if err := json.Unmarshal(explainOutput, &explainOutputJSON); err != nil { + explainError = fmt.Sprintf("%s", err) + } else if len(explainOutputJSON) != 1 { + explainError = fmt.Sprintf("Unexpected plan size: %d (expected 1)", len(explainOutputJSON)) + } else { + explainJSON = &explainOutputJSON[0] + } + } + } + + return +} + func getQuotedParamsStr(parameters []null.String) string { params := []string{} for i := 0; i < len(parameters); i++ { diff --git a/runner/logs.go b/runner/logs.go index 73b6ea241..8006c3002 100644 --- a/runner/logs.go +++ b/runner/logs.go @@ -285,7 +285,7 @@ func processLogStream(ctx context.Context, server *state.Server, logLines []stat func postprocessAndSendLogs(ctx context.Context, server *state.Server, globalCollectionOpts state.CollectionOpts, logger *util.Logger, transientLogState state.TransientLogState, grant state.Grant) (err error) { if server.Config.EnableLogExplain && len(transientLogState.QuerySamples) != 0 { - transientLogState.QuerySamples = postgres.RunExplain(ctx, server, transientLogState.QuerySamples, globalCollectionOpts, logger) + transientLogState.QuerySamples = postgres.RunExplainForSamples(ctx, server, transientLogState.QuerySamples, globalCollectionOpts, logger) } if server.Config.FilterQuerySample == "all" { diff --git a/runner/query_run.go b/runner/query_run.go index 5fd8c7c6f..e9cd41a1c 100644 --- a/runner/query_run.go +++ b/runner/query_run.go @@ -4,7 +4,6 @@ import ( "context" "errors" "fmt" - "strings" "time" "github.com/pganalyze/collector/input/postgres" @@ -88,22 +87,11 @@ func run(ctx context.Context, server *state.Server, collectionOpts state.Collect result := "" if query.Type == pganalyze_collector.QueryRunType_EXPLAIN { - sql := "BEGIN; EXPLAIN (ANALYZE, VERBOSE, BUFFERS, FORMAT JSON) " + comment + query.QueryText + "; ROLLBACK" - err = db.QueryRowContext(ctx, sql).Scan(&result) - firstErr = err + sql := comment + query.QueryText - // Run EXPLAIN ANALYZE a second time to get a warm cache result - err = db.QueryRowContext(ctx, sql).Scan(&result) - - // If the first run failed and the second run succeeded, run once more to get a warm cache result - if err == nil && firstErr != nil { - err = db.QueryRowContext(ctx, sql).Scan(&result) - } - - // If it timed out, capture a non-ANALYZE EXPLAIN instead - if err != nil && strings.Contains(err.Error(), "statement timeout") { - sql = "BEGIN; EXPLAIN (VERBOSE, FORMAT JSON) " + comment + query.QueryText + "; ROLLBACK" - err = db.QueryRowContext(ctx, sql).Scan(&result) + result, err, firstExecutionErr := postgres.RunExplainAnalyzeForQueryRun(ctx, server, sql, globalCollectionOpts, logger) + if firstExecutionErr != nil { + firstErr = firstExecutionErr } } else { err = errors.New("Unhandled query run type")