Repository: tobymao/sqlglot

Issue:
Tokenizer Error with raw string ending with \
```python
import sqlglot

sql = r'select r"\\apath.local\" '
print(sql)
sqlglot.parse_one(
    sql,
    dialect="databricks",
)

```
this creates the following error:
```
raceback (most recent call last):
  File "C:\Projects\DataHub\Fabricks.Runtime\.venv\Lib\site-packages\sqlglot\tokens.py", line 973, in tokenize
    self._scan()
  File "C:\Projects\DataHub\Fabricks.Runtime\.venv\Lib\site-packages\sqlglot\tokens.py", line 1006, in _scan
    self._scan_keywords()
  File "C:\Projects\DataHub\Fabricks.Runtime\.venv\Lib\site-packages\sqlglot\tokens.py", line 1137, in _scan_keywords
    if self._scan_string(word):
       ^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Projects\DataHub\Fabricks.Runtime\.venv\Lib\site-packages\sqlglot\tokens.py", line 1300, in _scan_string
    text = self._extract_string(end, unescape_sequences=token_type != TokenType.RAW_STRING)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Projects\DataHub\Fabricks.Runtime\.venv\Lib\site-packages\sqlglot\tokens.py", line 1379, in _extract_string
    raise TokenError(f"Missing {delimiter} from {self._line}:{self._start}")
sqlglot.errors.TokenError: Missing " from 1:7
```

The issue only happens if the raw string ends with \

Discussion:

Patch:
diff --git a/sqlglot/dialects/spark.py b/sqlglot/dialects/spark.py
index 0212352d..b71204b8 100644
--- a/sqlglot/dialects/spark.py
+++ b/sqlglot/dialects/spark.py
@@ -91,6 +91,8 @@ def _dateadd_sql(self: Spark.Generator, expression: exp.TsOrDsAdd | exp.Timestam
 
 class Spark(Spark2):
     class Tokenizer(Spark2.Tokenizer):
+        STRING_ESCAPES_ALLOWED_IN_RAW_STRINGS = False
+
         RAW_STRINGS = [
             (prefix + q, q)
             for q in t.cast(t.List[str], Spark2.Tokenizer.QUOTES)
diff --git a/sqlglot/tokens.py b/sqlglot/tokens.py
index 12f4a673..9a816a72 100644
--- a/sqlglot/tokens.py
+++ b/sqlglot/tokens.py
@@ -526,6 +526,7 @@ class _Tokenizer(type):
                     _TOKEN_TYPE_TO_INDEX[v] for v in klass.COMMAND_PREFIX_TOKENS
                 },
                 heredoc_tag_is_identifier=klass.HEREDOC_TAG_IS_IDENTIFIER,
+                string_escapes_allowed_in_raw_strings=klass.STRING_ESCAPES_ALLOWED_IN_RAW_STRINGS,
             )
             token_types = RsTokenTypeSettings(
                 bit_string=_TOKEN_TYPE_TO_INDEX[TokenType.BIT_STRING],
@@ -602,6 +603,9 @@ class Tokenizer(metaclass=_Tokenizer):
     # Token that we'll generate as a fallback if the heredoc prefix doesn't correspond to a heredoc
     HEREDOC_STRING_ALTERNATIVE = TokenType.VAR
 
+    # Whether string escape characters function as such when placed within raw strings
+    STRING_ESCAPES_ALLOWED_IN_RAW_STRINGS = True
+
     # Autofilled
     _COMMENTS: t.Dict[str, str] = {}
     _FORMAT_STRINGS: t.Dict[str, t.Tuple[str, TokenType]] = {}
@@ -1292,7 +1296,7 @@ class Tokenizer(metaclass=_Tokenizer):
                 else:
                     tag = self._extract_string(
                         end,
-                        unescape_sequences=False,
+                        raw_string=True,
                         raise_unmatched=not self.HEREDOC_TAG_IS_IDENTIFIER,
                     )
 
@@ -1309,7 +1313,7 @@ class Tokenizer(metaclass=_Tokenizer):
             return False
 
         self._advance(len(start))
-        text = self._extract_string(end, unescape_sequences=token_type != TokenType.RAW_STRING)
+        text = self._extract_string(end, raw_string=token_type == TokenType.RAW_STRING)
 
         if base:
             try:
@@ -1345,7 +1349,7 @@ class Tokenizer(metaclass=_Tokenizer):
         self,
         delimiter: str,
         escapes: t.Optional[t.Set[str]] = None,
-        unescape_sequences: bool = True,
+        raw_string: bool = False,
         raise_unmatched: bool = True,
     ) -> str:
         text = ""
@@ -1354,7 +1358,7 @@ class Tokenizer(metaclass=_Tokenizer):
 
         while True:
             if (
-                unescape_sequences
+                not raw_string
                 and self.dialect.UNESCAPED_SEQUENCES
                 and self._peek
                 and self._char in self.STRING_ESCAPES
@@ -1365,7 +1369,8 @@ class Tokenizer(metaclass=_Tokenizer):
                     text += unescaped_sequence
                     continue
             if (
-                self._char in escapes
+                (self.STRING_ESCAPES_ALLOWED_IN_RAW_STRINGS or not raw_string)
+                and self._char in escapes
                 and (self._peek == delimiter or self._peek in escapes)
                 and (self._char not in self._QUOTES or self._char == self._peek)
             ):
diff --git a/sqlglotrs/src/settings.rs b/sqlglotrs/src/settings.rs
index 4cacb9b5..9fb8fdae 100644
--- a/sqlglotrs/src/settings.rs
+++ b/sqlglotrs/src/settings.rs
@@ -76,6 +76,7 @@ pub struct TokenizerSettings {
     pub commands: HashSet<TokenType>,
     pub command_prefix_tokens: HashSet<TokenType>,
     pub heredoc_tag_is_identifier: bool,
+    pub string_escapes_allowed_in_raw_strings: bool,
 }
 
 #[pymethods]
@@ -98,6 +99,7 @@ impl TokenizerSettings {
         commands: HashSet<TokenType>,
         command_prefix_tokens: HashSet<TokenType>,
         heredoc_tag_is_identifier: bool,
+        string_escapes_allowed_in_raw_strings: bool,
     ) -> Self {
         let to_char = |v: &String| {
             if v.len() == 1 {
@@ -147,6 +149,7 @@ impl TokenizerSettings {
             commands,
             command_prefix_tokens,
             heredoc_tag_is_identifier,
+            string_escapes_allowed_in_raw_strings,
         }
     }
 }
diff --git a/sqlglotrs/src/tokenizer.rs b/sqlglotrs/src/tokenizer.rs
index 84ed25d9..ca5c44b5 100644
--- a/sqlglotrs/src/tokenizer.rs
+++ b/sqlglotrs/src/tokenizer.rs
@@ -424,7 +424,7 @@ impl<'a> TokenizerState<'a> {
                 let tag = if self.current_char.to_string() == *end {
                     String::from("")
                 } else {
-                    self.extract_string(end, false, false, !self.settings.heredoc_tag_is_identifier)?
+                    self.extract_string(end, false, true, !self.settings.heredoc_tag_is_identifier)?
                 };
 
                 if !tag.is_empty()
@@ -449,7 +449,7 @@ impl<'a> TokenizerState<'a> {
         };
 
         self.advance(start.len() as isize)?;
-        let text = self.extract_string(&end, false, token_type != self.token_types.raw_string, true)?;
+        let text = self.extract_string(&end, false, token_type == self.token_types.raw_string, true)?;
 
         if let Some(b) = base {
             if u64::from_str_radix(&text, b).is_err() {
@@ -595,7 +595,7 @@ impl<'a> TokenizerState<'a> {
 
     fn scan_identifier(&mut self, identifier_end: &str) -> Result<(), TokenizerError> {
         self.advance(1)?;
-        let text = self.extract_string(identifier_end, true, true, true)?;
+        let text = self.extract_string(identifier_end, true, false, true)?;
         self.add(self.token_types.identifier, Some(text))
     }
 
@@ -603,7 +603,7 @@ impl<'a> TokenizerState<'a> {
         &mut self,
         delimiter: &str,
         use_identifier_escapes: bool,
-        unescape_sequences: bool,
+        raw_string: bool,
         raise_unmatched: bool,
     ) -> Result<String, TokenizerError> {
         let mut text = String::from("");
@@ -616,7 +616,7 @@ impl<'a> TokenizerState<'a> {
             };
             let peek_char_str = self.peek_char.to_string();
 
-            if unescape_sequences
+            if !raw_string
                 && !self.dialect_settings.unescaped_sequences.is_empty()
                 && !self.peek_char.is_whitespace()
                 && self.settings.string_escapes.contains(&self.current_char)
@@ -631,7 +631,8 @@ impl<'a> TokenizerState<'a> {
                 }
             }
 
-            if escapes.contains(&self.current_char)
+            if (self.settings.string_escapes_allowed_in_raw_strings || !raw_string)
+                && escapes.contains(&self.current_char)
                 && (peek_char_str == delimiter || escapes.contains(&self.peek_char))
                 && (self.current_char == self.peek_char
                     || !self


Test patch:
diff --git a/tests/dialects/test_databricks.py b/tests/dialects/test_databricks.py
index c88679dc..785ea922 100644
--- a/tests/dialects/test_databricks.py
+++ b/tests/dialects/test_databricks.py
@@ -25,6 +25,7 @@ class TestDatabricks(Validator):
         self.validate_identity("CREATE FUNCTION a AS b")
         self.validate_identity("SELECT ${x} FROM ${y} WHERE ${z} > 1")
         self.validate_identity("CREATE TABLE foo (x DATE GENERATED ALWAYS AS (CAST(y AS DATE)))")
+        self.validate_identity("TRUNCATE TABLE t1 PARTITION(age = 10, name = 'test', address)")
         self.validate_identity(
             "CREATE TABLE IF NOT EXISTS db.table (a TIMESTAMP, b BOOLEAN GENERATED ALWAYS AS (NOT a IS NULL)) USING DELTA"
         )
@@ -37,22 +38,22 @@ class TestDatabricks(Validator):
         self.validate_identity(
             "SELECT * FROM sales UNPIVOT EXCLUDE NULLS (sales FOR quarter IN (q1 AS `Jan-Mar`))"
         )
-
         self.validate_identity(
             "CREATE FUNCTION add_one(x INT) RETURNS INT LANGUAGE PYTHON AS $$def add_one(x):\n  return x+1$$"
         )
-
         self.validate_identity(
             "CREATE FUNCTION add_one(x INT) RETURNS INT LANGUAGE PYTHON AS $FOO$def add_one(x):\n  return x+1$FOO$"
         )
-
-        self.validate_identity("TRUNCATE TABLE t1 PARTITION(age = 10, name = 'test', address)")
         self.validate_identity(
             "TRUNCATE TABLE t1 PARTITION(age = 10, name = 'test', city LIKE 'LA')"
         )
         self.validate_identity(
             "COPY INTO target FROM `s3://link` FILEFORMAT = AVRO VALIDATE = ALL FILES = ('file1', 'file2') FORMAT_OPTIONS ('opt1'='true', 'opt2'='test') COPY_OPTIONS ('mergeSchema'='true')"
         )
+        self.validate_identity(
+            r'SELECT r"\\foo.bar\"',
+            r"SELECT '\\\\foo.bar\\'",
+        )
 
         self.validate_all(
             "CREATE TABLE foo (x INT GENERATED ALWAYS AS (YEAR(y)))",
@@ -67,7 +68,6 @@ class TestDatabricks(Validator):
                 "teradata": "CREATE TABLE t1 AS (SELECT c FROM t2) WITH DATA",
             },
         )
-
         self.validate_all(
             "SELECT X'1A2B'",
             read={