fix(instrument): emit coverage signal after statement execution

Move the injected pg_notify() call from before the tracked statement to
after it so that coverage reflects successful execution rather than
attempted execution. Previously, if a statement raised an exception
caught by the caller, the signal had already been sent — recording false
coverage.

Terminal statements (RETURN, RAISE EXCEPTION, bare RAISE) keep the
signal before the statement because they exit the current scope and
any code placed after them would be unreachable. Non-fatal RAISE levels
(NOTICE, WARNING, INFO, LOG, DEBUG) are correctly identified as
non-terminal and receive after-execute signals.

Add isTerminalSegment() helper that inspects the first meaningful token
of a segment to distinguish terminal from non-terminal statements.

Add tests: TestIsTerminalSegment (17 cases),
TestInstrumentBody_CoverageAfterExecute
(verifies ordering for assignments vs RETURN),
TestInstrumentBody_RaiseExceptionBeforeNotify
(verifies RAISE EXCEPTION vs RAISE NOTICE ordering).

Author: pashagolub

internal/instrument/instrumenter.go

@@ -126,8 +126,13 @@ func instrumentBody(stmt *parser.Statement, filePath string, skipToBegin bool, n
 	segStart := -1
 
 	// emitSegment checks the segment between segStart..segEnd for
-	// executability and, if it qualifies, writes the gap + notify + segment
-	// into instrumentedBody.
+	// executability and, if it qualifies, writes the gap + signal + segment
+	// (or segment + signal) into instrumentedBody.
+	//
+	// For terminal statements (RETURN, RAISE EXCEPTION) the signal is
+	// emitted *before* the statement because nothing executes after them.
+	// For all other statements the signal is emitted *after* the statement
+	// so that coverage is recorded only on successful execution (B1 fix).
 	emitSegment := func(segEnd int) {
 		segText := bodyContent[segStart:segEnd]
 		if !isExecutableSegment(segText) {
@@ -160,11 +165,29 @@ func instrumentBody(stmt *parser.Statement, filePath string, skipToBegin bool, n
 			}
 		}
 
-		// Write notify call, then the original segment text.
-		fmt.Fprintf(&instrumentedBody, "%s%s pg_notify('pgcov', '%s');\n",
+		notifyCall := fmt.Sprintf("%s%s pg_notify('pgcov', '%s');",
 			indent, notifyCmd, strings.ReplaceAll(cp.SignalID, "'", "''"))
-		instrumentedBody.WriteString(segText)
-		lastWrittenPos = segEnd
+
+		if isTerminalSegment(segText) {
+			// Terminal statements (RETURN, RAISE EXCEPTION) exit the
+			// current scope — the signal must fire before the statement.
+			fmt.Fprintf(&instrumentedBody, "%s\n", notifyCall)
+			instrumentedBody.WriteString(segText)
+			lastWrittenPos = segEnd
+		} else {
+			// Non-terminal statements: emit the signal after the
+			// statement so coverage reflects successful execution.
+			instrumentedBody.WriteString(segText)
+			// Consume the semicolon that terminates this segment so
+			// the notify call sits between statement and next gap.
+			if segEnd < len(bodyContent) && bodyContent[segEnd] == ';' {
+				instrumentedBody.WriteByte(';')
+				lastWrittenPos = segEnd + 1
+			} else {
+				lastWrittenPos = segEnd
+			}
+			fmt.Fprintf(&instrumentedBody, "\n%s", notifyCall)
+		}
 	}
 
 	// Stream tokens one at a time – mirrors SplitStatements style.
@@ -251,6 +274,51 @@ func isExecutableSegment(segmentContent string) bool {
 	return true
 }
 
+// isTerminalSegment checks whether a segment starts with a terminal statement
+// (RETURN or RAISE EXCEPTION / bare RAISE) that exits the current scope.
+// NOTIFY calls placed after such statements would be unreachable.
+// Non-fatal RAISE levels (NOTICE, WARNING, INFO, LOG, DEBUG) are not terminal.
+func isTerminalSegment(segmentContent string) bool {
+	sc := pglex.NewScanner(segmentContent)
+
+	// Find the first non-comment token.
+	var first pglex.Token
+	for {
+		first = sc.Scan()
+		if first.Type == pglex.EOF {
+			return false
+		}
+		if first.Type != pglex.Comment {
+			break
+		}
+	}
+
+	if first.Type == pglex.KReturn {
+		return true
+	}
+
+	if first.Type == pglex.KRaise {
+		// RAISE is terminal unless followed by a non-fatal level.
+		for {
+			tok := sc.Scan()
+			if tok.Type == pglex.EOF {
+				return true // bare RAISE; — re-raise in exception handler
+			}
+			if tok.Type == pglex.Comment {
+				continue
+			}
+			switch tok.Type {
+			case pglex.KNotice, pglex.KWarning, pglex.KInfo, pglex.KLog, pglex.KDebug:
+				return false
+			default:
+				return true // RAISE EXCEPTION, RAISE 'message', etc.
+			}
+		}
+	}
+
+	return false
+}
+
 // getIndentation returns the leading whitespace of a line.
 func getIndentation(line string) string {
 	return line[:len(line)-len(strings.TrimLeft(line, " \t"))]

internal/instrument/plpgsql_ast_test.go

@@ -103,7 +103,8 @@ $$ LANGUAGE plpgsql;`
 		}
 	}
 
-	// Verify PERFORM statements are injected before the executable lines
+	// Verify PERFORM statements are injected after the executable lines (coverage-after-execute),
+	// except for terminal statements (RETURN) which keep PERFORM before.
 	// Note: After instrumentation, line numbers shift due to inserted PERFORM statements
 	// We just verify that PERFORM statements exist for each coverage point
 	for _, cp := range instrumented.Locations {
@@ -236,3 +237,133 @@ $$ LANGUAGE plpgsql;`
 	}
 	t.Logf("Coverage points: %d (may be 0 for malformed SQL)", len(instrumented.Locations))
 }
+
+func TestIsTerminalSegment(t *testing.T) {
+	tests := []struct {
+		name     string
+		segment  string
+		terminal bool
+	}{
+		{"RETURN value", "RETURN a + b", true},
+		{"RETURN bare", "RETURN", true},
+		{"RAISE EXCEPTION", "RAISE EXCEPTION 'error'", true},
+		{"RAISE with string (default EXCEPTION)", "RAISE 'something went wrong'", true},
+		{"RAISE bare re-raise", "RAISE", true},
+		{"RAISE NOTICE", "RAISE NOTICE 'hello'", false},
+		{"RAISE WARNING", "RAISE WARNING 'warn'", false},
+		{"RAISE INFO", "RAISE INFO 'info'", false},
+		{"RAISE LOG", "RAISE LOG 'log'", false},
+		{"RAISE DEBUG", "RAISE DEBUG 'debug'", false},
+		{"assignment", "discount_rate := 0.20", false},
+		{"IF block", "IF x > 0 THEN\n    y := 1", false},
+		{"PERFORM", "PERFORM some_function()", false},
+		{"SELECT", "SELECT 1 INTO result", false},
+		{"comment then RETURN", "-- comment\nRETURN 42", true},
+		{"comment then assignment", "-- comment\nx := 1", false},
+		{"empty", "", false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := isTerminalSegment(tt.segment)
+			if got != tt.terminal {
+				t.Errorf("isTerminalSegment(%q) = %v, want %v", tt.segment, got, tt.terminal)
+			}
+		})
+	}
+}
+
+func TestInstrumentBody_CoverageAfterExecute(t *testing.T) {
+	// Verify that for non-terminal statements, NOTIFY comes after the statement,
+	// and for RETURN, NOTIFY comes before.
+	sql := `CREATE OR REPLACE FUNCTION example(x INT)
+RETURNS INT AS $$
+BEGIN
+    x := x + 1;
+    RETURN x;
+END;
+$$ LANGUAGE plpgsql;`
+
+	stmts := parser.ParseStatements(sql)
+	if len(stmts) == 0 {
+		t.Fatal("ParseStatements() returned no statements")
+	}
+
+	instrumentedSQL, coveragePoints := instrumentBody(stmts[0], "test.sql", true, "PERFORM")
+	if len(coveragePoints) != 2 {
+		t.Fatalf("expected 2 coverage points, got %d", len(coveragePoints))
+	}
+
+	// The assignment "x := x + 1" should have NOTIFY *after* it.
+	assignSignal := coveragePoints[0].SignalID
+	assignNotify := fmt.Sprintf("PERFORM pg_notify('pgcov', '%s');", assignSignal)
+	assignIdx := strings.Index(instrumentedSQL, assignNotify)
+	assignStmtIdx := strings.Index(instrumentedSQL, "x := x + 1")
+	if assignIdx < 0 || assignStmtIdx < 0 {
+		t.Fatal("could not find assignment or its notify in instrumented SQL")
+	}
+	if assignIdx <= assignStmtIdx {
+		t.Errorf("assignment: NOTIFY at %d should come after statement at %d (coverage-after-execute)", assignIdx, assignStmtIdx)
+	}
+
+	// The RETURN should have NOTIFY *before* it (terminal statement).
+	returnSignal := coveragePoints[1].SignalID
+	returnNotify := fmt.Sprintf("PERFORM pg_notify('pgcov', '%s');", returnSignal)
+	returnIdx := strings.Index(instrumentedSQL, returnNotify)
+	returnStmtIdx := strings.Index(instrumentedSQL, "RETURN x")
+	if returnIdx < 0 || returnStmtIdx < 0 {
+		t.Fatal("could not find RETURN or its notify in instrumented SQL")
+	}
+	if returnIdx >= returnStmtIdx {
+		t.Errorf("RETURN: NOTIFY at %d should come before statement at %d (terminal statement)", returnIdx, returnStmtIdx)
+	}
+
+	t.Log(instrumentedSQL)
+}
+
+func TestInstrumentBody_RaiseExceptionBeforeNotify(t *testing.T) {
+	// Use standalone RAISE statements so each is its own segment.
+	sql := `CREATE OR REPLACE FUNCTION check_positive(x INT)
+RETURNS VOID AS $$
+BEGIN
+    RAISE NOTICE 'checking value: %', x;
+    RAISE EXCEPTION 'negative value: %', x;
+END;
+$$ LANGUAGE plpgsql;`
+
+	stmts := parser.ParseStatements(sql)
+	if len(stmts) == 0 {
+		t.Fatal("ParseStatements() returned no statements")
+	}
+
+	instrumentedSQL, coveragePoints := instrumentBody(stmts[0], "test.sql", true, "PERFORM")
+	if len(coveragePoints) != 2 {
+		t.Fatalf("expected 2 coverage points, got %d", len(coveragePoints))
+	}
+
+	// RAISE NOTICE is non-terminal — NOTIFY should come after it.
+	raiseNoticeSignal := coveragePoints[0].SignalID
+	raiseNoticeNotify := fmt.Sprintf("PERFORM pg_notify('pgcov', '%s');", raiseNoticeSignal)
+	raiseNoticeIdx := strings.Index(instrumentedSQL, raiseNoticeNotify)
+	raiseNoticeStmtIdx := strings.Index(instrumentedSQL, "RAISE NOTICE")
+	if raiseNoticeIdx < 0 || raiseNoticeStmtIdx < 0 {
+		t.Fatal("could not find RAISE NOTICE or its notify")
+	}
+	if raiseNoticeIdx <= raiseNoticeStmtIdx {
+		t.Errorf("RAISE NOTICE: NOTIFY at %d should come after statement at %d", raiseNoticeIdx, raiseNoticeStmtIdx)
+	}
+
+	// RAISE EXCEPTION is terminal — NOTIFY should come before it.
+	raiseExcSignal := coveragePoints[1].SignalID
+	raiseExcNotify := fmt.Sprintf("PERFORM pg_notify('pgcov', '%s');", raiseExcSignal)
+	raiseExcIdx := strings.Index(instrumentedSQL, raiseExcNotify)
+	raiseExcStmtIdx := strings.Index(instrumentedSQL, "RAISE EXCEPTION")
+	if raiseExcIdx < 0 || raiseExcStmtIdx < 0 {
+		t.Fatal("could not find RAISE EXCEPTION or its notify")
+	}
+	if raiseExcIdx >= raiseExcStmtIdx {
+		t.Errorf("RAISE EXCEPTION: NOTIFY at %d should come before statement at %d", raiseExcIdx, raiseExcStmtIdx)
+	}
+
+	t.Log(instrumentedSQL)
+}