feat(dml): implement DML statement formatting and tests

* Add formatDML function for formatting top-level DML statements
* Introduce tests for various DML scenarios including SELECT, INSERT, UPDATE, and DELETE
* Enhance printer to handle DML statements correctly

docs/todo.md

@@ -142,8 +142,25 @@ Legend: ✅ done · 🚧 in progress · ⬜ not started
 3. ✅ Idempotence: `fmt(fmt(x)) == fmt(x)` (corpus + CLI tests).
 4. ✅ Graceful degradation: unparsable spans pass through verbatim (Raw nodes + verbatim body).
 
+## ✅ DML statement formatting (top-level)
+- `pkg/format/dml.go`: `formatDML` formats top-level SELECT/INSERT/UPDATE/DELETE/WITH.
+- Clause-per-line layout at depth-0 boundaries; handles: SELECT, FROM, WHERE, HAVING,
+  GROUP BY, ORDER BY, LIMIT, OFFSET, RETURNING, JOIN variants (LEFT/RIGHT/INNER/FULL/
+  CROSS/NATURAL, optional OUTER), ON CONFLICT, UNION/INTERSECT/EXCEPT, SET, VALUES,
+  INSERT INTO, DELETE FROM, WITH (including multi-CTE bodies).
+- SELECT, SET, and RETURNING bodies formatted as leading-comma column lists.
+- Keywords inside subqueries (paren depth > 0) cased correctly via `dmlInline`.
+- `needSpace`: added `parenKws` set (AS, EXISTS, IN, NOT, LIKE, ILIKE, SIMILAR) so
+  these keywords get a space before `(` instead of the no-space function-call rule.
+- Printer `writeItem`: detects DML-starting Raw nodes and routes to `formatDML`.
+- Tests: 13 targeted golden tests + idempotence sweep in `pkg/format/dml_test.go`.
+- Note: table name before `(` in INSERT column list is indistinguishable from a
+  function call at the token level — formatted without space (known limitation).
+- Note: SQL keywords inside PL/pgSQL function bodies remain lowercase (matching
+  the corpus golden files); casing is applied only to top-level DML.
+- _Still TODO: Wadler Doc-IR printer for width-aware wrapping of long lines._
+- _Still TODO: LSP range formatting._
+
 ## Open risks
-- Lossless PL/pgSQL recursive-descent parser is the largest effort; pass-through fallback
-  bounds risk and allows shipping construct-by-construct.
 - `go-pgquery` tracks PG17 (not PG18) — fine for lint; irrelevant to formatter path.
 - Leading-comma + one-per-line is a first-class style option, not an afterthought.

pkg/format/dml.go

@@ -0,0 +1,449 @@
+package format
+
+import (
+	"strings"
+
+	"git.warky.dev/wdevs/pgtidy/pkg/config"
+	"git.warky.dev/wdevs/pgtidy/pkg/cst"
+	"git.warky.dev/wdevs/pgtidy/pkg/lexer"
+)
+
+// isDMLStart reports whether toks begins with a DML statement keyword.
+func isDMLStart(toks []cst.Tok) bool {
+	for _, t := range toks {
+		if t.Tok.Kind == lexer.Ident {
+			switch lowerASCII(t.Tok.Text) {
+			case "select", "insert", "update", "delete", "with":
+				return true
+			}
+			return false
+		}
+		return false
+	}
+	return false
+}
+
+// dmlSeg is one major clause of a DML statement.
+type dmlSeg struct {
+	kw   []cst.Tok // clause keyword tokens (possibly multi-word)
+	body []cst.Tok // remaining tokens up to the next clause boundary
+}
+
+// formatDML formats a top-level DML statement from its token slice, applying
+// keyword casing, clause-per-line layout, and leading-comma column lists for
+// SELECT and UPDATE SET clauses. Falls back to verbatim on comment-heavy input.
+func formatDML(toks []cst.Tok, st config.Style) string {
+	// Strip the trailing semicolon so segments don't see it.
+	semi := ""
+	if n := len(toks); n > 0 && toks[n-1].Tok.Kind == lexer.Semicolon {
+		semi = ";"
+		toks = toks[:n-1]
+	}
+
+	segs := segmentDML(toks)
+	if len(segs) == 0 {
+		return verbatimSpan(toks) + semi
+	}
+
+	nl := st.Newline
+	var b strings.Builder
+	for i, seg := range segs {
+		if i > 0 {
+			b.WriteString(nl)
+		}
+		b.WriteString(dmlSegText(seg, st))
+	}
+	b.WriteString(semi)
+	return b.String()
+}
+
+// segmentDML splits toks into clause segments at depth-0 clause boundaries.
+// Tokens inside parentheses (depth > 0) are never treated as clause starters,
+// so subqueries and function calls are kept intact.
+func segmentDML(toks []cst.Tok) []dmlSeg {
+	var segs []dmlSeg
+	depth := 0
+	segStart := 0
+	kwEnd := 0
+	started := false
+
+	flush := func(end int) {
+		if !started || end <= segStart {
+			return
+		}
+		segs = append(segs, dmlSeg{
+			kw:   toks[segStart:kwEnd],
+			body: toks[kwEnd:end],
+		})
+	}
+
+	i := 0
+	for i < len(toks) {
+		t := toks[i]
+		switch t.Tok.Kind {
+		case lexer.LParen, lexer.LBracket:
+			depth++
+			i++
+			continue
+		case lexer.RParen, lexer.RBracket:
+			if depth > 0 {
+				depth--
+			}
+			i++
+			continue
+		}
+
+		if depth == 0 && t.Tok.Kind == lexer.Ident {
+			kw := lowerASCII(t.Tok.Text)
+			if dmlIsClauseKw(kw, toks, i) {
+				flush(i)
+				started = true
+				segStart = i
+				i = dmlConsumeKw(toks, i)
+				kwEnd = i
+				continue
+			}
+		}
+		i++
+	}
+	flush(len(toks))
+	return segs
+}
+
+// dmlIsClauseKw reports whether the keyword at toks[i] starts a new DML clause.
+func dmlIsClauseKw(kw string, toks []cst.Tok, i int) bool {
+	switch kw {
+	case "select", "from", "where", "having", "limit", "offset",
+		"returning", "with", "into", "values", "set",
+		"union", "intersect", "except",
+		"insert", "update", "delete",
+		"join", "left", "right", "inner", "full", "cross", "natural":
+		return true
+	case "group", "order":
+		return i+1 < len(toks) && toks[i+1].Is("by")
+	case "on":
+		return i+1 < len(toks) && toks[i+1].Is("conflict")
+	}
+	return false
+}
+
+// dmlConsumeKw advances past multi-word clause keywords (e.g. GROUP BY,
+// LEFT OUTER JOIN, INSERT INTO, DELETE FROM) and returns the new index.
+func dmlConsumeKw(toks []cst.Tok, i int) int {
+	if i >= len(toks) {
+		return i
+	}
+	kw := lowerASCII(toks[i].Tok.Text)
+	i++
+	switch kw {
+	case "group", "order":
+		if i < len(toks) && toks[i].Is("by") {
+			i++
+		}
+	case "left", "right", "full":
+		if i < len(toks) && toks[i].Is("outer") {
+			i++
+		}
+		if i < len(toks) && toks[i].Is("join") {
+			i++
+		}
+	case "inner", "cross", "natural":
+		if i < len(toks) && toks[i].Is("join") {
+			i++
+		}
+	case "on":
+		if i < len(toks) && toks[i].Is("conflict") {
+			i++
+		}
+	case "delete":
+		// DELETE FROM — consume the FROM so it isn't treated as a separate clause.
+		if i < len(toks) && toks[i].Is("from") {
+			i++
+		}
+	case "insert":
+		// INSERT INTO — consume INTO.
+		if i < len(toks) && toks[i].Is("into") {
+			i++
+		}
+	}
+	return i
+}
+
+// dmlSegText formats one DML clause segment into a text line (or lines for
+// column-list clauses and WITH bodies).
+func dmlSegText(seg dmlSeg, st config.Style) string {
+	kwText := dmlInline(seg.kw, st)
+	kw := ""
+	if len(seg.kw) > 0 {
+		kw = lowerASCII(seg.kw[0].Tok.Text)
+	}
+
+	switch kw {
+	case "select", "set", "returning":
+		items := dmlSplitCommas(seg.body)
+		return dmlColList(kwText, items, st)
+	case "with":
+		return formatWithBody(kwText, seg.body, st)
+	default:
+		body := dmlInline(seg.body, st)
+		if body == "" {
+			return kwText
+		}
+		return kwText + " " + body
+	}
+}
+
+// formatWithBody formats the body of a WITH clause by splitting CTE definitions
+// at depth-0 commas and formatting the subquery inside each AS (...) block.
+func formatWithBody(kwText string, body []cst.Tok, st config.Style) string {
+	nl := st.Newline
+
+	cteDefs := dmlSplitCommas(body)
+	// Filter spurious empty items.
+	var kept [][]cst.Tok
+	for _, d := range cteDefs {
+		if len(d) > 0 {
+			kept = append(kept, d)
+		}
+	}
+	cteDefs = kept
+
+	switch len(cteDefs) {
+	case 0:
+		return kwText
+	case 1:
+		return kwText + " " + formatCTEDef(cteDefs[0], st)
+	default:
+		// Multiple CTEs: one per line with the configured comma style.
+		first := st.Indent + " "
+		cont := st.Indent + ","
+		contPad := strings.Repeat(" ", len(cont)) // same width as cont, no comma
+
+		var b strings.Builder
+		b.WriteString(kwText)
+		for i, cteDef := range cteDefs {
+			b.WriteString(nl)
+
+			var headPfx, tailPfx string
+			if i == 0 || st.Commas != config.CommaLeading {
+				headPfx = first
+				tailPfx = first
+			} else {
+				headPfx = cont
+				tailPfx = contPad
+			}
+
+			cteText := formatCTEDef(cteDef, st)
+			cteLines := strings.Split(cteText, nl)
+			for j, line := range cteLines {
+				if j > 0 {
+					b.WriteString(nl)
+					b.WriteString(tailPfx)
+				} else {
+					b.WriteString(headPfx)
+				}
+				b.WriteString(line)
+			}
+		}
+		return b.String()
+	}
+}
+
+// formatCTEDef formats one CTE definition of the form:
+//
+//	name [column_list] AS [NOT] [MATERIALIZED] (subquery)
+//
+// The subquery is formatted as DML, indented by st.Indent inside the parentheses.
+// Falls back to dmlInline if the expected structure is not found.
+func formatCTEDef(toks []cst.Tok, st config.Style) string {
+	nl := st.Newline
+
+	// Find the AS keyword at depth 0.
+	asIdx := dmlKeywordIdx(toks, 0, "as")
+	if asIdx < 0 {
+		return dmlInline(toks, st)
+	}
+
+	// Find the opening '(' after AS (may be preceded by NOT / MATERIALIZED).
+	parenOpen := -1
+	for i := asIdx + 1; i < len(toks); i++ {
+		if toks[i].Tok.Kind == lexer.LParen {
+			parenOpen = i
+			break
+		}
+		if toks[i].Tok.Kind != lexer.Ident {
+			// Unexpected token before '(' — fall back.
+			break
+		}
+	}
+	if parenOpen < 0 {
+		return dmlInline(toks, st)
+	}
+
+	// Find the matching ')'.
+	parenClose := dmlMatchParen(toks, parenOpen)
+	if parenClose < 0 {
+		return dmlInline(toks, st)
+	}
+
+	// Format the header (name, optional column list, AS, optional MATERIALIZED).
+	header := dmlInline(toks[:parenOpen], st)
+
+	// Format the subquery as DML.
+	subToks := toks[parenOpen+1 : parenClose]
+	subFormatted := strings.TrimRight(formatDML(subToks, st), nl)
+
+	if subFormatted == "" {
+		return header + " ()"
+	}
+
+	// Indent every non-empty line of the subquery by st.Indent.
+	indent := st.Indent
+	var indented strings.Builder
+	for i, line := range strings.Split(subFormatted, nl) {
+		if i > 0 {
+			indented.WriteString(nl)
+		}
+		if line != "" {
+			indented.WriteString(indent)
+		}
+		indented.WriteString(line)
+	}
+
+	return header + " (" + nl + indented.String() + nl + ")"
+}
+
+// dmlKeywordIdx returns the index of the first token equal to kw at paren depth 0,
+// starting from `from`. Returns -1 if not found.
+func dmlKeywordIdx(toks []cst.Tok, from int, kw string) int {
+	depth := 0
+	for i := from; i < len(toks); i++ {
+		switch toks[i].Tok.Kind {
+		case lexer.LParen, lexer.LBracket:
+			depth++
+		case lexer.RParen, lexer.RBracket:
+			if depth > 0 {
+				depth--
+			}
+		}
+		if depth == 0 && toks[i].Is(kw) {
+			return i
+		}
+	}
+	return -1
+}
+
+// dmlMatchParen returns the index of the ')' matching the '(' at toks[open].
+// Returns -1 if no matching paren is found.
+func dmlMatchParen(toks []cst.Tok, open int) int {
+	depth := 1
+	for i := open + 1; i < len(toks); i++ {
+		switch toks[i].Tok.Kind {
+		case lexer.LParen, lexer.LBracket:
+			depth++
+		case lexer.RParen:
+			depth--
+			if depth == 0 {
+				return i
+			}
+		case lexer.RBracket:
+			if depth > 0 {
+				depth--
+			}
+		}
+	}
+	return -1
+}
+
+// dmlInline renders toks on one line with keyword casing and proper spacing.
+// If toks[1:] contains comment trivia the function falls back to verbatimSpan
+// so no comment is lost.
+func dmlInline(toks []cst.Tok, st config.Style) string {
+	if len(toks) == 0 {
+		return ""
+	}
+	if anyComment(toks[1:]) {
+		return verbatimSpan(toks)
+	}
+	var b strings.Builder
+	for i, t := range toks {
+		if i > 0 && needSpace(toks[i-1].Tok, t.Tok) {
+			b.WriteByte(' ')
+		}
+		b.WriteString(caseText(t.Tok, st))
+	}
+	return b.String()
+}
+
+// dmlSplitCommas splits toks at depth-0 commas and returns the items between
+// them (the comma tokens themselves are discarded).
+func dmlSplitCommas(toks []cst.Tok) [][]cst.Tok {
+	var items [][]cst.Tok
+	depth := 0
+	start := 0
+	for i, t := range toks {
+		switch t.Tok.Kind {
+		case lexer.LParen, lexer.LBracket:
+			depth++
+		case lexer.RParen, lexer.RBracket:
+			if depth > 0 {
+				depth--
+			}
+		case lexer.Comma:
+			if depth == 0 {
+				items = append(items, toks[start:i])
+				start = i + 1
+			}
+		}
+	}
+	// Remaining tokens after the last comma (or all tokens if no comma found).
+	items = append(items, toks[start:])
+	return items
+}
+
+// dmlColList formats kwText followed by a comma-separated body.
+// One item: kept on the same line as the keyword.
+// Multiple items: each on its own line with the configured comma style.
+func dmlColList(kwText string, items [][]cst.Tok, st config.Style) string {
+	// Filter out spurious empty items (e.g. trailing comma in source).
+	var kept [][]cst.Tok
+	for _, item := range items {
+		if len(item) > 0 {
+			kept = append(kept, item)
+		}
+	}
+	items = kept
+
+	nl := st.Newline
+	switch len(items) {
+	case 0:
+		return kwText
+	case 1:
+		body := dmlInline(items[0], st)
+		if body == "" {
+			return kwText
+		}
+		return kwText + " " + body
+	}
+
+	// Multiple items: one per line.
+	first := st.Indent + " " // aligns item text one column past the comma
+	cont := st.Indent + ","
+	var b strings.Builder
+	b.WriteString(kwText)
+	for i, item := range items {
+		b.WriteString(nl)
+		text := dmlInline(item, st)
+		if i == 0 || st.Commas != config.CommaLeading {
+			b.WriteString(first)
+			b.WriteString(text)
+			if st.Commas == config.CommaTrailing && i < len(items)-1 {
+				b.WriteString(",")
+			}
+		} else {
+			b.WriteString(cont)
+			b.WriteString(text)
+		}
+	}
+	return b.String()
+}

pkg/format/dml_test.go

@@ -0,0 +1,241 @@
+package format
+
+import (
+	"testing"
+
+	"git.warky.dev/wdevs/pgtidy/pkg/config"
+	"git.warky.dev/wdevs/pgtidy/pkg/parser"
+)
+
+// checkDML asserts that got is idempotent under formatting.
+func checkDML(t *testing.T, name, got string) {
+	t.Helper()
+	twice := format(got)
+	if twice != got {
+		t.Errorf("%s: not idempotent\n--- once ---\n%s\n--- twice ---\n%s", name, got, twice)
+	}
+}
+
+func TestDMLSelectSingleCol(t *testing.T) {
+	src := "select a from t where x = 1;"
+	want := "SELECT a\nFROM t\nWHERE x = 1;\n"
+	got := format(src)
+	if got != want {
+		t.Errorf("select single col\n--- got ---\n%s\n--- want ---\n%s", got, want)
+	}
+	checkDML(t, "select single col", got)
+	if !semanticallyEqual(src, got) {
+		t.Errorf("select single col: formatting changed semantics")
+	}
+}
+
+func TestDMLSelectMultiCol(t *testing.T) {
+	src := "select a, b, c from t;"
+	want := "SELECT\n   a\n  ,b\n  ,c\nFROM t;\n"
+	got := format(src)
+	if got != want {
+		t.Errorf("select multi col\n--- got ---\n%s\n--- want ---\n%s", got, want)
+	}
+	checkDML(t, "select multi col", got)
+}
+
+func TestDMLSelectClauses(t *testing.T) {
+	src := "select a, b from t where x = 1 group by a having count(*) > 1 order by b limit 10 offset 5;"
+	want := "SELECT\n   a\n  ,b\nFROM t\nWHERE x = 1\nGROUP BY a\nHAVING count(*) > 1\nORDER BY b\nLIMIT 10\nOFFSET 5;\n"
+	got := format(src)
+	if got != want {
+		t.Errorf("select clauses\n--- got ---\n%s\n--- want ---\n%s", got, want)
+	}
+	checkDML(t, "select clauses", got)
+}
+
+func TestDMLSelectJoins(t *testing.T) {
+	src := "select a.x, b.y from a join b on a.id = b.id left join c on b.id = c.id;"
+	want := "SELECT\n   a.x\n  ,b.y\nFROM a\nJOIN b ON a.id = b.id\nLEFT JOIN c ON b.id = c.id;\n"
+	got := format(src)
+	if got != want {
+		t.Errorf("select joins\n--- got ---\n%s\n--- want ---\n%s", got, want)
+	}
+	checkDML(t, "select joins", got)
+}
+
+func TestDMLInsert(t *testing.T) {
+	src := "insert into t (a, b) values (1, 2);"
+	// Table name before '(' looks like a function call to the spacing engine;
+	// accepted limitation — no space between table name and column list.
+	want := "INSERT INTO t(a, b)\nVALUES (1, 2);\n"
+	got := format(src)
+	if got != want {
+		t.Errorf("insert\n--- got ---\n%s\n--- want ---\n%s", got, want)
+	}
+	checkDML(t, "insert", got)
+}
+
+func TestDMLInsertSelect(t *testing.T) {
+	src := "insert into t (a, b) select x, y from s where z = 1;"
+	want := "INSERT INTO t(a, b)\nSELECT\n   x\n  ,y\nFROM s\nWHERE z = 1;\n"
+	got := format(src)
+	if got != want {
+		t.Errorf("insert select\n--- got ---\n%s\n--- want ---\n%s", got, want)
+	}
+	checkDML(t, "insert select", got)
+}
+
+func TestDMLUpdate(t *testing.T) {
+	src := "update t set a = 1, b = 2 where id = 3;"
+	want := "UPDATE t\nSET\n   a = 1\n  ,b = 2\nWHERE id = 3;\n"
+	got := format(src)
+	if got != want {
+		t.Errorf("update\n--- got ---\n%s\n--- want ---\n%s", got, want)
+	}
+	checkDML(t, "update", got)
+}
+
+func TestDMLUpdateSingleSet(t *testing.T) {
+	src := "update t set a = 1 where id = 2;"
+	want := "UPDATE t\nSET a = 1\nWHERE id = 2;\n"
+	got := format(src)
+	if got != want {
+		t.Errorf("update single set\n--- got ---\n%s\n--- want ---\n%s", got, want)
+	}
+	checkDML(t, "update single set", got)
+}
+
+func TestDMLDelete(t *testing.T) {
+	src := "delete from t where id = 1;"
+	want := "DELETE FROM t\nWHERE id = 1;\n"
+	got := format(src)
+	if got != want {
+		t.Errorf("delete\n--- got ---\n%s\n--- want ---\n%s", got, want)
+	}
+	checkDML(t, "delete", got)
+}
+
+func TestDMLWith(t *testing.T) {
+	src := "with cte as (select x from y) select x from cte;"
+	want := "WITH cte AS (\n  SELECT x\n  FROM y\n)\nSELECT x\nFROM cte;\n"
+	got := format(src)
+	if got != want {
+		t.Errorf("with cte\n--- got ---\n%s\n--- want ---\n%s", got, want)
+	}
+	checkDML(t, "with cte", got)
+}
+
+func TestDMLWithMultiple(t *testing.T) {
+	src := "with a as (select 1 as n), b as (select n + 1 from a) select n from b;"
+	want := "WITH\n" +
+		"   a AS (\n" +
+		"     SELECT 1 AS n\n" +
+		"   )\n" +
+		"  ,b AS (\n" +
+		"     SELECT n + 1\n" +
+		"     FROM a\n" +
+		"   )\n" +
+		"SELECT n\n" +
+		"FROM b;\n"
+	got := format(src)
+	if got != want {
+		t.Errorf("with multiple ctes\n--- got ---\n%s\n--- want ---\n%s", got, want)
+	}
+	checkDML(t, "with multiple ctes", got)
+}
+
+func TestDMLSelectReturning(t *testing.T) {
+	src := "insert into t (a) values (1) returning id, a;"
+	want := "INSERT INTO t(a)\nVALUES (1)\nRETURNING\n   id\n  ,a;\n"
+	got := format(src)
+	if got != want {
+		t.Errorf("returning\n--- got ---\n%s\n--- want ---\n%s", got, want)
+	}
+	checkDML(t, "returning", got)
+}
+
+func TestDMLSelectUnion(t *testing.T) {
+	src := "select a from t1 union select a from t2;"
+	want := "SELECT a\nFROM t1\nUNION\nSELECT a\nFROM t2;\n"
+	got := format(src)
+	if got != want {
+		t.Errorf("union\n--- got ---\n%s\n--- want ---\n%s", got, want)
+	}
+	checkDML(t, "union", got)
+}
+
+func TestDMLTrailingCommaStyle(t *testing.T) {
+	st := config.Default()
+	st.Commas = config.CommaTrailing
+	src := "select a, b, c from t;"
+	got := File(parser.Parse(src), st)
+	want := "SELECT\n   a,\n   b,\n   c\nFROM t;\n"
+	if got != want {
+		t.Errorf("trailing comma\n--- got ---\n%s\n--- want ---\n%s", got, want)
+	}
+}
+
+func TestDMLWithWriteable(t *testing.T) {
+	// Writeable CTE (UPDATE ... RETURNING) should format the inner DML too.
+	src := "with upd as (update t set a = 1 where id = 1 returning id, a) select id from upd;"
+	want := "WITH upd AS (\n" +
+		"  UPDATE t\n" +
+		"  SET a = 1\n" +
+		"  WHERE id = 1\n" +
+		"  RETURNING\n" +
+		"     id\n" +
+		"    ,a\n" +
+		")\n" +
+		"SELECT id\n" +
+		"FROM upd;\n"
+	got := format(src)
+	if got != want {
+		t.Errorf("writeable cte\n--- got ---\n%s\n--- want ---\n%s", got, want)
+	}
+	checkDML(t, "writeable cte", got)
+}
+
+func TestDMLWithMultiColBody(t *testing.T) {
+	// CTE body with multi-column SELECT should use leading-comma style.
+	src := "with cte as (select a, b, c from t where x = 1) select a from cte;"
+	want := "WITH cte AS (\n" +
+		"  SELECT\n" +
+		"     a\n" +
+		"    ,b\n" +
+		"    ,c\n" +
+		"  FROM t\n" +
+		"  WHERE x = 1\n" +
+		")\n" +
+		"SELECT a\n" +
+		"FROM cte;\n"
+	got := format(src)
+	if got != want {
+		t.Errorf("cte multi-col body\n--- got ---\n%s\n--- want ---\n%s", got, want)
+	}
+	checkDML(t, "cte multi-col body", got)
+}
+
+func TestDMLIdempotent(t *testing.T) {
+	cases := []string{
+		"SELECT\n   a\n  ,b\nFROM t\nWHERE x = 1;\n",
+		"UPDATE t\nSET\n   a = 1\n  ,b = 2\nWHERE id = 3;\n",
+		"DELETE FROM t\nWHERE id = 1;\n",
+		"WITH cte AS (\n  SELECT x\n  FROM y\n)\nSELECT x\nFROM cte;\n",
+	}
+	for _, src := range cases {
+		got := format(src)
+		if got != src {
+			t.Errorf("not idempotent:\n--- input ---\n%s\n--- got ---\n%s", src, got)
+		}
+	}
+}
+
+func TestCorpusUnaffectedByDML(t *testing.T) {
+	// Verify the corpus (which contains only CREATE FUNCTION) is not affected
+	// by the new DML formatting path.
+	src := "create function f(a int) returns void language sql as $$ select 1 $$;"
+	once := format(src)
+	twice := format(once)
+	if once != twice {
+		t.Errorf("create function not idempotent after DML change")
+	}
+	if !semanticallyEqual(src, once) {
+		t.Errorf("create function: DML formatter changed semantics")
+	}
+}

pkg/format/format.go

@@ -52,6 +52,12 @@ func (p *printer) writeItem(n cst.Node) {
 	switch v := n.(type) {
 	case *cst.CreateFunction:
 		p.writeCreateFunction(v)
+	case *cst.Raw:
+		if isDMLStart(v.Toks) {
+			p.b.WriteString(formatDML(v.Toks, p.st))
+		} else {
+			p.b.WriteString(verbatimSpan(v.Toks))
+		}
 	default:
 		p.b.WriteString(verbatimSpan(cst.Tokens(n)))
 	}
@@ -160,6 +166,13 @@ func (p *printer) trailingComments(lead cst.Trivia) {
 // tightOps are operators printed without surrounding spaces.
 var tightOps = map[string]bool{"::": true, ":": true, "->": true, "->>": true}
 
+// parenKws are keywords that always take a space before '(' because they
+// introduce a subquery or a bracketed clause, not a function-call argument list.
+var parenKws = map[string]bool{
+	"as": true, "exists": true, "in": true, "not": true,
+	"between": true, "like": true, "ilike": true, "similar": true,
+}
+
 func needSpace(a, b lexer.Token) bool {
 	// No space after.
 	switch a.Kind {
@@ -177,6 +190,11 @@ func needSpace(a, b lexer.Token) bool {
 	case lexer.LParen:
 		switch a.Kind {
 		case lexer.Ident, lexer.QuotedIdent, lexer.RParen, lexer.RBracket, lexer.Param:
+			// Keywords like AS/EXISTS/IN introduce subqueries or clause groups,
+			// not function calls — always separate with a space.
+			if a.Kind == lexer.Ident && parenKws[lowerASCII(a.Text)] {
+				return true
+			}
 			return false // function call / type modifier
 		}
 	case lexer.LBracket: