feat(lint): lint code-like dollar-quoted SQL strings

pkg/lint/lint.go

@@ -94,6 +94,7 @@ func (e *Engine) Check(sql, file string) ([]diagnostics.Diagnostic, error) {
 	}
 
 	bodies := e.checkBodies(sql, result.Stmts)
+	bodies = append(bodies, e.checkEmbeddedLiterals(sql, result.Stmts)...)
 	for i := range bodies {
 		bodies[i].File = file
 	}

pkg/lint/lint_test.go

@@ -154,3 +154,23 @@ func TestPlpgsqlBodyLint(t *testing.T) {
 		t.Errorf("PLPGSQL syntax error: %+v", pl)
 	}
 }
+
+func TestEmbeddedDollarSQLLint(t *testing.T) {
+	src := "create function f() returns void language plpgsql as $$\nbegin\n  execute $q$\n    select * from t\n  $q$;\nend;\n$$;\n" +
+		"create function s() returns int language sql as $$ select * from u $$;\n" +
+		"create function p() returns void language plpython3u as $$\nselect * from nothing\n$$;\n" +
+		"select format($f$select %I from t$f$, 'a');\n"
+	diags, err := lint.New().Check(src, "x.sql")
+	if err != nil {
+		t.Fatal(err)
+	}
+	var got [][2]int
+	for _, d := range diags {
+		if d.RuleID == "COR001" {
+			got = append(got, [2]int{d.Line, d.Col})
+		}
+	}
+	if len(got) != 2 || got[0] != [2]int{4, 12} || got[1] != [2]int{8, 59} {
+		t.Errorf("COR001 positions: %v", got)
+	}
+}

pkg/lint/plpgsql.go

@@ -9,6 +9,7 @@ import (
 	waspg "github.com/wasilibs/go-pgquery"
 
 	"git.warky.dev/wdevs/pgtidy/pkg/diagnostics"
+	"git.warky.dev/wdevs/pgtidy/pkg/lexer"
 	"git.warky.dev/wdevs/pgtidy/pkg/pgast"
 )
 
@@ -200,3 +201,132 @@ func (e *Engine) checkEmbedded(q bodyQuery, routineStartLine int, fileLines []st
 	}
 	return out
 }
+
+// checkEmbeddedLiterals lints SQL held in dollar-quoted strings: anything whose
+// content starts with SELECT/INSERT/UPDATE/DELETE/WITH (the same "looks like
+// code" test the formatter uses), e.g. EXECUTE $q$ select … $q$ or a
+// LANGUAGE sql function body. Literals are searched recursively, including
+// inside PL/pgSQL bodies. Bodies of non-sql, non-plpgsql routines are opaque
+// and skipped, as are format()-style templates.
+func (e *Engine) checkEmbeddedLiterals(sql string, stmts []*pg_query.RawStmt) []diagnostics.Diagnostic {
+	if !strings.Contains(sql, "$") {
+		return nil
+	}
+	var opaque [][2]int
+	for _, raw := range stmts {
+		if raw.Stmt == nil {
+			continue
+		}
+		lang, has := "", false
+		switch n := raw.Stmt.GetNode().(type) {
+		case *pg_query.Node_CreateFunctionStmt:
+			for _, o := range n.CreateFunctionStmt.Options {
+				if d := o.GetDefElem(); d != nil && d.Defname == "language" {
+					lang, has = strings.ToLower(d.Arg.GetString_().GetSval()), true
+				}
+			}
+		case *pg_query.Node_DoStmt:
+			for _, a := range n.DoStmt.Args {
+				if d := a.GetDefElem(); d != nil && d.Defname == "language" {
+					lang, has = strings.ToLower(d.Arg.GetString_().GetSval()), true
+				}
+			}
+		default:
+			continue
+		}
+		if has && lang != "sql" && lang != "plpgsql" {
+			s := int(raw.StmtLocation)
+			end := len(sql)
+			if raw.StmtLen > 0 {
+				end = s + int(raw.StmtLen)
+			}
+			opaque = append(opaque, [2]int{s, end})
+		}
+	}
+	var out []diagnostics.Diagnostic
+	var walk func(text string, base int)
+	walk = func(text string, base int) {
+		for _, t := range lexer.Lex(text) {
+			if t.Kind != lexer.DollarString {
+				continue
+			}
+			abs := base + t.Off
+			skip := false
+			for _, r := range opaque {
+				if abs >= r[0] && abs < r[1] {
+					skip = true
+				}
+			}
+			if skip {
+				continue
+			}
+			inner, innerOff, ok := dollarInner(t.Text)
+			if !ok {
+				continue
+			}
+			if looksLikeSQL(inner) && !formatPlaceholderRe.MatchString(inner) {
+				out = append(out, e.checkLiteral(sql, inner, abs+innerOff)...)
+			}
+			walk(inner, abs+innerOff)
+		}
+	}
+	walk(sql, 0)
+	return out
+}
+
+var formatPlaceholderRe = regexp.MustCompile(`%(\d+\$)?-?\d*[sILlx]`)
+
+// dollarInner splits a dollar-quoted literal into its content and the content's
+// offset within the literal.
+func dollarInner(lit string) (inner string, off int, ok bool) {
+	if len(lit) < 4 || lit[0] != '$' {
+		return "", 0, false
+	}
+	e := strings.IndexByte(lit[1:], '$')
+	if e < 0 {
+		return "", 0, false
+	}
+	tag := lit[:e+2]
+	if !strings.HasSuffix(lit, tag) || len(lit) < 2*len(tag) {
+		return "", 0, false
+	}
+	return lit[len(tag) : len(lit)-len(tag)], len(tag), true
+}
+
+func looksLikeSQL(inner string) bool {
+	for _, t := range lexer.Lex(inner) {
+		if t.IsTrivia() || t.Kind == lexer.EOF {
+			continue
+		}
+		if t.Kind != lexer.Ident {
+			return false
+		}
+		switch strings.ToLower(t.Text) {
+		case "select", "insert", "update", "delete", "with":
+			return true
+		}
+		return false
+	}
+	return false
+}
+
+// checkLiteral lints one SQL string whose first byte sits at file offset off.
+func (e *Engine) checkLiteral(sql, inner string, off int) []diagnostics.Diagnostic {
+	res, err := pgast.Parse(inner)
+	if err != nil {
+		return nil
+	}
+	startLine, startCol := pgast.LocationToLineCol(sql, off)
+	var out []diagnostics.Diagnostic
+	for _, r := range e.rules {
+		for _, d := range r.Check(res.Stmts, inner) {
+			d.Fix = nil
+			if d.Line <= 1 {
+				d.Col = startCol + d.Col - 1
+			}
+			d.Line = startLine + d.Line - 1
+			out = append(out, d)
+		}
+	}
+	return out
+}