Fix: handle unicode heredoc tags & Rust grapheme clusters properly (#6024)

Author: georgesittas

sqlglot/tokens.py

@@ -1421,7 +1421,11 @@ def _scan_string(self, start: str) -> bool:
                         raise_unmatched=not self.HEREDOC_TAG_IS_IDENTIFIER,
                     )
 
-                if tag and self.HEREDOC_TAG_IS_IDENTIFIER and (self._end or not tag.isidentifier()):
+                if (
+                    tag
+                    and self.HEREDOC_TAG_IS_IDENTIFIER
+                    and (self._end or tag.isdigit() or any(c.isspace() for c in tag))
+                ):
                     if not self._end:
                         self._advance(-1)
 

sqlglotrs/src/tokenizer.rs

@@ -444,13 +444,13 @@ impl<'a> TokenizerState<'a> {
 
                 if !tag.is_empty()
                     && self.settings.heredoc_tag_is_identifier
-                    && (self.is_end || !self.is_identifier(&tag))
+                    && (self.is_end || tag.chars().all(|c| c.is_ascii_digit()) || tag.chars().any(|c| c.is_whitespace()))
                 {
                     if !self.is_end {
                         self.advance(-1)?;
                     }
 
-                    self.advance(-(tag.len() as isize))?;
+                    self.advance(-(tag.chars().count() as isize))?;
                     self.add(self.token_types.heredoc_string_alternative, None)?;
                     return Ok(true);
                 }
@@ -687,9 +687,10 @@ impl<'a> TokenizerState<'a> {
                     continue;
                 }
             }
-            if self.chars(delimiter.len()) == delimiter {
-                if delimiter.len() > 1 {
-                    self.advance((delimiter.len() - 1) as isize)?;
+            let delimiter_char_count = delimiter.chars().count();
+            if self.chars(delimiter_char_count) == delimiter {
+                if delimiter_char_count > 1 {
+                    self.advance((delimiter_char_count - 1) as isize)?;
                 }
                 break;
             }
@@ -720,16 +721,6 @@ impl<'a> TokenizerState<'a> {
         name.is_alphabetic() || name == '_'
     }
 
-    fn is_identifier(&self, s: &str) -> bool {
-        s.chars().enumerate().all(|(i, c)| {
-            if i == 0 {
-                self.is_alphabetic_or_underscore(c)
-            } else {
-                self.is_alphabetic_or_underscore(c) || c.is_ascii_digit()
-            }
-        })
-    }
-
     fn is_numeric(&self, s: &str) -> bool {
         s.chars().all(|c| c.is_ascii_digit())
     }

tests/dialects/test_duckdb.py

@@ -363,6 +363,10 @@ def test_duckdb(self):
         self.validate_identity(
             """SELECT '{ "family": "anatidae", "species": [ "duck", "goose", "swan", null ] }' ->> ['$.family', '$.species']""",
         )
+        self.validate_identity(
+            "SELECT $🦆$foo$🦆$",
+            "SELECT 'foo'",
+        )
         self.validate_identity(
             "SELECT * FROM t PIVOT(FIRST(t) AS t, FOR quarter IN ('Q1', 'Q2'))",
             "SELECT * FROM t PIVOT(FIRST(t) AS t FOR quarter IN ('Q1', 'Q2'))",