Repository: tobymao/sqlglot

Issue:
Count distinct syntax error while transpiling to trino sql
**Before you file an issue**
- Make sure you specify the "read" dialect eg. parse_one(sql, read="spark")
- Check if the issue still exists on main 

**Fully reproducible code snippet**
```python
import sqlglot

sql = """
with tbl as (
	select 1 as id, 'eggy' as name
	union all
	select null as id, 'jake' as name
)
select count(distinct id, name) as cnt from tbl
"""
sqlglot.transpile(sql, read="hive", write="trino")[0]
```

**Details**
The result of code sinppet is ⬇️ 
```
"WITH tbl AS (SELECT 1 AS id, 'eggy' AS name UNION ALL SELECT NULL AS id, 'jake' AS name) SELECT COUNT(DISTINCT id, name) AS cnt FROM tbl"
```
and I run this sql in the trino cli, it would raise error below, as I konwn that trino/presto dose not support syntax `count(distinct col_1, col_2, ...)` but for wrapping multi-columns into **row** type like `count(distinct (col_1, col_2, ...))`.
<img width="808" alt="image" src="https://github.com/tobymao/sqlglot/assets/60967034/82f76a45-1519-401a-9133-b749317ba0d5">




Comments:

Patch:
diff --git a/sqlglot/dialects/duckdb.py b/sqlglot/dialects/duckdb.py
index d7ba729c..e61ac4fd 100644
--- a/sqlglot/dialects/duckdb.py
+++ b/sqlglot/dialects/duckdb.py
@@ -333,6 +333,7 @@ class DuckDB(Dialect):
         IGNORE_NULLS_IN_FUNC = True
         JSON_PATH_BRACKETED_KEY_SUPPORTED = False
         SUPPORTS_CREATE_TABLE_LIKE = False
+        MULTI_ARG_DISTINCT = False
 
         TRANSFORMS = {
             **generator.Generator.TRANSFORMS,
diff --git a/sqlglot/dialects/postgres.py b/sqlglot/dialects/postgres.py
index 0404c78f..68e2c6de 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,
@@ -381,6 +384,7 @@ class Postgres(Dialect):
         JSON_TYPE_REQUIRED_FOR_EXTRACTION = True
         SUPPORTS_UNLOGGED_TABLES = True
         LIKE_PROPERTY_INSIDE_SCHEMA = True
+        MULTI_ARG_DISTINCT = False
 
         SUPPORTED_JSON_PATH_PARTS = {
             exp.JSONPathKey,
diff --git a/sqlglot/dialects/presto.py b/sqlglot/dialects/presto.py
index 8691192b..609103e5 100644
--- a/sqlglot/dialects/presto.py
+++ b/sqlglot/dialects/presto.py
@@ -292,6 +292,7 @@ class Presto(Dialect):
         LIMIT_ONLY_LITERALS = True
         SUPPORTS_SINGLE_ARG_CONCAT = False
         LIKE_PROPERTY_INSIDE_SCHEMA = True
+        MULTI_ARG_DISTINCT = False
 
         PROPERTIES_LOCATION = {
             **generator.Generator.PROPERTIES_LOCATION,
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index 81af56d8..eff8aaa2 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -296,6 +296,10 @@ class Generator(metaclass=_Generator):
     # Whether or not the LikeProperty needs to be specified inside of the schema clause
     LIKE_PROPERTY_INSIDE_SCHEMA = False
 
+    # Whether or not DISTINCT can be followed by multiple args in an AggFunc. If not, it will be
+    # transpiled into a series of CASE-WHEN-ELSE, ultimately using a tuple conseisting of the args
+    MULTI_ARG_DISTINCT = True
+
     # Whether or not the JSON extraction operators expect a value of type JSON
     JSON_TYPE_REQUIRED_FOR_EXTRACTION = False
 
@@ -2837,6 +2841,13 @@ class Generator(metaclass=_Generator):
 
     def distinct_sql(self, expression: exp.Distinct) -> str:
         this = self.expressions(expression, flat=True)
+
+        if not self.MULTI_ARG_DISTINCT and len(expression.expressions) > 1:
+            case = exp.case()
+            for arg in expression.expressions:
+                case = case.when(arg.is_(exp.null()), exp.null())
+            this = self.sql(case.else_(f"({this})"))
+
         this = f" {this}" if this else ""
 
         on = self.sql(expression, "on")
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)")
diff --git a/tests/dialects/test_spark.py b/tests/dialects/test_spark.py
index a02a735c..75bb91af 100644
--- a/tests/dialects/test_spark.py
+++ b/tests/dialects/test_spark.py
@@ -277,6 +277,21 @@ TBLPROPERTIES (
             "SELECT STR_TO_MAP('a:1,b:2,c:3', ',', ':')",
         )
 
+        self.validate_all(
+            "WITH tbl AS (SELECT 1 AS id, 'eggy' AS name UNION ALL SELECT NULL AS id, 'jake' AS name) SELECT COUNT(DISTINCT id, name) AS cnt FROM tbl",
+            write={
+                "clickhouse": "WITH tbl AS (SELECT 1 AS id, 'eggy' AS name UNION ALL SELECT NULL AS id, 'jake' AS name) SELECT COUNT(DISTINCT id, name) AS cnt FROM tbl",
+                "databricks": "WITH tbl AS (SELECT 1 AS id, 'eggy' AS name UNION ALL SELECT NULL AS id, 'jake' AS name) SELECT COUNT(DISTINCT id, name) AS cnt FROM tbl",
+                "doris": "WITH tbl AS (SELECT 1 AS id, 'eggy' AS name UNION ALL SELECT NULL AS id, 'jake' AS name) SELECT COUNT(DISTINCT id, name) AS cnt FROM tbl",
+                "duckdb": "WITH tbl AS (SELECT 1 AS id, 'eggy' AS name UNION ALL SELECT NULL AS id, 'jake' AS name) SELECT COUNT(DISTINCT CASE WHEN id IS NULL THEN NULL WHEN name IS NULL THEN NULL ELSE (id, name) END) AS cnt FROM tbl",
+                "hive": "WITH tbl AS (SELECT 1 AS id, 'eggy' AS name UNION ALL SELECT NULL AS id, 'jake' AS name) SELECT COUNT(DISTINCT id, name) AS cnt FROM tbl",
+                "mysql": "WITH tbl AS (SELECT 1 AS id, 'eggy' AS name UNION ALL SELECT NULL AS id, 'jake' AS name) SELECT COUNT(DISTINCT id, name) AS cnt FROM tbl",
+                "postgres": "WITH tbl AS (SELECT 1 AS id, 'eggy' AS name UNION ALL SELECT NULL AS id, 'jake' AS name) SELECT COUNT(DISTINCT CASE WHEN id IS NULL THEN NULL WHEN name IS NULL THEN NULL ELSE (id, name) END) AS cnt FROM tbl",
+                "presto": "WITH tbl AS (SELECT 1 AS id, 'eggy' AS name UNION ALL SELECT NULL AS id, 'jake' AS name) SELECT COUNT(DISTINCT CASE WHEN id IS NULL THEN NULL WHEN name IS NULL THEN NULL ELSE (id, name) END) AS cnt FROM tbl",
+                "snowflake": "WITH tbl AS (SELECT 1 AS id, 'eggy' AS name UNION ALL SELECT NULL AS id, 'jake' AS name) SELECT COUNT(DISTINCT id, name) AS cnt FROM tbl",
+                "spark": "WITH tbl AS (SELECT 1 AS id, 'eggy' AS name UNION ALL SELECT NULL AS id, 'jake' AS name) SELECT COUNT(DISTINCT id, name) AS cnt FROM tbl",
+            },
+        )
         self.validate_all(
             "SELECT TO_UTC_TIMESTAMP('2016-08-31', 'Asia/Seoul')",
             write={