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:

Resolving patch:
diff --git a/sqlglot/dialects/postgres.py b/sqlglot/dialects/postgres.py
--- a/sqlglot/dialects/postgres.py
+++ b/sqlglot/dialects/postgres.py
@@ -232,6 +232,9 @@ class Tokenizer(tokens.Tokenizer):
         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
--- a/sqlglot/tokens.py
+++ b/sqlglot/tokens.py
@@ -504,6 +504,7 @@ def _quotes_to_format(
                 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 @@ def _quotes_to_format(
                 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 @@ def _scan_string(self, start: str) -> bool:
             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}"


Test patch:
diff --git a/tests/dialects/test_clickhouse.py b/tests/dialects/test_clickhouse.py
--- a/tests/dialects/test_clickhouse.py
+++ b/tests/dialects/test_clickhouse.py
@@ -77,6 +77,10 @@ def test_clickhouse(self):
         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
--- a/tests/dialects/test_postgres.py
+++ b/tests/dialects/test_postgres.py
@@ -33,6 +33,7 @@ def test_postgres(self):
         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)")