fix(instrument): use CTE-based pg_notify for SQL-language functions

SELECT pg_notify(...) injected as a standalone statement inside
SQL-language function bodies produces an extra result set that
conflicts with the function's declared return type, causing runtime
errors.

Replace standalone SELECT pg_notify(...) with a CTE prefix:

  WITH _pgcov_signal AS (SELECT pg_notify('pgcov', '<signal>')) <stmt>

This wraps each statement without adding extra result rows, preserving
the function's return type semantics.

Changes:
- instrumentBody() gains a useCTE parameter; when true, coverage
  signals are injected as CTE prefixes instead of separate statements
- instrumentStatement() passes useCTE=true for sql-language functions
- Add unit tests for CTE instrumentation (single and multi-statement)
- Add testdata/sqlfunc/ with SQL-language function sources and tests
- Add TestSQLFunctionInstrumentation e2e test verifying instrumented
  SQL functions deploy and return correct values against real Postgres

Author: pashagolub

internal/instrument/instrumenter.go

@@ -74,10 +74,10 @@ func instrumentStatement(stmt *parser.Statement, filePath string) (string, []Cov
 	case parser.StmtFunction, parser.StmtProcedure, parser.StmtDO:
 		switch stmt.Language {
 		case "plpgsql":
-			instrumented, locs := instrumentBody(stmt, filePath, true, "PERFORM")
+			instrumented, locs := instrumentBody(stmt, filePath, true, false)
 			return instrumented, locs
 		case "sql":
-			instrumented, locs := instrumentBody(stmt, filePath, false, "SELECT")
+			instrumented, locs := instrumentBody(stmt, filePath, false, true)
 			return instrumented, locs
 		default:
 			// Unknown language, mark as implicitly covered
@@ -101,8 +101,12 @@ func instrumentStatement(stmt *parser.Statement, filePath string) (string, []Cov
 //
 // For PL/pgSQL (skipToBegin=true), tokens before the first BEGIN are skipped.
 // For SQL functions (skipToBegin=false), instrumentation starts immediately.
-// notifyCmd is "PERFORM" for PL/pgSQL or "SELECT" for SQL functions.
-func instrumentBody(stmt *parser.Statement, filePath string, skipToBegin bool, notifyCmd string) (string, []CoveragePoint) {
+// When useCTE is true, coverage signals are injected as a CTE prefix
+// (WITH _pgcov_signal AS (SELECT pg_notify(...)) <original statement>)
+// instead of a standalone statement, avoiding extra result sets that break
+// SQL-language function return types.
+// When useCTE is false, signals use PERFORM pg_notify(...) (PL/pgSQL).
+func instrumentBody(stmt *parser.Statement, filePath string, skipToBegin bool, useCTE bool) (string, []CoveragePoint) {
 	bodyContent := stmt.Body
 	if bodyContent == "" {
 		return stmt.RawSQL, nil
@@ -132,7 +136,7 @@ func instrumentBody(stmt *parser.Statement, filePath string, skipToBegin bool, n
 	// 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).
+	// so that coverage is recorded only on successful execution.
 	emitSegment := func(segEnd int) {
 		segText := bodyContent[segStart:segEnd]
 		if !isExecutableSegment(segText) {
@@ -167,16 +171,25 @@ func instrumentBody(stmt *parser.Statement, filePath string, skipToBegin bool, n
 
 		escapedSignal := strings.ReplaceAll(cp.SignalID, "'", "''")
 
-		if termPos := findTerminalPos(segText); termPos >= 0 {
+		if useCTE {
+			// SQL-language functions: inject coverage signal as a CTE
+			// prefix so we don't produce an extra result set that would
+			// conflict with the function's declared return type (B6).
+			ctePrefix := fmt.Sprintf("WITH _pgcov_signal AS (SELECT pg_notify('pgcov', '%s')) ",
+				escapedSignal)
+			instrumentedBody.WriteString(ctePrefix)
+			instrumentedBody.WriteString(segText)
+			lastWrittenPos = segEnd
+		} else if termPos := findTerminalPos(segText); termPos >= 0 {
 			// Segment contains (or starts with) a terminal statement
 			// (RETURN, RAISE EXCEPTION).  Place the signal before the
 			// terminal so it fires before the scope exits.
 			// When termPos == 0 the segment starts with the terminal;
 			// when termPos > 0 it is nested inside a control structure
 			// (e.g. IF … THEN RETURN …).
 			termIndent := indentOf(segText[termPos:])
-			notifyCall := fmt.Sprintf("%s%s pg_notify('pgcov', '%s');",
-				termIndent, notifyCmd, escapedSignal)
+			notifyCall := fmt.Sprintf("%sPERFORM pg_notify('pgcov', '%s');",
+				termIndent, escapedSignal)
 			instrumentedBody.WriteString(segText[:termPos])
 			fmt.Fprintf(&instrumentedBody, "%s\n", notifyCall)
 			instrumentedBody.WriteString(segText[termPos:])
@@ -193,8 +206,8 @@ func instrumentBody(stmt *parser.Statement, filePath string, skipToBegin bool, n
 			} else {
 				lastWrittenPos = segEnd
 			}
-			notifyCall := fmt.Sprintf("%s%s pg_notify('pgcov', '%s');",
-				indent, notifyCmd, escapedSignal)
+			notifyCall := fmt.Sprintf("%sPERFORM pg_notify('pgcov', '%s');",
+				indent, escapedSignal)
 			fmt.Fprintf(&instrumentedBody, "\n%s", notifyCall)
 		}
 	}

internal/instrument/instrumenter_test.go

@@ -32,7 +32,7 @@ $$ LANGUAGE plpgsql;`
 	}
 	stmt := stmts[0]
 
-	instrumentedSQL, coveragePoints := instrumentBody(stmt, "test.sql", true, "PERFORM")
+	instrumentedSQL, coveragePoints := instrumentBody(stmt, "test.sql", true, false)
 	if instrumentedSQL == "" {
 		t.Error("instrumentWithLexer() returned empty instrumented SQL")
 	}
@@ -127,6 +127,99 @@ $$ LANGUAGE plpgsql;`
 	}
 }
 
+func TestInstrument_SQLFunction_UsesCTE(t *testing.T) {
+	sql := `CREATE OR REPLACE FUNCTION double_val(x INT)
+RETURNS INT AS $$
+    SELECT x * 2;
+$$ LANGUAGE sql;`
+
+	tmpDir := t.TempDir()
+	tmpFile := filepath.Join(tmpDir, "sqlfunc.sql")
+	if err := os.WriteFile(tmpFile, []byte(sql), 0644); err != nil {
+		t.Fatalf("failed to write temp file: %v", err)
+	}
+
+	file := &discovery.DiscoveredFile{
+		Path:         tmpFile,
+		RelativePath: "sqlfunc.sql",
+		Type:         discovery.FileTypeSource,
+	}
+
+	parsed, err := parser.Parse(file)
+	if err != nil {
+		t.Fatalf("Parse() error = %v", err)
+	}
+
+	instrumented, err := GenerateCoverageInstrument(parsed)
+	if err != nil {
+		t.Fatalf("Instrument() error = %v", err)
+	}
+
+	// Should have coverage points
+	if len(instrumented.Locations) == 0 {
+		t.Fatal("Instrument() produced no coverage points for SQL function")
+	}
+
+	// Should use CTE-based instrumentation, not standalone SELECT pg_notify
+	if !strings.Contains(instrumented.InstrumentedText, "WITH _pgcov_signal AS (SELECT pg_notify(") {
+		t.Error("SQL function should use CTE-based instrumentation")
+	}
+
+	// Should NOT have a standalone SELECT pg_notify (without CTE wrapper)
+	// Check that 'SELECT pg_notify' only appears inside CTE definitions
+	text := instrumented.InstrumentedText
+	cteRemoved := strings.ReplaceAll(text, "WITH _pgcov_signal AS (SELECT pg_notify(", "")
+	if strings.Contains(cteRemoved, "SELECT pg_notify(") {
+		t.Error("SQL function should not have standalone SELECT pg_notify calls")
+	}
+
+	// Should NOT use PERFORM (that's for PL/pgSQL)
+	if strings.Contains(instrumented.InstrumentedText, "PERFORM pg_notify") {
+		t.Error("SQL function should not use PERFORM")
+	}
+
+	t.Logf("Instrumented SQL:\n%s", instrumented.InstrumentedText)
+}
+
+func TestInstrument_SQLFunction_MultipleStatements(t *testing.T) {
+	// SQL functions can have multiple statements; the last one determines the return value
+	sql := `CREATE OR REPLACE FUNCTION insert_and_count()
+RETURNS BIGINT AS $$
+    INSERT INTO log(msg) VALUES ('hello');
+    SELECT count(*) FROM log;
+$$ LANGUAGE sql;`
+
+	tmpDir := t.TempDir()
+	tmpFile := filepath.Join(tmpDir, "sqlfunc_multi.sql")
+	if err := os.WriteFile(tmpFile, []byte(sql), 0644); err != nil {
+		t.Fatalf("failed to write temp file: %v", err)
+	}
+
+	file := &discovery.DiscoveredFile{
+		Path:         tmpFile,
+		RelativePath: "sqlfunc_multi.sql",
+		Type:         discovery.FileTypeSource,
+	}
+
+	parsed, err := parser.Parse(file)
+	if err != nil {
+		t.Fatalf("Parse() error = %v", err)
+	}
+
+	instrumented, err := GenerateCoverageInstrument(parsed)
+	if err != nil {
+		t.Fatalf("Instrument() error = %v", err)
+	}
+
+	// Both statements should get CTE-based instrumentation
+	cteCount := strings.Count(instrumented.InstrumentedText, "WITH _pgcov_signal AS (SELECT pg_notify(")
+	if cteCount < 2 {
+		t.Errorf("Expected at least 2 CTE injections for multi-statement SQL function, got %d", cteCount)
+	}
+
+	t.Logf("Instrumented SQL:\n%s", instrumented.InstrumentedText)
+}
+
 func TestInstrument_MultipleStatements(t *testing.T) {
 	sql := `SELECT 1;
 SELECT 2;