fix: parser treats 'atomic' in identifiers as BEGIN ATOMIC keyword

pkg/parser/state.go

@@ -47,17 +47,16 @@
 		fallthrough
 	case 'C':
 		offset := len(data) - len(BEGIN_ATOMIC)
-		if offset >= 0 && strings.EqualFold(string(data[offset:]), BEGIN_ATOMIC) {
+		if offset > 0 && strings.EqualFold(string(data[offset:]), BEGIN_ATOMIC) && unicode.IsSpace(rune(data[offset-1])) {
 			return &AtomicState{prev: s, delimiter: []byte(END_ATOMIC)}
 		}
-	}
 	return s
 }
 
 // Opened a line comment
 type CommentState struct{}
 
 func (s *CommentState) Next(r rune, data []byte) State {
 	if r == '-' {
 		// No characters are escaped in comments, which is the same as dollar
 		return &DollarState{delimiter: []byte{'\n'}}

pkg/parser/state_test.go

@@ -167,4 +167,27 @@ END
 ;`}
 		checkSplit(t, sql)
 	})
+
+	t.Run("atomic in table name", func(t *testing.T) {
+		sql := []string{"create table public.atomic_test (id int);", "\nselect 1;"}
+		checkSplit(t, sql)
+	})
+
+	t.Run("atomic in function name", func(t *testing.T) {
+		sql := []string{
+			"CREATE OR REPLACE FUNCTION public.my_function_atomic(p_id uuid)\nRETURNS void\nLANGUAGE plpgsql\nAS $$\nBEGIN\n  RAISE NOTICE 'hello';\nEND;\n$$;",
+			"\nGRANT EXECUTE ON FUNCTION public.my_function_atomic(uuid) TO authenticated;",
+		}
+		checkSplit(t, sql)
+	})
+
+	t.Run("atomic in column name", func(t *testing.T) {
+		sql := []string{"select is_atomic from flags;", "\nselect 1;"}
+		checkSplit(t, sql)
+	})
+
+	t.Run("digit before atomic", func(t *testing.T) {
+		sql := []string{"select col1atomic from t;", "\nselect 1;"}
+		checkSplit(t, sql)
+	})
 }