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
This commit is contained in:
@@ -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()
|
||||
}
|
||||
Reference in New Issue
Block a user