From 0f9c62c9641bc2a91d2e3031ed7a9117420ffa25 Mon Sep 17 00:00:00 2001
From: 7ttp <117663341+7ttp@users.noreply.github.com>
Date: Sat, 11 Apr 2026 02:25:18 +0530
Subject: [PATCH 1/2] fix

---
 pkg/parser/state.go | 30 ++++++++++++++++++++++++++++--
 1 file changed, 28 insertions(+), 2 deletions(-)

diff --git a/pkg/parser/state.go b/pkg/parser/state.go
index f32a671315..47775390d1 100644
--- a/pkg/parser/state.go
+++ b/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{}
 

From c05e35b1ac590d84efb7f7206987873a016b23d3 Mon Sep 17 00:00:00 2001
From: 7ttp <117663341+7ttp@users.noreply.github.com>
Date: Sat, 11 Apr 2026 02:25:51 +0530
Subject: [PATCH 2/2] test

---
 pkg/parser/state_test.go | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/pkg/parser/state_test.go b/pkg/parser/state_test.go
index bae10fe190..ad6db9d26a 100644
--- a/pkg/parser/state_test.go
+++ b/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)
+	})
 }