Repository: tobymao/sqlglot

Issue:
Parameter `$1` for `postgres` is not supported
In version `v18.8.0`, heredoc strings support was added, and `"$": TokenType.PARAMETER` was replaced with `"$": TokenType.HEREDOC_STRING`. Since then, sql query with a parameter like `$1` cannot be parsed.

Here is the original PR [Feat!: add support for heredoc strings (Postgres, ClickHouse) #2328](https://github.com/tobymao/sqlglot/pull/2328)

**Fully reproducible code snippet**

```
import sqlglot
from sqlglot.optimizer.annotate_types import annotate_types
from sqlglot.optimizer.qualify import qualify

schema = {"t": {"x": "text"}}
sql = "select x from t where $1::text = 'ok'"
expression = sqlglot.parse_one(sql, dialect="postgres")
print(expression.sql(dialect="postgres"))
qualified_expr = qualify(expression, schema=schema, dialect="postgres")
annotated_expr = annotate_types(qualified_expr, schema=schema)
print(annotated_expr.selects[0].type)
```
Exception was raised. Expected output is `TEXT`.
```
Traceback (most recent call last):
  File "/Users/eric/venv/lib/python3.10/site-packages/sqlglot/tokens.py", line 836, in tokenize
    self._scan()
  File "/Users/eric/venv/lib/python3.10/site-packages/sqlglot/tokens.py", line 859, in _scan
    self._scan_keywords()
  File "/Users/eric/venv/lib/python3.10/site-packages/sqlglot/tokens.py", line 994, in _scan_keywords
    if self._scan_string(word):
  File "/Users/eric/venv/lib/python3.10/site-packages/sqlglot/tokens.py", line 1136, in _scan_string
    tag = "" if self._char == end else self._extract_string(end)
  File "/Users/eric/venv/lib/python3.10/site-packages/sqlglot/tokens.py", line 1203, in _extract_string
    raise TokenError(f"Missing {delimiter} from {self._line}:{self._start}")
sqlglot.errors.TokenError: Missing $ from 1:22

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/eric/tmp/test.py", line 20, in <module>
    expression = sqlglot.parse_one(sql, dialect="postgres")
  File "/Users/eric/venv/lib/python3.10/site-packages/sqlglot/__init__.py", line 125, in parse_one
    result = dialect.parse(sql, **opts)
  File "/Users/eric/venv/lib/python3.10/site-packages/sqlglot/dialects/dialect.py", line 304, in parse
    return self.parser(**opts).parse(self.tokenize(sql), sql)
  File "/Users/eric/venv/lib/python3.10/site-packages/sqlglot/dialects/dialect.py", line 318, in tokenize
    return self.tokenizer.tokenize(sql)
  File "/Users/eric/venv/lib/python3.10/site-packages/sqlglot/tokens.py", line 841, in tokenize
    raise TokenError(f"Error tokenizing '{context}'") from e
sqlglot.errors.TokenError: Error tokenizing 'select x from t where $1::text = 'ok'
```




Discussion:
georgesittas: Addressed feedback:

- Python side now mimics Rust side w.r.t. the `HEREDOC_STRING_ALTERNATIVE` logic
- Factored out the `is_alphabetic() ||  ... == '_'` logic so it resembles the Python side more.

Patch:
diff --git a/sqlglot/dialects/postgres.py b/sqlglot/dialects/postgres.py
index 0404c78f..6a6825e4 100644
--- a/sqlglot/dialects/postgres.py
+++ b/sqlglot/dialects/postgres.py
@@ -232,6 +232,9 @@ class Postgres(Dialect):
         BYTE_STRINGS = [("e'", "'"), ("E'", "'")]
         HEREDOC_STRINGS = ["$"]
 
+        HEREDOC_TAG_IS_IDENTIFIER = True
+        HEREDOC_STRING_ALTERNATIVE = TokenType.PARAMETER
+
         KEYWORDS = {
             **tokens.Tokenizer.KEYWORDS,
             "~~": TokenType.LIKE,
diff --git a/sqlglot/tokens.py b/sqlglot/tokens.py
index 87a49240..b0649578 100644
--- a/sqlglot/tokens.py
+++ b/sqlglot/tokens.py
@@ -504,6 +504,7 @@ class _Tokenizer(type):
                 command_prefix_tokens={
                     _TOKEN_TYPE_TO_INDEX[v] for v in klass.COMMAND_PREFIX_TOKENS
                 },
+                heredoc_tag_is_identifier=klass.HEREDOC_TAG_IS_IDENTIFIER,
             )
             token_types = RsTokenTypeSettings(
                 bit_string=_TOKEN_TYPE_TO_INDEX[TokenType.BIT_STRING],
@@ -517,6 +518,7 @@ class _Tokenizer(type):
                 semicolon=_TOKEN_TYPE_TO_INDEX[TokenType.SEMICOLON],
                 string=_TOKEN_TYPE_TO_INDEX[TokenType.STRING],
                 var=_TOKEN_TYPE_TO_INDEX[TokenType.VAR],
+                heredoc_string_alternative=_TOKEN_TYPE_TO_INDEX[klass.HEREDOC_STRING_ALTERNATIVE],
             )
             klass._RS_TOKENIZER = RsTokenizer(settings, token_types)
         else:
@@ -573,6 +575,12 @@ class Tokenizer(metaclass=_Tokenizer):
     STRING_ESCAPES = ["'"]
     VAR_SINGLE_TOKENS: t.Set[str] = set()
 
+    # Whether or not the heredoc tags follow the same lexical rules as unquoted identifiers
+    HEREDOC_TAG_IS_IDENTIFIER = False
+
+    # Token that we'll generate as a fallback if the heredoc prefix doesn't correspond to a heredoc
+    HEREDOC_STRING_ALTERNATIVE = TokenType.VAR
+
     # Autofilled
     _COMMENTS: t.Dict[str, str] = {}
     _FORMAT_STRINGS: t.Dict[str, t.Tuple[str, TokenType]] = {}
@@ -1249,6 +1257,18 @@ class Tokenizer(metaclass=_Tokenizer):
             elif token_type == TokenType.BIT_STRING:
                 base = 2
             elif token_type == TokenType.HEREDOC_STRING:
+                if (
+                    self.HEREDOC_TAG_IS_IDENTIFIER
+                    and not self._peek.isidentifier()
+                    and not self._peek == end
+                ):
+                    if self.HEREDOC_STRING_ALTERNATIVE != token_type.VAR:
+                        self._add(self.HEREDOC_STRING_ALTERNATIVE)
+                    else:
+                        self._scan_var()
+
+                    return True
+
                 self._advance()
                 tag = "" if self._char == end else self._extract_string(end)
                 end = f"{start}{tag}{end}"
diff --git a/sqlglotrs/src/settings.rs b/sqlglotrs/src/settings.rs
index 32575c63..c6e76a70 100644
--- a/sqlglotrs/src/settings.rs
+++ b/sqlglotrs/src/settings.rs
@@ -17,6 +17,7 @@ pub struct TokenTypeSettings {
     pub semicolon: TokenType,
     pub string: TokenType,
     pub var: TokenType,
+    pub heredoc_string_alternative: TokenType,
 }
 
 #[pymethods]
@@ -34,6 +35,7 @@ impl TokenTypeSettings {
         semicolon: TokenType,
         string: TokenType,
         var: TokenType,
+        heredoc_string_alternative: TokenType,
     ) -> Self {
         TokenTypeSettings {
             bit_string,
@@ -47,6 +49,7 @@ impl TokenTypeSettings {
             semicolon,
             string,
             var,
+            heredoc_string_alternative,
         }
     }
 }
@@ -69,6 +72,7 @@ pub struct TokenizerSettings {
     pub var_single_tokens: HashSet<char>,
     pub commands: HashSet<TokenType>,
     pub command_prefix_tokens: HashSet<TokenType>,
+    pub heredoc_tag_is_identifier: bool,
 }
 
 #[pymethods]
@@ -90,6 +94,7 @@ impl TokenizerSettings {
         var_single_tokens: HashSet<String>,
         commands: HashSet<TokenType>,
         command_prefix_tokens: HashSet<TokenType>,
+        heredoc_tag_is_identifier: bool,
     ) -> Self {
         let to_char = |v: &String| {
             if v.len() == 1 {
@@ -138,6 +143,7 @@ impl TokenizerSettings {
             var_single_tokens: var_single_tokens_native,
             commands,
             command_prefix_tokens,
+            heredoc_tag_is_identifier,
         }
     }
 }
diff --git a/sqlglotrs/src/tokenizer.rs b/sqlglotrs/src/tokenizer.rs
index 920a5b5c..94a8b084 100644
--- a/sqlglotrs/src/tokenizer.rs
+++ b/sqlglotrs/src/tokenizer.rs
@@ -399,6 +399,19 @@ impl<'a> TokenizerState<'a> {
             } else if *token_type == self.token_types.bit_string {
                 (Some(2), *token_type, end.clone())
             } else if *token_type == self.token_types.heredoc_string {
+                if self.settings.heredoc_tag_is_identifier
+                    && !self.is_identifier(self.peek_char)
+                    && self.peek_char.to_string() != *end
+                {
+                    if self.token_types.heredoc_string_alternative != self.token_types.var {
+                        self.add(self.token_types.heredoc_string_alternative, None)?
+                    } else {
+                        self.scan_var()?
+                    };
+
+                    return Ok(true)
+                };
+
                 self.advance(1)?;
                 let tag = if self.current_char.to_string() == *end {
                     String::from("")
@@ -469,7 +482,7 @@ impl<'a> TokenizerState<'a> {
             } else if self.peek_char.to_ascii_uppercase() == 'E' && scientific == 0 {
                 scientific += 1;
                 self.advance(1)?;
-            } else if self.peek_char.is_alphabetic() || self.peek_char == '_' {
+            } else if self.is_identifier(self.peek_char) {
                 let number_text = self.text();
                 let mut literal = String::from("");
 
@@ -643,6 +656,10 @@ impl<'a> TokenizerState<'a> {
         Ok(text)
     }
 
+    fn is_identifier(&mut self, name: char) -> bool {
+        name.is_alphabetic() || name == '_'
+    }
+
     fn extract_value(&mut self) -> Result<String, TokenizerError> {
         loop {
             if !self.peek_char.is_whitespace()


Test patch:
diff --git a/tests/dialects/test_clickhouse.py b/tests/dialects/test_clickhouse.py
index dd73bae3..7351f6a0 100644
--- a/tests/dialects/test_clickhouse.py
+++ b/tests/dialects/test_clickhouse.py
@@ -77,6 +77,10 @@ class TestClickhouse(Validator):
         self.validate_identity("""SELECT JSONExtractString('{"x": {"y": 1}}', 'x', 'y')""")
         self.validate_identity("SELECT * FROM table LIMIT 1 BY a, b")
         self.validate_identity("SELECT * FROM table LIMIT 2 OFFSET 1 BY a, b")
+        self.validate_identity(
+            "SELECT $1$foo$1$",
+            "SELECT 'foo'",
+        )
         self.validate_identity(
             "SELECT * FROM table LIMIT 1, 2 BY a, b",
             "SELECT * FROM table LIMIT 2 OFFSET 1 BY a, b",
diff --git a/tests/dialects/test_postgres.py b/tests/dialects/test_postgres.py
index 9c4246e5..61421e5f 100644
--- a/tests/dialects/test_postgres.py
+++ b/tests/dialects/test_postgres.py
@@ -33,6 +33,7 @@ class TestPostgres(Validator):
         self.assertIsInstance(expr, exp.AlterTable)
         self.assertEqual(expr.sql(dialect="postgres"), alter_table_only)
 
+        self.validate_identity("SELECT x FROM t WHERE CAST($1 AS TEXT) = 'ok'")
         self.validate_identity("SELECT * FROM t TABLESAMPLE SYSTEM (50) REPEATABLE (55)")
         self.validate_identity("x @@ y")
         self.validate_identity("CAST(x AS MONEY)")