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{} 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) + }) }