fix(cli): respect atomic END boundaries (CLI-2405)

apps/cli-go/pkg/parser/state.go

@@ -26,7 +26,12 @@ type ReadyState struct{}
 func (s *ReadyState) Next(r rune, data []byte) State {
 	switch r {
 	case '$':
+		// A $ after an identifier rune continues the identifier (pending$$foo$), not a
+		// dollar quote. A digit counts too (1$$), unlike PostgreSQL; valid SQL never has that.
 		offset := len(data) - utf8.RuneLen(r)
+		if hasIdentifierRuneBefore(data, offset) {
+			return s
+		}
 		return &TagState{offset: offset}
 	case '\'':
 		fallthrough
@@ -42,42 +47,85 @@ func (s *ReadyState) Next(r rune, data []byte) State {
 		// Emit token
 		return nil
 	case '(':
-		return &AtomicState{prev: s, delimiter: []byte{')'}}
+		return &ParenState{prev: s}
 	case 'c':
 		fallthrough
 	case 'C':
 		if isBeginAtomic(data) {
-			return &AtomicState{prev: s, delimiter: []byte(END_ATOMIC)}
+			return &AtomicState{prev: s, statementStart: len(data)}
 		}
 	}
 	return s
 }
 
 func isBeginAtomic(data []byte) bool {
-	offset := len(data) - len(BEGIN_ATOMIC)
-	if offset < 0 || !strings.EqualFold(string(data[offset:]), BEGIN_ATOMIC) {
+	if !endsWithKeyword(data, BEGIN_ATOMIC) {
 		return false
 	}
-	if offset > 0 {
-		r, _ := utf8.DecodeLastRune(data[:offset])
-		if isIdentifierRune(r) {
-			return false
-		}
-	}
-	prefix := bytes.TrimRightFunc(data[:offset], unicode.IsSpace)
-	offset = len(prefix) - len("BEGIN")
-	if offset < 0 || !strings.EqualFold(string(prefix[offset:]), "BEGIN") {
+	prefix := bytes.TrimRight(data[:len(data)-len(BEGIN_ATOMIC)], sqlWhitespace)
+	return endsWithKeyword(prefix, "BEGIN")
+}
+
+// PostgreSQL's scan.l treats every byte at or above 0x80 as an identifier/dollar-tag
+// character (ident_cont/dolq_cont), whatever its Unicode category.
+func isIdentifierRune(r rune) bool {
+	return r >= utf8.RuneSelf || unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_' || r == '$'
+}
+
+func hasIdentifierRuneBefore(data []byte, offset int) bool {
+	if offset <= 0 {
 		return false
 	}
-	if offset == 0 {
-		return true
+	r, _ := utf8.DecodeLastRune(data[:offset])
+	return isIdentifierRune(r)
+}
+
+func endsWithKeyword(data []byte, keyword string) bool {
+	offset := len(data) - len(keyword)
+	if offset < 0 || !strings.EqualFold(string(data[offset:]), keyword) {
+		return false
 	}
-	r, _ := utf8.DecodeLastRune(prefix[:offset])
-	return !isIdentifierRune(r)
+	return !hasIdentifierRuneBefore(data, offset)
 }
 
-func isIdentifierRune(r rune) bool {
-	return unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_' || r == '$'
+const sqlWhitespace = " \t\n\r\f\v"
+
+func isSqlWhitespace(b byte) bool {
+	return bytes.IndexByte([]byte(sqlWhitespace), b) >= 0
+}
+
+func isCommentsAndWhitespace(text []byte) bool {
+	for i := 0; i < len(text); {
+		switch {
+		case isSqlWhitespace(text[i]):
+			i++
+		case bytes.HasPrefix(text[i:], []byte("--")):
+			newline := bytes.IndexByte(text[i+2:], '\n')
+			if newline == -1 {
+				return true
+			}
+			i += 2 + newline + 1
+		case bytes.HasPrefix(text[i:], []byte("/*")):
+			// Match BlockState's sliding-window scan so both agree on overlapping delimiters.
+			depth := 1
+			i += 2
+			for i < len(text) && depth > 0 {
+				switch {
+				case bytes.HasPrefix(text[i-1:], []byte("/*")):
+					depth++
+				case bytes.HasPrefix(text[i-1:], []byte("*/")):
+					depth--
+				}
+				i++
+			}
+			if depth > 0 {
+				return true
+			}
+		default:
+			return false
+		}
+	}
+	return true
 }
 
 // Opened a line comment
@@ -173,7 +221,7 @@ func (s *TagState) Next(r rune, data []byte) State {
 		return &dollar
 	}
 	// Valid tag: https://www.postgresql.org/docs/current/sql-syntax-lexical.html
-	if unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_' {
+	if isIdentifierRune(r) {
 		return s
 	}
 	// Break out of tag state
@@ -188,22 +236,60 @@ func (s *EscapeState) Next(r rune, data []byte) State {
 	return &ReadyState{}
 }
 
+// Opened a parenthesis group
+type ParenState struct {
+	prev State
+}
+
+func (s *ParenState) Next(r rune, data []byte) State {
+	curr := s.prev.Next(r, data)
+	if curr == nil {
+		s.prev = &ReadyState{}
+		return s
+	}
+	s.prev = curr
+	if _, ok := s.prev.(*ReadyState); !ok {
+		return s
+	}
+	if r == ')' {
+		return &ReadyState{}
+	}
+	return s
+}
+
 // Opened BEGIN ATOMIC function body
 type AtomicState struct {
-	prev      State
-	delimiter []byte
+	prev                State
+	pendingEnd          bool
+	statementStart      int
+	statementHasContent bool
 }
 
 func (s *AtomicState) Next(r rune, data []byte) State {
-	// If we are in a quoted state, the current delimiter doesn't count.
-	if curr := s.prev.Next(r, data); curr != nil {
-		s.prev = curr
-	}
-	if _, ok := s.prev.(*ReadyState); ok {
-		window := data[len(data)-len(s.delimiter):]
-		// Treat delimiter as case insensitive
-		if strings.EqualFold(string(window), string(s.delimiter)) {
-			return &ReadyState{}
+	pendingEnd := s.pendingEnd
+	s.pendingEnd = false
+	if pendingEnd && !isIdentifierRune(r) {
+		return (&ReadyState{}).Next(r, data)
+	}
+	// An END inside a nested quote/comment doesn't count.
+	curr := s.prev.Next(r, data)
+	if curr == nil {
+		s.prev = &ReadyState{}
+		s.statementStart = len(data)
+		s.statementHasContent = false
+		return s
+	}
+	s.prev = curr
+	if _, ok := s.prev.(*ReadyState); !ok {
+		return s
+	}
+	// PostgreSQL requires each inner statement to end with ';', so the closing END is
+	// always the first token of a statement; a later END is expression text.
+	if !s.statementHasContent && endsWithKeyword(data, END_ATOMIC) {
+		if isCommentsAndWhitespace(data[s.statementStart : len(data)-len(END_ATOMIC)]) {
+			s.pendingEnd = true
+		} else {
+			s.statementHasContent = true
 		}
 	}
 	return s

apps/cli-go/pkg/parser/state_test.go

@@ -77,6 +77,14 @@ func TestDollarQuote(t *testing.T) {
 		checkSplit(t, sql)
 	})
 
+	t.Run("non-ASCII named tag", func(t *testing.T) {
+		for _, tag := range []string{"a²", "a😀", "á"} {
+			t.Run(tag, func(t *testing.T) {
+				checkSplit(t, []string{"$" + tag + "$ any ; END; string$" + tag + "$;", " SELECT 2;"})
+			})
+		}
+	})
+
 	t.Run("anonymous tag", func(t *testing.T) {
 		sql := []string{"$$\"Dane's horse\"$$"}
 		checkSplit(t, sql)
@@ -193,6 +201,81 @@ SELECT 1;`,
 		}
 	})
 
+	t.Run("ignores end inside identifiers", func(t *testing.T) {
+		for _, name := range []string{"pending", "pending_change", "append", "legend", "𐐀end", "😀end", "́end", "²end", "pending$$foo$"} {
+			t.Run(name, func(t *testing.T) {
+				body := `CREATE FUNCTION f() RETURNS int LANGUAGE sql BEGIN ATOMIC SELECT 1 AS ` + name + `; END;`
+				checkSplit(t, []string{body, ` SELECT 2;`})
+			})
+		}
+	})
+
+	t.Run("ignores identifiers starting with end", func(t *testing.T) {
+		for _, name := range []string{"endpoint", "end_date", "ended_at", "end𐐀", "end😀", "end́", "end²", "end$$foo$"} {
+			t.Run(name, func(t *testing.T) {
+				body := `CREATE FUNCTION f() RETURNS int LANGUAGE sql BEGIN ATOMIC SELECT ` + name + `; SELECT 1; END;`
+				checkSplit(t, []string{body, ` SELECT 2;`})
+			})
+		}
+	})
+
+	t.Run("ignores end inside inner statements", func(t *testing.T) {
+		for _, expr := range []string{
+			"CASE WHEN true THEN 1 ELSE 0 END",
+			"case when true then 1 end",
+			"CASE WHEN true THEN 1 END AS ended",
+			"CASE WHEN CASE WHEN true THEN true END THEN 1 END",
+			"(CASE WHEN (true) THEN 1 END)",
+			"coalesce(CASE WHEN length('a') > 0 THEN 1 END, 0)",
+			"CASE(1)WHEN 1 THEN 1 END",
+			"1 AS case",
+			"1 case",
+			"1 AS end",
+			"1 end",
+			"'end'",
+		} {
+			t.Run(expr, func(t *testing.T) {
+				body := `CREATE FUNCTION f() RETURNS int LANGUAGE sql BEGIN ATOMIC SELECT ` + expr + `; SELECT 1; END;`
+				checkSplit(t, []string{body, ` SELECT 2;`})
+			})
+		}
+	})
+
+	t.Run("closes at end preceded only by comments", func(t *testing.T) {
+		for _, comment := range []string{"-- note END\n", "/* note; */ ", "\n/* a /* b; */ */ -- c\n"} {
+			t.Run(comment, func(t *testing.T) {
+				body := `CREATE FUNCTION f() RETURNS int LANGUAGE sql BEGIN ATOMIC SELECT 1; ` + comment + `END;`
+				checkSplit(t, []string{body, ` SELECT 2;`})
+			})
+		}
+	})
+
+	t.Run("ignores non-ASCII begin atomic lookalikes", func(t *testing.T) {
+		// atomıc (dotless ı) case-folds to ATOMIC but is a plain identifier in SQL.
+		checkSplit(t, []string{"BEGIN atomıc;", " SELECT 'end';", " SELECT 2;"})
+	})
+
+	t.Run("closes atomic body at end right after a positional parameter", func(t *testing.T) {
+		checkSplit(t, []string{"CREATE FUNCTION f(int) RETURNS int LANGUAGE sql BEGIN ATOMIC SELECT $1;END;", " SELECT 2;"})
+	})
+
+	t.Run("requires sql whitespace between begin and atomic", func(t *testing.T) {
+		for _, gap := range []string{"\u00A0", "\uFEFF", "\u0085"} {
+			t.Run(gap, func(t *testing.T) {
+				checkSplit(t, []string{"BEGIN " + gap + " ATOMIC;", " SELECT 1;", " end;", " SELECT 2;"})
+			})
+		}
+	})
+
+	t.Run("closes nested atomic body inside parentheses", func(t *testing.T) {
+		checkSplit(t, []string{"DO (BEGIN ATOMIC SELECT 1; END; );", " SELECT 2;"})
+	})
+
+	t.Run("ignores end after overlapping block comment", func(t *testing.T) {
+		body := `CREATE FUNCTION f() RETURNS int LANGUAGE sql BEGIN ATOMIC SELECT 1; /* a /*/ b */ SELECT 2 END; SELECT 3; END;`
+		checkSplit(t, []string{body, ` SELECT 4;`})
+	})
+
 	t.Run("does not treat schema-qualified atomic function names as begin atomic", func(t *testing.T) {
 		sql := []string{`CREATE OR REPLACE FUNCTION public.atomic_example()
 RETURNS INTEGER