fix: atomic parser (#5064)

* fix

* test

---------

Co-authored-by: Andrew Valleteau <avallete@users.noreply.github.com>

Author: 7ttp

pkg/parser/state.go

@@ -46,14 +46,40 @@ func (s *ReadyState) Next(r rune, data []byte) State {
 	case 'c':
 		fallthrough
 	case 'C':
-		offset := len(data) - len(BEGIN_ATOMIC)
-		if offset >= 0 && strings.EqualFold(string(data[offset:]), BEGIN_ATOMIC) {
+		if isBeginAtomic(data) {
 			return &AtomicState{prev: s, delimiter: []byte(END_ATOMIC)}
 		}
 	}
 	return s
 }
 
+func isBeginAtomic(data []byte) bool {
+	offset := len(data) - len(BEGIN_ATOMIC)
+	if offset < 0 || !strings.EqualFold(string(data[offset:]), 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") {
+		return false
+	}
+	if offset == 0 {
+		return true
+	}
+	r, _ := utf8.DecodeLastRune(prefix[:offset])
+	return !isIdentifierRune(r)
+}
+
+func isIdentifierRune(r rune) bool {
+	return unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_' || r == '$'
+}
+
 // Opened a line comment
 type CommentState struct{}
 

pkg/parser/state_test.go

@@ -167,4 +167,44 @@ END
 ;`}
 		checkSplit(t, sql)
 	})
+
+	t.Run("ignores atomic in identifiers", func(t *testing.T) {
+		names := []string{
+			"fn_atomic",
+			"atomic_fn",
+			"my_atomic_thing",
+			"xatomicx",
+			"fn_ATomiC",
+		}
+		for _, name := range names {
+			t.Run(name, func(t *testing.T) {
+				sql := []string{
+					`CREATE OR REPLACE FUNCTION ` + name + `()
+RETURNS void LANGUAGE plpgsql AS $$
+BEGIN
+  NULL;
+END;
+$$;`,
+					`
+SELECT 1;`,
+				}
+				checkSplit(t, sql)
+			})
+		}
+	})
+
+	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
+LANGUAGE plpgsql
+AS $$
+BEGIN
+  RETURN 1;
+END;
+$$;`,
+			`
+GRANT EXECUTE ON FUNCTION public.atomic_example() TO authenticated;`,
+		}
+		checkSplit(t, sql)
+	})
 }