Repository: tobymao/sqlglot

Issue:
BigQuery raw string support has regressed
This issue is present on the current release and on main.

In the following:
```
>>> q = "SELECT r'\s'"
>>> print(sqlglot.parse_one(q, "bigquery").sql("bigquery"))
SELECT '\s'
```

`SELECT r'\s'` is not the same as `SELECT '\s'` - the `r` prefix indicates the following is a "raw string," mostly used in REGEXP_REPLACE/REGEXP_EXTRACT ([the docs](https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical) say it is also referred to as a "regex string")

It looks like this was addressed in  #218, but it appears to have regressed in https://github.com/tobymao/sqlglot/commit/d2377e04eacdb1e9793d77629baa9c0a751fd50c

The test [here in test_bigquery.py](https://github.com/tobymao/sqlglot/blob/eb6eaf5340bb7001e047583dc53735b54612c860/tests/dialects/test_bigquery.py#L110) is now incorrect:
```
        self.validate_all(
            r'R"""/\*.*\*/"""',
            write={
                "bigquery": r"'/\*.*\*/'",
                "duckdb": r"'/\*.*\*/'",
                "presto": r"'/\*.*\*/'",
                "hive": r"'/\*.*\*/'",
                "spark": r"'/\*.*\*/'",
            },
        )
```
The query string `R"""/\*.*\*/"""` starts with `R"`, so if the leading R is removed, all the backslashes in the resulting strings should be double escaped.

I'm working on wrapping my head around the tokenizer to see if I can fix this, but I thought I'd post an issue in the interim in case I'm called away from working on this before I can get to a PR.

**Fully reproducible code snippet**
`assert sqlglot.parse_one(r"SELECT r'\s'", "bigquery").sql("bigquery") == r"SELECT '\\s'"`


**Official Documentation**
https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#quoted_literals


Discussion:

Patch:
diff --git a/sqlglot/dialects/bigquery.py b/sqlglot/dialects/bigquery.py
index 8e622476..d05d99ff 100644
--- a/sqlglot/dialects/bigquery.py
+++ b/sqlglot/dialects/bigquery.py
@@ -116,16 +116,22 @@ class BigQuery(Dialect):
     }
 
     class Tokenizer(tokens.Tokenizer):
-        QUOTES = [
-            (prefix + quote, quote) if prefix else quote
-            for quote in ["'", '"', '"""', "'''"]
-            for prefix in ["", "r", "R"]
-        ]
+        QUOTES = ["'", '"', '"""', "'''"]
         COMMENTS = ["--", "#", ("/*", "*/")]
         IDENTIFIERS = ["`"]
         STRING_ESCAPES = ["\\"]
-        HEX_STRINGS = [("0x", ""), ("0X", "")]
-        BYTE_STRINGS = [("b'", "'"), ("B'", "'")]
+
+        HEX_STRINGS = [
+            (prefix + q, q) for q in [""] + t.cast(t.List[str], QUOTES) for prefix in ("0x", "0X")
+        ]
+
+        BYTE_STRINGS = [
+            (prefix + q, q) for q in t.cast(t.List[str], QUOTES) for prefix in ("b", "B")
+        ]
+
+        RAW_STRINGS = [
+            (prefix + q, q) for q in t.cast(t.List[str], QUOTES) for prefix in ("r", "R")
+        ]
 
         KEYWORDS = {
             **tokens.Tokenizer.KEYWORDS,
diff --git a/sqlglot/dialects/dialect.py b/sqlglot/dialects/dialect.py
index 55890cdf..013dd2ae 100644
--- a/sqlglot/dialects/dialect.py
+++ b/sqlglot/dialects/dialect.py
@@ -8,7 +8,7 @@ from sqlglot.generator import Generator
 from sqlglot.helper import flatten, seq_get
 from sqlglot.parser import Parser
 from sqlglot.time import format_time
-from sqlglot.tokens import Token, Tokenizer
+from sqlglot.tokens import Token, Tokenizer, TokenType
 from sqlglot.trie import new_trie
 
 if t.TYPE_CHECKING:
@@ -84,17 +84,20 @@ class _Dialect(type):
             klass.tokenizer_class._IDENTIFIERS.items()
         )[0]
 
-        klass.bit_start, klass.bit_end = seq_get(
-            list(klass.tokenizer_class._BIT_STRINGS.items()), 0
-        ) or (None, None)
-
-        klass.hex_start, klass.hex_end = seq_get(
-            list(klass.tokenizer_class._HEX_STRINGS.items()), 0
-        ) or (None, None)
+        def get_start_end(token_type: TokenType) -> t.Tuple[t.Optional[str], t.Optional[str]]:
+            return next(
+                (
+                    (s, e)
+                    for s, (e, t) in klass.tokenizer_class._FORMAT_STRINGS.items()
+                    if t == token_type
+                ),
+                (None, None),
+            )
 
-        klass.byte_start, klass.byte_end = seq_get(
-            list(klass.tokenizer_class._BYTE_STRINGS.items()), 0
-        ) or (None, None)
+        klass.bit_start, klass.bit_end = get_start_end(TokenType.BIT_STRING)
+        klass.hex_start, klass.hex_end = get_start_end(TokenType.HEX_STRING)
+        klass.byte_start, klass.byte_end = get_start_end(TokenType.BYTE_STRING)
+        klass.raw_start, klass.raw_end = get_start_end(TokenType.RAW_STRING)
 
         return klass
 
@@ -212,6 +215,8 @@ class Dialect(metaclass=_Dialect):
                 "hex_end": self.hex_end,
                 "byte_start": self.byte_start,
                 "byte_end": self.byte_end,
+                "raw_start": self.raw_start,
+                "raw_end": self.raw_end,
                 "identifier_start": self.identifier_start,
                 "identifier_end": self.identifier_end,
                 "string_escape": self.tokenizer_class.STRING_ESCAPES[0],
diff --git a/sqlglot/dialects/hive.py b/sqlglot/dialects/hive.py
index a9a893ff..f4dddfb4 100644
--- a/sqlglot/dialects/hive.py
+++ b/sqlglot/dialects/hive.py
@@ -368,7 +368,7 @@ class Hive(Dialect):
             exp.SerdeProperties: lambda self, e: self.properties(e, prefix="WITH SERDEPROPERTIES"),
             exp.NumberToStr: rename_func("FORMAT_NUMBER"),
             exp.LastDateOfMonth: rename_func("LAST_DAY"),
-            exp.National: lambda self, e: self.sql(e, "this"),
+            exp.National: lambda self, e: self.national_sql(e, prefix=""),
         }
 
         PROPERTIES_LOCATION = {
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 595632c4..94a9220f 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -1091,6 +1091,10 @@ class ByteString(Condition):
     pass
 
 
+class RawString(Condition):
+    pass
+
+
 class Column(Condition):
     arg_types = {"this": True, "table": False, "db": False, "catalog": False, "join_mark": False}
 
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index 7bd0d36a..2c9c7085 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -31,6 +31,8 @@ class Generator:
         hex_end (str): specifies which ending character to use to delimit hex literals. Default: None.
         byte_start (str): specifies which starting character to use to delimit byte literals. Default: None.
         byte_end (str): specifies which ending character to use to delimit byte literals. Default: None.
+        raw_start (str): specifies which starting character to use to delimit raw literals. Default: None.
+        raw_end (str): specifies which ending character to use to delimit raw literals. Default: None.
         identify (bool | str): 'always': always quote, 'safe': quote identifiers if they don't contain an upcase, True defaults to always.
         normalize (bool): if set to True all identifiers will lower cased
         string_escape (str): specifies a string escape character. Default: '.
@@ -249,6 +251,8 @@ class Generator:
         "hex_end",
         "byte_start",
         "byte_end",
+        "raw_start",
+        "raw_end",
         "identify",
         "normalize",
         "string_escape",
@@ -286,6 +290,8 @@ class Generator:
         hex_end=None,
         byte_start=None,
         byte_end=None,
+        raw_start=None,
+        raw_end=None,
         identify=False,
         normalize=False,
         string_escape=None,
@@ -318,6 +324,8 @@ class Generator:
         self.hex_end = hex_end
         self.byte_start = byte_start
         self.byte_end = byte_end
+        self.raw_start = raw_start
+        self.raw_end = raw_end
         self.identify = identify
         self.normalize = normalize
         self.string_escape = string_escape or "'"
@@ -793,6 +801,11 @@ class Generator:
             return f"{self.byte_start}{this}{self.byte_end}"
         return this
 
+    def rawstring_sql(self, expression: exp.RawString) -> str:
+        if self.raw_start:
+            return f"{self.raw_start}{expression.name}{self.raw_end}"
+        return self.sql(exp.Literal.string(expression.name.replace("\\", "\\\\")))
+
     def datatypesize_sql(self, expression: exp.DataTypeSize) -> str:
         this = self.sql(expression, "this")
         specifier = self.sql(expression, "expression")
@@ -903,8 +916,9 @@ class Generator:
         output_format = f"OUTPUTFORMAT {output_format}" if output_format else ""
         return self.sep().join((input_format, output_format))
 
-    def national_sql(self, expression: exp.National) -> str:
-        return f"N{self.sql(expression, 'this')}"
+    def national_sql(self, expression: exp.National, prefix: str = "N") -> str:
+        string = self.sql(exp.Literal.string(expression.name))
+        return f"{prefix}{string}"
 
     def partition_sql(self, expression: exp.Partition) -> str:
         return f"PARTITION({self.expressions(expression)})"
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index c879e588..478b2c76 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -521,7 +521,10 @@ class Parser(metaclass=_Parser):
         TokenType.HEX_STRING: lambda self, token: self.expression(exp.HexString, this=token.text),
         TokenType.BYTE_STRING: lambda self, token: self.expression(exp.ByteString, this=token.text),
         TokenType.INTRODUCER: lambda self, token: self._parse_introducer(token),
-        TokenType.NATIONAL: lambda self, token: self._parse_national(token),
+        TokenType.NATIONAL_STRING: lambda self, token: self.expression(
+            exp.National, this=token.text
+        ),
+        TokenType.RAW_STRING: lambda self, token: self.expression(exp.RawString, this=token.text),
         TokenType.SESSION_PARAMETER: lambda self, _: self._parse_session_parameter(),
     }
 
@@ -3154,9 +3157,6 @@ class Parser(metaclass=_Parser):
 
         return self.expression(exp.Identifier, this=token.text)
 
-    def _parse_national(self, token: Token) -> exp.Expression:
-        return self.expression(exp.National, this=exp.Literal.string(token.text))
-
     def _parse_session_parameter(self) -> exp.Expression:
         kind = None
         this = self._parse_id_var() or self._parse_primary()
diff --git a/sqlglot/tokens.py b/sqlglot/tokens.py
index 9e37e774..ad329d26 100644
--- a/sqlglot/tokens.py
+++ b/sqlglot/tokens.py
@@ -51,7 +51,6 @@ class TokenType(AutoName):
     DOLLAR = auto()
     PARAMETER = auto()
     SESSION_PARAMETER = auto()
-    NATIONAL = auto()
     DAMP = auto()
 
     BLOCK_START = auto()
@@ -72,6 +71,8 @@ class TokenType(AutoName):
     BIT_STRING = auto()
     HEX_STRING = auto()
     BYTE_STRING = auto()
+    NATIONAL_STRING = auto()
+    RAW_STRING = auto()
 
     # types
     BIT = auto()
@@ -350,15 +351,31 @@ class _Tokenizer(type):
     def __new__(cls, clsname, bases, attrs):
         klass = super().__new__(cls, clsname, bases, attrs)
 
-        klass._QUOTES = {
-            f"{prefix}{s}": e
-            for s, e in cls._delimeter_list_to_dict(klass.QUOTES).items()
-            for prefix in (("",) if s[0].isalpha() else ("", "n", "N"))
+        def _convert_quotes(arr: t.List[str | t.Tuple[str, str]]) -> t.Dict[str, str]:
+            return dict(
+                (item, item) if isinstance(item, str) else (item[0], item[1]) for item in arr
+            )
+
+        def _quotes_to_format(
+            token_type: TokenType, arr: t.List[str | t.Tuple[str, str]]
+        ) -> t.Dict[str, t.Tuple[str, TokenType]]:
+            return {k: (v, token_type) for k, v in _convert_quotes(arr).items()}
+
+        klass._QUOTES = _convert_quotes(klass.QUOTES)
+        klass._IDENTIFIERS = _convert_quotes(klass.IDENTIFIERS)
+
+        klass._FORMAT_STRINGS = {
+            **{
+                p + s: (e, TokenType.NATIONAL_STRING)
+                for s, e in klass._QUOTES.items()
+                for p in ("n", "N")
+            },
+            **_quotes_to_format(TokenType.BIT_STRING, klass.BIT_STRINGS),
+            **_quotes_to_format(TokenType.BYTE_STRING, klass.BYTE_STRINGS),
+            **_quotes_to_format(TokenType.HEX_STRING, klass.HEX_STRINGS),
+            **_quotes_to_format(TokenType.RAW_STRING, klass.RAW_STRINGS),
         }
-        klass._BIT_STRINGS = cls._delimeter_list_to_dict(klass.BIT_STRINGS)
-        klass._HEX_STRINGS = cls._delimeter_list_to_dict(klass.HEX_STRINGS)
-        klass._BYTE_STRINGS = cls._delimeter_list_to_dict(klass.BYTE_STRINGS)
-        klass._IDENTIFIERS = cls._delimeter_list_to_dict(klass.IDENTIFIERS)
+
         klass._STRING_ESCAPES = set(klass.STRING_ESCAPES)
         klass._IDENTIFIER_ESCAPES = set(klass.IDENTIFIER_ESCAPES)
         klass._COMMENTS = dict(
@@ -372,19 +389,13 @@ class _Tokenizer(type):
                 *klass.KEYWORDS,
                 *klass._COMMENTS,
                 *klass._QUOTES,
-                *klass._BIT_STRINGS,
-                *klass._HEX_STRINGS,
-                *klass._BYTE_STRINGS,
+                *klass._FORMAT_STRINGS,
             )
             if " " in key or any(single in key for single in klass.SINGLE_TOKENS)
         )
 
         return klass
 
-    @staticmethod
-    def _delimeter_list_to_dict(list: t.List[str | t.Tuple[str, str]]) -> t.Dict[str, str]:
-        return dict((item, item) if isinstance(item, str) else (item[0], item[1]) for item in list)
-
 
 class Tokenizer(metaclass=_Tokenizer):
     SINGLE_TOKENS = {
@@ -425,6 +436,7 @@ class Tokenizer(metaclass=_Tokenizer):
     BIT_STRINGS: t.List[str | t.Tuple[str, str]] = []
     BYTE_STRINGS: t.List[str | t.Tuple[str, str]] = []
     HEX_STRINGS: t.List[str | t.Tuple[str, str]] = []
+    RAW_STRINGS: t.List[str | t.Tuple[str, str]] = []
     IDENTIFIERS: t.List[str | t.Tuple[str, str]] = ['"']
     IDENTIFIER_ESCAPES = ['"']
     QUOTES: t.List[t.Tuple[str, str] | str] = ["'"]
@@ -432,9 +444,7 @@ class Tokenizer(metaclass=_Tokenizer):
     VAR_SINGLE_TOKENS: t.Set[str] = set()
 
     _COMMENTS: t.Dict[str, str] = {}
-    _BIT_STRINGS: t.Dict[str, str] = {}
-    _BYTE_STRINGS: t.Dict[str, str] = {}
-    _HEX_STRINGS: t.Dict[str, str] = {}
+    _FORMAT_STRINGS: t.Dict[str, t.Tuple[str, TokenType]] = {}
     _IDENTIFIERS: t.Dict[str, str] = {}
     _IDENTIFIER_ESCAPES: t.Set[str] = set()
     _QUOTES: t.Dict[str, str] = {}
@@ -898,8 +908,6 @@ class Tokenizer(metaclass=_Tokenizer):
 
         if self._scan_string(word):
             return
-        if self._scan_formatted_string(word):
-            return
         if self._scan_comment(word):
             return
 
@@ -943,9 +951,9 @@ class Tokenizer(metaclass=_Tokenizer):
         if self._char == "0":
             peek = self._peek.upper()
             if peek == "B":
-                return self._scan_bits() if self._BIT_STRINGS else self._add(TokenType.NUMBER)
+                return self._scan_bits() if self.BIT_STRINGS else self._add(TokenType.NUMBER)
             elif peek == "X":
-                return self._scan_hex() if self._HEX_STRINGS else self._add(TokenType.NUMBER)
+                return self._scan_hex() if self.HEX_STRINGS else self._add(TokenType.NUMBER)
 
         decimal = False
         scientific = 0
@@ -1014,37 +1022,24 @@ class Tokenizer(metaclass=_Tokenizer):
 
         return self._text
 
-    def _scan_string(self, quote: str) -> bool:
-        quote_end = self._QUOTES.get(quote)
-        if quote_end is None:
-            return False
+    def _scan_string(self, start: str) -> bool:
+        base = None
+        token_type = TokenType.STRING
 
-        self._advance(len(quote))
-        text = self._extract_string(quote_end)
-        text = text.encode(self.ENCODE).decode(self.ENCODE) if self.ENCODE else text
-        self._add(TokenType.NATIONAL if quote[0].upper() == "N" else TokenType.STRING, text)
-        return True
+        if start in self._QUOTES:
+            end = self._QUOTES[start]
+        elif start in self._FORMAT_STRINGS:
+            end, token_type = self._FORMAT_STRINGS[start]
 
-    # X'1234', b'0110', E'\\\\\' etc.
-    def _scan_formatted_string(self, string_start: str) -> bool:
-        if string_start in self._HEX_STRINGS:
-            delimiters = self._HEX_STRINGS
-            token_type = TokenType.HEX_STRING
-            base = 16
-        elif string_start in self._BIT_STRINGS:
-            delimiters = self._BIT_STRINGS
-            token_type = TokenType.BIT_STRING
-            base = 2
-        elif string_start in self._BYTE_STRINGS:
-            delimiters = self._BYTE_STRINGS
-            token_type = TokenType.BYTE_STRING
-            base = None
+            if token_type == TokenType.HEX_STRING:
+                base = 16
+            elif token_type == TokenType.BIT_STRING:
+                base = 2
         else:
             return False
 
-        self._advance(len(string_start))
-        string_end = delimiters[string_start]
-        text = self._extract_string(string_end)
+        self._advance(len(start))
+        text = self._extract_string(end)
 
         if base:
             try:
@@ -1053,6 +1048,8 @@ class Tokenizer(metaclass=_Tokenizer):
                 raise RuntimeError(
                     f"Numeric string contains invalid characters from {self._line}:{self._start}"
                 )
+        else:
+            text = text.encode(self.ENCODE).decode(self.ENCODE) if self.ENCODE else text
 
         self._add(token_type, text)
         return True


Test patch:
diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
index 59ed7eb6..72d09933 100644
--- a/tests/dialects/test_bigquery.py
+++ b/tests/dialects/test_bigquery.py
@@ -85,16 +85,6 @@ class TestBigQuery(Validator):
                 "spark": "'x\\''",
             },
         )
-        self.validate_all(
-            r'r"""/\*.*\*/"""',
-            write={
-                "bigquery": r"'/\*.*\*/'",
-                "duckdb": r"'/\*.*\*/'",
-                "presto": r"'/\*.*\*/'",
-                "hive": r"'/\*.*\*/'",
-                "spark": r"'/\*.*\*/'",
-            },
-        )
         with self.assertRaises(ValueError):
             transpile("'\\'", read="bigquery")
 
@@ -107,14 +97,24 @@ class TestBigQuery(Validator):
                 "hive": r"'\\'",
             },
         )
+        self.validate_all(
+            r'r"""/\*.*\*/"""',
+            write={
+                "bigquery": r"r'/\*.*\*/'",
+                "duckdb": r"'/\\*.*\\*/'",
+                "presto": r"'/\\*.*\\*/'",
+                "hive": r"'/\\*.*\\*/'",
+                "spark": r"'/\\*.*\\*/'",
+            },
+        )
         self.validate_all(
             r'R"""/\*.*\*/"""',
             write={
-                "bigquery": r"'/\*.*\*/'",
-                "duckdb": r"'/\*.*\*/'",
-                "presto": r"'/\*.*\*/'",
-                "hive": r"'/\*.*\*/'",
-                "spark": r"'/\*.*\*/'",
+                "bigquery": r"r'/\*.*\*/'",
+                "duckdb": r"'/\\*.*\\*/'",
+                "presto": r"'/\\*.*\\*/'",
+                "hive": r"'/\\*.*\\*/'",
+                "spark": r"'/\\*.*\\*/'",
             },
         )
         self.validate_all(