Repository: tobymao/sqlglot

Issue:
T-SQL COLLATE argument being interpreted as a column
SQLGlot versions: tested with 25.9.0 and 25.8.1


When parsing a code block, I'm finding that the parsed query is returning the COLLATE collation argument as a column.

Example SQL:
```
select
    sql_text
    , some_id
from
    my_catalog.my_schema.my_table a
    inner join my_catalog.my_schema.my_other_table b
        on a.id = b.id
WHERE
    MySpecialColumn LIKE '%test%' COLLATE LATIN1_GENERAL_BIN

```

Given the above sql as variable "sql" - 

```
parsed_query = parse_one(sql=sql, dialect="tsql")
x = list(parsed_query.find_all(exp.Column))
for c in x:
    print(c)
```

I'm seeing "LATIN1_GENERAL_BIN" being emitted as a column, which surprised me. I wasn't sure if this was the desired behavior or not.


Discussion:

Patch:
diff --git a/sqlglot/dialects/dialect.py b/sqlglot/dialects/dialect.py
index d2eef99d..b3df9a71 100644
--- a/sqlglot/dialects/dialect.py
+++ b/sqlglot/dialects/dialect.py
@@ -647,6 +647,9 @@ class Dialect(metaclass=_Dialect):
         exp.GenerateDateArray: lambda self, e: self._annotate_with_type(
             e, exp.DataType.build("ARRAY<DATE>")
         ),
+        exp.GenerateTimestampArray: lambda self, e: self._annotate_with_type(
+            e, exp.DataType.build("ARRAY<TIMESTAMP>")
+        ),
         exp.If: lambda self, e: self._annotate_by_args(e, "true", "false"),
         exp.Interval: lambda self, e: self._annotate_with_type(e, exp.DataType.Type.INTERVAL),
         exp.Least: lambda self, e: self._annotate_by_args(e, "expressions"),
diff --git a/sqlglot/dialects/duckdb.py b/sqlglot/dialects/duckdb.py
index 70edc1e5..ddd252cc 100644
--- a/sqlglot/dialects/duckdb.py
+++ b/sqlglot/dialects/duckdb.py
@@ -209,17 +209,41 @@ def _arrow_json_extract_sql(self: DuckDB.Generator, expression: JSON_EXTRACT_TYP
     return arrow_sql
 
 
-def _implicit_date_cast(arg: t.Optional[exp.Expression]) -> t.Optional[exp.Expression]:
-    return exp.cast(arg, exp.DataType.Type.DATE) if isinstance(arg, exp.Literal) else arg
+def _implicit_datetime_cast(
+    arg: t.Optional[exp.Expression], type: exp.DataType.Type = exp.DataType.Type.DATE
+) -> t.Optional[exp.Expression]:
+    return exp.cast(arg, type) if isinstance(arg, exp.Literal) else arg
 
 
 def _date_diff_sql(self: DuckDB.Generator, expression: exp.DateDiff) -> str:
-    this = _implicit_date_cast(expression.this)
-    expr = _implicit_date_cast(expression.expression)
+    this = _implicit_datetime_cast(expression.this)
+    expr = _implicit_datetime_cast(expression.expression)
 
     return self.func("DATE_DIFF", unit_to_str(expression), expr, this)
 
 
+def _generate_datetime_array_sql(
+    self: DuckDB.Generator, expression: t.Union[exp.GenerateDateArray, exp.GenerateTimestampArray]
+) -> str:
+    is_generate_date_array = isinstance(expression, exp.GenerateDateArray)
+
+    type = exp.DataType.Type.DATE if is_generate_date_array else exp.DataType.Type.TIMESTAMP
+    start = _implicit_datetime_cast(expression.args.get("start"), type=type)
+    end = _implicit_datetime_cast(expression.args.get("end"), type=type)
+
+    # BQ's GENERATE_DATE_ARRAY & GENERATE_TIMESTAMP_ARRAY are transformed to DuckDB'S GENERATE_SERIES
+    gen_series: t.Union[exp.GenerateSeries, exp.Cast] = exp.GenerateSeries(
+        start=start, end=end, step=expression.args.get("interval")
+    )
+
+    if is_generate_date_array:
+        # The GENERATE_SERIES result type is TIMESTAMP array, so to match BQ's semantics for
+        # GENERATE_DATE_ARRAY we must cast it back to DATE array
+        gen_series = exp.cast(gen_series, exp.DataType.build("ARRAY<DATE>"))
+
+    return self.sql(gen_series)
+
+
 class DuckDB(Dialect):
     NULL_ORDERING = "nulls_are_last"
     SUPPORTS_USER_DEFINED_TYPES = False
@@ -476,6 +500,8 @@ class DuckDB(Dialect):
             exp.DiToDate: lambda self,
             e: f"CAST(STRPTIME(CAST({self.sql(e, 'this')} AS TEXT), {DuckDB.DATEINT_FORMAT}) AS DATE)",
             exp.Encode: lambda self, e: encode_decode_sql(self, e, "ENCODE", replace=False),
+            exp.GenerateDateArray: _generate_datetime_array_sql,
+            exp.GenerateTimestampArray: _generate_datetime_array_sql,
             exp.Explode: rename_func("UNNEST"),
             exp.IntDiv: lambda self, e: self.binary(e, "//"),
             exp.IsInf: rename_func("ISINF"),
@@ -855,15 +881,3 @@ class DuckDB(Dialect):
                 return self.func("STRUCT_PACK", kv_sql)
 
             return self.func("STRUCT_INSERT", this, kv_sql)
-
-        def generatedatearray_sql(self, expression: exp.GenerateDateArray) -> str:
-            start = _implicit_date_cast(expression.args.get("start"))
-            end = _implicit_date_cast(expression.args.get("end"))
-
-            # BQ's GENERATE_DATE_ARRAY is transformed to DuckDB'S GENERATE_SERIES
-            gen_series = exp.GenerateSeries(
-                start=start, end=end, step=expression.args.get("interval")
-            )
-
-            # The result is TIMESTAMP array, so to match BQ's semantics we must cast it back to DATE array
-            return self.sql(exp.cast(gen_series, exp.DataType.build("ARRAY<DATE>")))
diff --git a/sqlglot/dialects/mysql.py b/sqlglot/dialects/mysql.py
index 2fa355e5..145a7d1c 100644
--- a/sqlglot/dialects/mysql.py
+++ b/sqlglot/dialects/mysql.py
@@ -302,6 +302,9 @@ class MySQL(Dialect):
 
         FUNCTIONS = {
             **parser.Parser.FUNCTIONS,
+            "CONVERT_TZ": lambda args: exp.ConvertTimezone(
+                source_tz=seq_get(args, 1), target_tz=seq_get(args, 2), timestamp=seq_get(args, 0)
+            ),
             "DATE": lambda args: exp.TsOrDsToDate(this=seq_get(args, 0)),
             "DATE_ADD": build_date_delta_with_interval(exp.DateAdd),
             "DATE_FORMAT": build_formatted_time(exp.TimeToStr, "mysql"),
@@ -1213,3 +1216,10 @@ class MySQL(Dialect):
             dateadd = build_date_delta_with_interval(exp.DateAdd)([start_ts, interval])
 
             return self.sql(dateadd)
+
+        def converttimezone_sql(self, expression: exp.ConvertTimezone) -> str:
+            from_tz = expression.args.get("source_tz")
+            to_tz = expression.args.get("target_tz")
+            dt = expression.args.get("timestamp")
+
+            return self.func("CONVERT_TZ", dt, from_tz, to_tz)
diff --git a/sqlglot/dialects/redshift.py b/sqlglot/dialects/redshift.py
index 5e1fa432..b39a9a50 100644
--- a/sqlglot/dialects/redshift.py
+++ b/sqlglot/dialects/redshift.py
@@ -17,6 +17,7 @@ from sqlglot.dialects.dialect import (
 from sqlglot.dialects.postgres import Postgres
 from sqlglot.helper import seq_get
 from sqlglot.tokens import TokenType
+from sqlglot.parser import build_convert_timezone
 
 if t.TYPE_CHECKING:
     from sqlglot._typing import E
@@ -63,6 +64,7 @@ class Redshift(Postgres):
                 unit=exp.var("month"),
                 return_type=exp.DataType.build("TIMESTAMP"),
             ),
+            "CONVERT_TIMEZONE": lambda args: build_convert_timezone(args, "UTC"),
             "DATEADD": _build_date_delta(exp.TsOrDsAdd),
             "DATE_ADD": _build_date_delta(exp.TsOrDsAdd),
             "DATEDIFF": _build_date_delta(exp.TsOrDsDiff),
@@ -155,6 +157,7 @@ class Redshift(Postgres):
         HEX_FUNC = "TO_HEX"
         PARSE_JSON_NAME = "JSON_PARSE"
         ARRAY_CONCAT_IS_VAR_LEN = False
+        SUPPORTS_CONVERT_TIMEZONE = True
 
         # Redshift doesn't have `WITH` as part of their with_properties so we remove it
         WITH_PROPERTIES_PREFIX = " "
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
index 81db3190..58fdb620 100644
--- a/sqlglot/dialects/snowflake.py
+++ b/sqlglot/dialects/snowflake.py
@@ -122,12 +122,6 @@ def _regexpilike_sql(self: Snowflake.Generator, expression: exp.RegexpILike) ->
     )
 
 
-def _build_convert_timezone(args: t.List) -> t.Union[exp.Anonymous, exp.AtTimeZone]:
-    if len(args) == 3:
-        return exp.Anonymous(this="CONVERT_TIMEZONE", expressions=args)
-    return exp.AtTimeZone(this=seq_get(args, 1), zone=seq_get(args, 0))
-
-
 def _build_regexp_replace(args: t.List) -> exp.RegexpReplace:
     regexp_replace = exp.RegexpReplace.from_arg_list(args)
 
@@ -268,7 +262,6 @@ class Snowflake(Dialect):
             "BITXOR": binary_from_function(exp.BitwiseXor),
             "BIT_XOR": binary_from_function(exp.BitwiseXor),
             "BOOLXOR": binary_from_function(exp.Xor),
-            "CONVERT_TIMEZONE": _build_convert_timezone,
             "DATE": _build_datetime("DATE", exp.DataType.Type.DATE),
             "DATE_TRUNC": _date_trunc_to_time,
             "DATEADD": _build_date_time_add(exp.DateAdd),
@@ -694,6 +687,7 @@ class Snowflake(Dialect):
         STAR_EXCEPT = "EXCLUDE"
         SUPPORTS_EXPLODING_PROJECTIONS = False
         ARRAY_CONCAT_IS_VAR_LEN = False
+        SUPPORTS_CONVERT_TIMEZONE = True
 
         TRANSFORMS = {
             **generator.Generator.TRANSFORMS,
diff --git a/sqlglot/dialects/spark.py b/sqlglot/dialects/spark.py
index 35944d52..410c12e0 100644
--- a/sqlglot/dialects/spark.py
+++ b/sqlglot/dialects/spark.py
@@ -132,6 +132,7 @@ class Spark(Spark2):
     class Generator(Spark2.Generator):
         SUPPORTS_TO_NUMBER = True
         PAD_FILL_PATTERN_IS_REQUIRED = False
+        SUPPORTS_CONVERT_TIMEZONE = True
 
         TYPE_MAPPING = {
             **Spark2.Generator.TYPE_MAPPING,
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 6f0bcc2a..4e21f15d 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -4845,6 +4845,10 @@ class Convert(Func):
     arg_types = {"this": True, "expression": True, "style": False}
 
 
+class ConvertTimezone(Func):
+    arg_types = {"source_tz": False, "target_tz": True, "timestamp": True}
+
+
 class GenerateSeries(Func):
     arg_types = {"start": True, "end": True, "step": False, "is_end_exclusive": False}
 
@@ -5340,10 +5344,16 @@ class GapFill(Func):
     }
 
 
+# https://cloud.google.com/bigquery/docs/reference/standard-sql/array_functions#generate_date_array
 class GenerateDateArray(Func):
     arg_types = {"start": True, "end": True, "interval": False}
 
 
+# https://cloud.google.com/bigquery/docs/reference/standard-sql/array_functions#generate_timestamp_array
+class GenerateTimestampArray(Func):
+    arg_types = {"start": True, "end": True, "interval": True}
+
+
 class Greatest(Func):
     arg_types = {"this": True, "expressions": False}
     is_var_len_args = True
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index 04b805ee..d5aff88e 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -380,6 +380,9 @@ class Generator(metaclass=_Generator):
     # Whether ARRAY_CONCAT can be generated with varlen args or if it should be reduced to 2-arg version
     ARRAY_CONCAT_IS_VAR_LEN = True
 
+    # Whether CONVERT_TIMEZONE() is supported; if not, it will be generated as exp.AtTimeZone
+    SUPPORTS_CONVERT_TIMEZONE = False
+
     # The name to generate for the JSONPath expression. If `None`, only `this` will be generated
     PARSE_JSON_NAME: t.Optional[str] = "PARSE_JSON"
 
@@ -4083,3 +4086,20 @@ class Generator(metaclass=_Generator):
             rhs = self.expressions(expression)
 
         return self.func(name, expression.this, rhs)
+
+    def converttimezone_sql(self, expression: exp.ConvertTimezone) -> str:
+        if self.SUPPORTS_CONVERT_TIMEZONE:
+            return self.function_fallback_sql(expression)
+
+        source_tz = expression.args.get("source_tz")
+        target_tz = expression.args.get("target_tz")
+        timestamp = expression.args.get("timestamp")
+
+        if source_tz and timestamp:
+            timestamp = exp.AtTimeZone(
+                this=exp.cast(timestamp, exp.DataType.Type.TIMESTAMPNTZ), zone=source_tz
+            )
+
+        expr = exp.AtTimeZone(this=timestamp, zone=target_tz)
+
+        return self.sql(expr)
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index 3a0f529a..548002ac 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -128,6 +128,18 @@ def build_array_constructor(
     return array_exp
 
 
+def build_convert_timezone(
+    args: t.List, default_source_tz: t.Optional[str] = None
+) -> t.Union[exp.ConvertTimezone, exp.Anonymous]:
+    if len(args) == 2:
+        source_tz = exp.Literal.string(default_source_tz) if default_source_tz else None
+        return exp.ConvertTimezone(
+            source_tz=source_tz, target_tz=seq_get(args, 0), timestamp=seq_get(args, 1)
+        )
+
+    return exp.ConvertTimezone.from_arg_list(args)
+
+
 class _Parser(type):
     def __new__(cls, clsname, bases, attrs):
         klass = super().__new__(cls, clsname, bases, attrs)
@@ -166,6 +178,7 @@ class Parser(metaclass=_Parser):
             safe=not dialect.STRICT_STRING_CONCAT,
             coalesce=dialect.CONCAT_COALESCE,
         ),
+        "CONVERT_TIMEZONE": build_convert_timezone,
         "DATE_TO_DATE_STR": lambda args: exp.Cast(
             this=seq_get(args, 0),
             to=exp.DataType(this=exp.DataType.Type.TEXT),
@@ -4369,7 +4382,26 @@ class Parser(metaclass=_Parser):
         return this
 
     def _parse_term(self) -> t.Optional[exp.Expression]:
-        return self._parse_tokens(self._parse_factor, self.TERM)
+        this = self._parse_factor()
+
+        while self._match_set(self.TERM):
+            klass = self.TERM[self._prev.token_type]
+            comments = self._prev_comments
+            expression = self._parse_factor()
+
+            this = self.expression(klass, this=this, comments=comments, expression=expression)
+
+            if isinstance(this, exp.Collate):
+                expr = this.expression
+
+                # Preserve collations such as pg_catalog."default" (Postgres) as columns, otherwise
+                # fallback to Identifier / Var
+                if isinstance(expr, exp.Column) and len(expr.parts) == 1:
+                    ident = expr.this
+                    if isinstance(ident, exp.Identifier):
+                        this.set("expression", ident if ident.quoted else exp.var(ident.name))
+
+        return this
 
     def _parse_factor(self) -> t.Optional[exp.Expression]:
         parse_method = self._parse_exponent if self.EXPONENT else self._parse_unary


Test patch:
diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
index d32f1bef..909981b6 100644
--- a/tests/dialects/test_bigquery.py
+++ b/tests/dialects/test_bigquery.py
@@ -1446,6 +1446,13 @@ WHERE
                 "bigquery": "SELECT GENERATE_DATE_ARRAY('2016-10-05', '2016-10-08', INTERVAL '1' MONTH)",
             },
         )
+        self.validate_all(
+            "SELECT GENERATE_TIMESTAMP_ARRAY('2016-10-05 00:00:00', '2016-10-07 00:00:00', INTERVAL '1' DAY)",
+            write={
+                "duckdb": "SELECT GENERATE_SERIES(CAST('2016-10-05 00:00:00' AS TIMESTAMP), CAST('2016-10-07 00:00:00' AS TIMESTAMP), INTERVAL '1' DAY)",
+                "bigquery": "SELECT GENERATE_TIMESTAMP_ARRAY('2016-10-05 00:00:00', '2016-10-07 00:00:00', INTERVAL '1' DAY)",
+            },
+        )
 
     def test_errors(self):
         with self.assertRaises(TokenError):
diff --git a/tests/dialects/test_redshift.py b/tests/dialects/test_redshift.py
index 0359c67e..9117aa95 100644
--- a/tests/dialects/test_redshift.py
+++ b/tests/dialects/test_redshift.py
@@ -437,6 +437,11 @@ ORDER BY
         self.validate_identity("SELECT ARRAY(1, 2, 3)")
         self.validate_identity("SELECT ARRAY[1, 2, 3]")
 
+        self.validate_identity(
+            """SELECT CONVERT_TIMEZONE('America/New_York', '2024-08-06 09:10:00.000')""",
+            """SELECT CONVERT_TIMEZONE('UTC', 'America/New_York', '2024-08-06 09:10:00.000')""",
+        )
+
     def test_values(self):
         # Test crazy-sized VALUES clause to UNION ALL conversion to ensure we don't get RecursionError
         values = [str(v) for v in range(0, 10000)]
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
index a84d82a7..0062aa4d 100644
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -869,6 +869,28 @@ WHERE
             """SELECT CAST(['foo'] AS VARIANT)[0]""",
         )
 
+        self.validate_all(
+            "SELECT CONVERT_TIMEZONE('America/New_York', '2024-08-06 09:10:00.000')",
+            write={
+                "snowflake": "SELECT CONVERT_TIMEZONE('America/New_York', '2024-08-06 09:10:00.000')",
+                "spark": "SELECT CONVERT_TIMEZONE('America/New_York', '2024-08-06 09:10:00.000')",
+                "databricks": "SELECT CONVERT_TIMEZONE('America/New_York', '2024-08-06 09:10:00.000')",
+                "redshift": "SELECT CONVERT_TIMEZONE('America/New_York', '2024-08-06 09:10:00.000')",
+            },
+        )
+
+        self.validate_all(
+            "SELECT CONVERT_TIMEZONE('America/Los_Angeles', 'America/New_York', '2024-08-06 09:10:00.000')",
+            write={
+                "snowflake": "SELECT CONVERT_TIMEZONE('America/Los_Angeles', 'America/New_York', '2024-08-06 09:10:00.000')",
+                "spark": "SELECT CONVERT_TIMEZONE('America/Los_Angeles', 'America/New_York', '2024-08-06 09:10:00.000')",
+                "databricks": "SELECT CONVERT_TIMEZONE('America/Los_Angeles', 'America/New_York', '2024-08-06 09:10:00.000')",
+                "redshift": "SELECT CONVERT_TIMEZONE('America/Los_Angeles', 'America/New_York', '2024-08-06 09:10:00.000')",
+                "mysql": "SELECT CONVERT_TZ('2024-08-06 09:10:00.000', 'America/Los_Angeles', 'America/New_York')",
+                "duckdb": "SELECT CAST('2024-08-06 09:10:00.000' AS TIMESTAMP) AT TIME ZONE 'America/Los_Angeles' AT TIME ZONE 'America/New_York'",
+            },
+        )
+
     def test_null_treatment(self):
         self.validate_all(
             r"SELECT FIRST_VALUE(TABLE1.COLUMN1) OVER (PARTITION BY RANDOM_COLUMN1, RANDOM_COLUMN2 ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) AS MY_ALIAS FROM TABLE1",
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index 31542432..d84d2974 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -1212,6 +1212,19 @@ FROM READ_CSV('tests/fixtures/optimizer/tpc-h/nation.csv.gz', 'delimiter', '|')
             exp.DataType.build("date"),
         )
 
+        self.assertEqual(
+            annotate_types(
+                optimizer.qualify.qualify(
+                    parse_one(
+                        "SELECT x FROM UNNEST(GENERATE_TIMESTAMP_ARRAY('2016-10-05 00:00:00', '2016-10-06 02:00:00', interval 1 day)) AS x"
+                    )
+                )
+            )
+            .selects[0]
+            .type,
+            exp.DataType.build("timestamp"),
+        )
+
     def test_map_annotation(self):
         # ToMap annotation
         expression = annotate_types(parse_one("SELECT MAP {'x': 1}", read="duckdb"))
diff --git a/tests/test_parser.py b/tests/test_parser.py
index f360b431..039631ae 100644
--- a/tests/test_parser.py
+++ b/tests/test_parser.py
@@ -903,3 +903,18 @@ class TestParser(unittest.TestCase):
 
     def test_parse_prop_eq(self):
         self.assertIsInstance(parse_one("x(a := b and c)").expressions[0], exp.PropertyEQ)
+
+    def test_collate(self):
+        collates = [
+            ('pg_catalog."default"', exp.Column),
+            ('"en_DE"', exp.Identifier),
+            ("LATIN1_GENERAL_BIN", exp.Var),
+            ("'en'", exp.Literal),
+        ]
+
+        for collate_pair in collates:
+            collate_node = parse_one(
+                f"""SELECT * FROM t WHERE foo LIKE '%bar%' COLLATE {collate_pair[0]}"""
+            ).find(exp.Collate)
+            self.assertIsInstance(collate_node, exp.Collate)
+            self.assertIsInstance(collate_node.expression, collate_pair[1])