Handle newline within auto_explain query text with text format (#542)

* Handle newline within auto_explain query text with text format

* Run workaround only with Heroku

Author: keiko713

logs/analyze_test.go

@@ -4196,3 +4196,62 @@ func TestAnalyzeLogLines(t *testing.T) {
 		}
 	}
 }
+
+var testsHeroku = []testpair{{ // Heroku text explain + new line in Query Text (due to the query being too long)
+	[]state.LogLine{{
+		Content: "duration: 1681.452 ms  plan:\n" +
+			"\tQuery Text: UPDATE pgbench_branches SET bbalance = bbalance +\n 2656 WHERE bid = 59;\n" +
+			"\tUpdate on public.pgbench_branches  (cost=0.27..8.29 rows=1 width=370) (actual rows=0 loops=1)\n" +
+			"\t  Buffers: shared hit=7\n" +
+			"\t  ->  Index Scan using pgbench_branches_pkey on public.pgbench_branches  (cost=0.27..8.29 rows=1 width=370) (actual rows=1 loops=1)\n" +
+			"\t        Output: bid, (bbalance + 2656), filler, ctid\n" +
+			"\t        Index Cond: (pgbench_branches.bid = 59)",
+	}},
+	[]state.LogLine{{
+		Query:              "UPDATE pgbench_branches SET bbalance = bbalance +\n 2656 WHERE bid = 59;",
+		Classification:     pganalyze_collector.LogLineInformation_STATEMENT_AUTO_EXPLAIN,
+		ReviewedForSecrets: true,
+		SecretMarkers: []state.LogSecretMarker{{
+			ByteStart: 30,
+			ByteEnd:   469,
+			Kind:      state.StatementTextLogSecret,
+		}},
+	}},
+	[]state.PostgresQuerySample{{
+		Query:         "UPDATE pgbench_branches SET bbalance = bbalance +\n 2656 WHERE bid = 59;",
+		RuntimeMs:     1681.452,
+		HasExplain:    true,
+		ExplainSource: pganalyze_collector.QuerySample_AUTO_EXPLAIN_EXPLAIN_SOURCE,
+		ExplainFormat: pganalyze_collector.QuerySample_TEXT_EXPLAIN_FORMAT,
+		ExplainOutputText: "Update on public.pgbench_branches  (cost=0.27..8.29 rows=1 width=370) (actual rows=0 loops=1)\n" +
+			"\t  Buffers: shared hit=7\n" +
+			"\t  ->  Index Scan using pgbench_branches_pkey on public.pgbench_branches  (cost=0.27..8.29 rows=1 width=370) (actual rows=1 loops=1)\n" +
+			"\t        Output: bid, (bbalance + 2656), filler, ctid\n" +
+			"\t        Index Cond: (pgbench_branches.bid = 59)",
+	}},
+},
+}
+
+func TestAnalyzeLogLinesHeroku(t *testing.T) {
+	t.Setenv("DYNO", "dummy")
+	t.Setenv("PORT", "dummy")
+	for _, pair := range testsHeroku {
+		l, s := logs.AnalyzeLogLines(pair.logLinesIn)
+
+		cfg := pretty.CompareConfig
+		cfg.SkipZeroFields = true
+
+		if diff := cfg.Compare(pair.logLinesOut, l); diff != "" {
+			t.Errorf("For %v: log lines diff: (-want +got)\n%s", pair.logLinesIn, diff)
+		}
+		if diff := cfg.Compare(pair.samplesOut, s); diff != "" {
+			t.Errorf("For %v: query samples diff: (-want +got)\n%s", pair.samplesOut, diff)
+		}
+
+		for idx, line := range pair.logLinesOut {
+			if !line.ReviewedForSecrets && line.LogLevel != pganalyze_collector.LogLineInformation_STATEMENT && line.LogLevel != pganalyze_collector.LogLineInformation_QUERY {
+				t.Errorf("Missing secret review for:\n%s %s\n", pair.logLinesIn[idx].LogLevel, pair.logLinesIn[idx].Content)
+			}
+		}
+	}
+}

logs/querysample/querysample.go

@@ -11,6 +11,7 @@ import (
 	"github.com/pganalyze/collector/logs/util"
 	"github.com/pganalyze/collector/output/pganalyze_collector"
 	"github.com/pganalyze/collector/state"
+	cUtil "github.com/pganalyze/collector/util"
 )
 
 func TransformAutoExplainToQuerySample(logLine state.LogLine, explainText string, queryRuntime string) (state.PostgresQuerySample, error) {
@@ -54,12 +55,22 @@ func transformExplainJSONToQuerySample(logLine state.LogLine, explainText string
 }
 
 var autoExplainTextPlanDetailsRegexp = regexp.MustCompile(`^Query Text: (.+)\s+([\s\S]+)`)
+var herokuAutoExplainWithTabRegexp = regexp.MustCompile(`^Query Text: ([^\t]+)\t([\s\S]+)`)
 
 func transformExplainTextToQuerySample(logLine state.LogLine, explainText string, queryRuntimeMs float64) (state.PostgresQuerySample, error) {
 	explainParts := autoExplainTextPlanDetailsRegexp.FindStringSubmatch(explainText)
 	if len(explainParts) != 3 {
 		return state.PostgresQuerySample{}, fmt.Errorf("auto_explain output doesn't match expected format")
 	}
+	// If EXPLAIN output's first char is not a capital letter (e.g. not something like "Update on" or "Index Scan"),
+	// likely it's hitting the Heroku's newline break in "Query Text:" chunk
+	// Handle the separation of the query and the explain output text with the tab for these cases
+	explainOutputFirstChar := explainParts[2][0]
+	if cUtil.IsHeroku() && !(explainOutputFirstChar >= 'A' && explainOutputFirstChar <= 'Z') {
+		if parts := herokuAutoExplainWithTabRegexp.FindStringSubmatch(explainText); len(parts) == 3 {
+			explainParts = parts
+		}
+	}
 	return state.PostgresQuerySample{
 		Query:             strings.TrimSpace(explainParts[1]),
 		RuntimeMs:         queryRuntimeMs,