sqlglot.dialects.snowflake

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp, generator, parser, tokens, transforms
  6from sqlglot.dialects.dialect import (
  7    Dialect,
  8    date_trunc_to_time,
  9    datestrtodate_sql,
 10    format_time_lambda,
 11    inline_array_sql,
 12    max_or_greatest,
 13    min_or_least,
 14    rename_func,
 15    timestamptrunc_sql,
 16    timestrtotime_sql,
 17    ts_or_ds_to_date_sql,
 18    var_map_sql,
 19)
 20from sqlglot.expressions import Literal
 21from sqlglot.helper import seq_get
 22from sqlglot.parser import binary_range_parser
 23from sqlglot.tokens import TokenType
 24
 25
 26def _check_int(s: str) -> bool:
 27    if s[0] in ("-", "+"):
 28        return s[1:].isdigit()
 29    return s.isdigit()
 30
 31
 32# from https://docs.snowflake.com/en/sql-reference/functions/to_timestamp.html
 33def _parse_to_timestamp(args: t.List) -> t.Union[exp.StrToTime, exp.UnixToTime]:
 34    if len(args) == 2:
 35        first_arg, second_arg = args
 36        if second_arg.is_string:
 37            # case: <string_expr> [ , <format> ]
 38            return format_time_lambda(exp.StrToTime, "snowflake")(args)
 39
 40        # case: <numeric_expr> [ , <scale> ]
 41        if second_arg.name not in ["0", "3", "9"]:
 42            raise ValueError(
 43                f"Scale for snowflake numeric timestamp is {second_arg}, but should be 0, 3, or 9"
 44            )
 45
 46        if second_arg.name == "0":
 47            timescale = exp.UnixToTime.SECONDS
 48        elif second_arg.name == "3":
 49            timescale = exp.UnixToTime.MILLIS
 50        elif second_arg.name == "9":
 51            timescale = exp.UnixToTime.MICROS
 52
 53        return exp.UnixToTime(this=first_arg, scale=timescale)
 54
 55    from sqlglot.optimizer.simplify import simplify_literals
 56
 57    # The first argument might be an expression like 40 * 365 * 86400, so we try to
 58    # reduce it using `simplify_literals` first and then check if it's a Literal.
 59    first_arg = seq_get(args, 0)
 60    if not isinstance(simplify_literals(first_arg, root=True), Literal):
 61        # case: <variant_expr>
 62        return format_time_lambda(exp.StrToTime, "snowflake", default=True)(args)
 63
 64    if first_arg.is_string:
 65        if _check_int(first_arg.this):
 66            # case: <integer>
 67            return exp.UnixToTime.from_arg_list(args)
 68
 69        # case: <date_expr>
 70        return format_time_lambda(exp.StrToTime, "snowflake", default=True)(args)
 71
 72    # case: <numeric_expr>
 73    return exp.UnixToTime.from_arg_list(args)
 74
 75
 76def _parse_object_construct(args: t.List) -> t.Union[exp.StarMap, exp.Struct]:
 77    expression = parser.parse_var_map(args)
 78
 79    if isinstance(expression, exp.StarMap):
 80        return expression
 81
 82    return exp.Struct(
 83        expressions=[
 84            t.cast(exp.Condition, k).eq(v) for k, v in zip(expression.keys, expression.values)
 85        ]
 86    )
 87
 88
 89def _parse_datediff(args: t.List) -> exp.DateDiff:
 90    return exp.DateDiff(this=seq_get(args, 2), expression=seq_get(args, 1), unit=seq_get(args, 0))
 91
 92
 93def _unix_to_time_sql(self: generator.Generator, expression: exp.UnixToTime) -> str:
 94    scale = expression.args.get("scale")
 95    timestamp = self.sql(expression, "this")
 96    if scale in [None, exp.UnixToTime.SECONDS]:
 97        return f"TO_TIMESTAMP({timestamp})"
 98    if scale == exp.UnixToTime.MILLIS:
 99        return f"TO_TIMESTAMP({timestamp}, 3)"
100    if scale == exp.UnixToTime.MICROS:
101        return f"TO_TIMESTAMP({timestamp}, 9)"
102
103    raise ValueError("Improper scale for timestamp")
104
105
106# https://docs.snowflake.com/en/sql-reference/functions/date_part.html
107# https://docs.snowflake.com/en/sql-reference/functions-date-time.html#label-supported-date-time-parts
108def _parse_date_part(self: parser.Parser) -> t.Optional[exp.Expression]:
109    this = self._parse_var() or self._parse_type()
110
111    if not this:
112        return None
113
114    self._match(TokenType.COMMA)
115    expression = self._parse_bitwise()
116
117    name = this.name.upper()
118    if name.startswith("EPOCH"):
119        if name.startswith("EPOCH_MILLISECOND"):
120            scale = 10**3
121        elif name.startswith("EPOCH_MICROSECOND"):
122            scale = 10**6
123        elif name.startswith("EPOCH_NANOSECOND"):
124            scale = 10**9
125        else:
126            scale = None
127
128        ts = self.expression(exp.Cast, this=expression, to=exp.DataType.build("TIMESTAMP"))
129        to_unix: exp.Expression = self.expression(exp.TimeToUnix, this=ts)
130
131        if scale:
132            to_unix = exp.Mul(this=to_unix, expression=exp.Literal.number(scale))
133
134        return to_unix
135
136    return self.expression(exp.Extract, this=this, expression=expression)
137
138
139# https://docs.snowflake.com/en/sql-reference/functions/div0
140def _div0_to_if(args: t.List) -> exp.If:
141    cond = exp.EQ(this=seq_get(args, 1), expression=exp.Literal.number(0))
142    true = exp.Literal.number(0)
143    false = exp.Div(this=seq_get(args, 0), expression=seq_get(args, 1))
144    return exp.If(this=cond, true=true, false=false)
145
146
147# https://docs.snowflake.com/en/sql-reference/functions/zeroifnull
148def _zeroifnull_to_if(args: t.List) -> exp.If:
149    cond = exp.Is(this=seq_get(args, 0), expression=exp.Null())
150    return exp.If(this=cond, true=exp.Literal.number(0), false=seq_get(args, 0))
151
152
153# https://docs.snowflake.com/en/sql-reference/functions/zeroifnull
154def _nullifzero_to_if(args: t.List) -> exp.If:
155    cond = exp.EQ(this=seq_get(args, 0), expression=exp.Literal.number(0))
156    return exp.If(this=cond, true=exp.Null(), false=seq_get(args, 0))
157
158
159def _datatype_sql(self: generator.Generator, expression: exp.DataType) -> str:
160    if expression.is_type("array"):
161        return "ARRAY"
162    elif expression.is_type("map"):
163        return "OBJECT"
164    return self.datatype_sql(expression)
165
166
167def _parse_convert_timezone(args: t.List) -> t.Union[exp.Anonymous, exp.AtTimeZone]:
168    if len(args) == 3:
169        return exp.Anonymous(this="CONVERT_TIMEZONE", expressions=args)
170    return exp.AtTimeZone(this=seq_get(args, 1), zone=seq_get(args, 0))
171
172
173def _parse_regexp_replace(args: t.List) -> exp.RegexpReplace:
174    regexp_replace = exp.RegexpReplace.from_arg_list(args)
175
176    if not regexp_replace.args.get("replacement"):
177        regexp_replace.set("replacement", exp.Literal.string(""))
178
179    return regexp_replace
180
181
182class Snowflake(Dialect):
183    # https://docs.snowflake.com/en/sql-reference/identifiers-syntax
184    RESOLVES_IDENTIFIERS_AS_UPPERCASE = True
185    NULL_ORDERING = "nulls_are_large"
186    TIME_FORMAT = "'YYYY-MM-DD HH24:MI:SS'"
187
188    TIME_MAPPING = {
189        "YYYY": "%Y",
190        "yyyy": "%Y",
191        "YY": "%y",
192        "yy": "%y",
193        "MMMM": "%B",
194        "mmmm": "%B",
195        "MON": "%b",
196        "mon": "%b",
197        "MM": "%m",
198        "mm": "%m",
199        "DD": "%d",
200        "dd": "%-d",
201        "DY": "%a",
202        "dy": "%w",
203        "HH24": "%H",
204        "hh24": "%H",
205        "HH12": "%I",
206        "hh12": "%I",
207        "MI": "%M",
208        "mi": "%M",
209        "SS": "%S",
210        "ss": "%S",
211        "FF": "%f",
212        "ff": "%f",
213        "FF6": "%f",
214        "ff6": "%f",
215    }
216
217    class Parser(parser.Parser):
218        IDENTIFY_PIVOT_STRINGS = True
219
220        FUNCTIONS = {
221            **parser.Parser.FUNCTIONS,
222            "ARRAYAGG": exp.ArrayAgg.from_arg_list,
223            "ARRAY_CONSTRUCT": exp.Array.from_arg_list,
224            "ARRAY_TO_STRING": exp.ArrayJoin.from_arg_list,
225            "CONVERT_TIMEZONE": _parse_convert_timezone,
226            "DATE_TRUNC": date_trunc_to_time,
227            "DATEADD": lambda args: exp.DateAdd(
228                this=seq_get(args, 2), expression=seq_get(args, 1), unit=seq_get(args, 0)
229            ),
230            "DATEDIFF": _parse_datediff,
231            "DIV0": _div0_to_if,
232            "IFF": exp.If.from_arg_list,
233            "NULLIFZERO": _nullifzero_to_if,
234            "OBJECT_CONSTRUCT": _parse_object_construct,
235            "REGEXP_REPLACE": _parse_regexp_replace,
236            "REGEXP_SUBSTR": exp.RegexpExtract.from_arg_list,
237            "RLIKE": exp.RegexpLike.from_arg_list,
238            "SQUARE": lambda args: exp.Pow(this=seq_get(args, 0), expression=exp.Literal.number(2)),
239            "TIMEDIFF": _parse_datediff,
240            "TIMESTAMPDIFF": _parse_datediff,
241            "TO_ARRAY": exp.Array.from_arg_list,
242            "TO_TIMESTAMP": _parse_to_timestamp,
243            "TO_VARCHAR": exp.ToChar.from_arg_list,
244            "ZEROIFNULL": _zeroifnull_to_if,
245        }
246
247        FUNCTION_PARSERS = {
248            **parser.Parser.FUNCTION_PARSERS,
249            "DATE_PART": _parse_date_part,
250        }
251        FUNCTION_PARSERS.pop("TRIM")
252
253        FUNC_TOKENS = {
254            *parser.Parser.FUNC_TOKENS,
255            TokenType.TABLE,
256        }
257
258        COLUMN_OPERATORS = {
259            **parser.Parser.COLUMN_OPERATORS,
260            TokenType.COLON: lambda self, this, path: self.expression(
261                exp.Bracket, this=this, expressions=[path]
262            ),
263        }
264
265        TIMESTAMPS = parser.Parser.TIMESTAMPS.copy() - {TokenType.TIME}
266
267        RANGE_PARSERS = {
268            **parser.Parser.RANGE_PARSERS,
269            TokenType.LIKE_ANY: binary_range_parser(exp.LikeAny),
270            TokenType.ILIKE_ANY: binary_range_parser(exp.ILikeAny),
271        }
272
273        ALTER_PARSERS = {
274            **parser.Parser.ALTER_PARSERS,
275            "SET": lambda self: self._parse_set(tag=self._match_text_seq("TAG")),
276            "UNSET": lambda self: self.expression(
277                exp.Set,
278                tag=self._match_text_seq("TAG"),
279                expressions=self._parse_csv(self._parse_id_var),
280                unset=True,
281            ),
282        }
283
284        def _parse_id_var(
285            self,
286            any_token: bool = True,
287            tokens: t.Optional[t.Collection[TokenType]] = None,
288        ) -> t.Optional[exp.Expression]:
289            if self._match_text_seq("IDENTIFIER", "("):
290                identifier = (
291                    super()._parse_id_var(any_token=any_token, tokens=tokens)
292                    or self._parse_string()
293                )
294                self._match_r_paren()
295                return self.expression(exp.Anonymous, this="IDENTIFIER", expressions=[identifier])
296
297            return super()._parse_id_var(any_token=any_token, tokens=tokens)
298
299    class Tokenizer(tokens.Tokenizer):
300        QUOTES = ["'"]
301        STRING_ESCAPES = ["\\", "'"]
302        HEX_STRINGS = [("x'", "'"), ("X'", "'")]
303        RAW_STRINGS = ["$$"]
304        COMMENTS = ["--", "//", ("/*", "*/")]
305
306        KEYWORDS = {
307            **tokens.Tokenizer.KEYWORDS,
308            "BYTEINT": TokenType.INT,
309            "CHAR VARYING": TokenType.VARCHAR,
310            "CHARACTER VARYING": TokenType.VARCHAR,
311            "EXCLUDE": TokenType.EXCEPT,
312            "ILIKE ANY": TokenType.ILIKE_ANY,
313            "LIKE ANY": TokenType.LIKE_ANY,
314            "MATCH_RECOGNIZE": TokenType.MATCH_RECOGNIZE,
315            "MINUS": TokenType.EXCEPT,
316            "NCHAR VARYING": TokenType.VARCHAR,
317            "PUT": TokenType.COMMAND,
318            "RENAME": TokenType.REPLACE,
319            "SAMPLE": TokenType.TABLE_SAMPLE,
320            "TIMESTAMP_LTZ": TokenType.TIMESTAMPLTZ,
321            "TIMESTAMP_NTZ": TokenType.TIMESTAMP,
322            "TIMESTAMP_TZ": TokenType.TIMESTAMPTZ,
323            "TIMESTAMPNTZ": TokenType.TIMESTAMP,
324            "TOP": TokenType.TOP,
325        }
326
327        SINGLE_TOKENS = {
328            **tokens.Tokenizer.SINGLE_TOKENS,
329            "$": TokenType.PARAMETER,
330        }
331
332        VAR_SINGLE_TOKENS = {"$"}
333
334    class Generator(generator.Generator):
335        PARAMETER_TOKEN = "$"
336        MATCHED_BY_SOURCE = False
337        SINGLE_STRING_INTERVAL = True
338        JOIN_HINTS = False
339        TABLE_HINTS = False
340        QUERY_HINTS = False
341
342        TRANSFORMS = {
343            **generator.Generator.TRANSFORMS,
344            exp.Array: inline_array_sql,
345            exp.ArrayConcat: rename_func("ARRAY_CAT"),
346            exp.ArrayJoin: rename_func("ARRAY_TO_STRING"),
347            exp.AtTimeZone: lambda self, e: self.func(
348                "CONVERT_TIMEZONE", e.args.get("zone"), e.this
349            ),
350            exp.DateAdd: lambda self, e: self.func("DATEADD", e.text("unit"), e.expression, e.this),
351            exp.DateDiff: lambda self, e: self.func(
352                "DATEDIFF", e.text("unit"), e.expression, e.this
353            ),
354            exp.DateStrToDate: datestrtodate_sql,
355            exp.DataType: _datatype_sql,
356            exp.DayOfWeek: rename_func("DAYOFWEEK"),
357            exp.Extract: rename_func("DATE_PART"),
358            exp.If: rename_func("IFF"),
359            exp.LogicalAnd: rename_func("BOOLAND_AGG"),
360            exp.LogicalOr: rename_func("BOOLOR_AGG"),
361            exp.Map: lambda self, e: var_map_sql(self, e, "OBJECT_CONSTRUCT"),
362            exp.Max: max_or_greatest,
363            exp.Min: min_or_least,
364            exp.PartitionedByProperty: lambda self, e: f"PARTITION BY {self.sql(e, 'this')}",
365            exp.Select: transforms.preprocess([transforms.eliminate_distinct_on]),
366            exp.StarMap: rename_func("OBJECT_CONSTRUCT"),
367            exp.StartsWith: rename_func("STARTSWITH"),
368            exp.StrPosition: lambda self, e: self.func(
369                "POSITION", e.args.get("substr"), e.this, e.args.get("position")
370            ),
371            exp.StrToTime: lambda self, e: f"TO_TIMESTAMP({self.sql(e, 'this')}, {self.format_time(e)})",
372            exp.Struct: lambda self, e: self.func(
373                "OBJECT_CONSTRUCT",
374                *(arg for expression in e.expressions for arg in expression.flatten()),
375            ),
376            exp.TimestampTrunc: timestamptrunc_sql,
377            exp.TimeStrToTime: timestrtotime_sql,
378            exp.TimeToStr: lambda self, e: self.func(
379                "TO_CHAR", exp.cast(e.this, "timestamp"), self.format_time(e)
380            ),
381            exp.TimeToUnix: lambda self, e: f"EXTRACT(epoch_second FROM {self.sql(e, 'this')})",
382            exp.ToChar: lambda self, e: self.function_fallback_sql(e),
383            exp.Trim: lambda self, e: self.func("TRIM", e.this, e.expression),
384            exp.TsOrDsToDate: ts_or_ds_to_date_sql("snowflake"),
385            exp.UnixToTime: _unix_to_time_sql,
386            exp.VarMap: lambda self, e: var_map_sql(self, e, "OBJECT_CONSTRUCT"),
387            exp.WeekOfYear: rename_func("WEEKOFYEAR"),
388        }
389
390        TYPE_MAPPING = {
391            **generator.Generator.TYPE_MAPPING,
392            exp.DataType.Type.TIMESTAMP: "TIMESTAMPNTZ",
393        }
394
395        STAR_MAPPING = {
396            "except": "EXCLUDE",
397            "replace": "RENAME",
398        }
399
400        PROPERTIES_LOCATION = {
401            **generator.Generator.PROPERTIES_LOCATION,
402            exp.SetProperty: exp.Properties.Location.UNSUPPORTED,
403            exp.VolatileProperty: exp.Properties.Location.UNSUPPORTED,
404        }
405
406        def regexpextract_sql(self, expression: exp.RegexpExtract) -> str:
407            # Other dialects don't support all of the following parameters, so we need to
408            # generate default values as necessary to ensure the transpilation is correct
409            group = expression.args.get("group")
410            parameters = expression.args.get("parameters") or (group and exp.Literal.string("c"))
411            occurrence = expression.args.get("occurrence") or (parameters and exp.Literal.number(1))
412            position = expression.args.get("position") or (occurrence and exp.Literal.number(1))
413
414            return self.func(
415                "REGEXP_SUBSTR",
416                expression.this,
417                expression.expression,
418                position,
419                occurrence,
420                parameters,
421                group,
422            )
423
424        def except_op(self, expression: exp.Except) -> str:
425            if not expression.args.get("distinct", False):
426                self.unsupported("EXCEPT with All is not supported in Snowflake")
427            return super().except_op(expression)
428
429        def intersect_op(self, expression: exp.Intersect) -> str:
430            if not expression.args.get("distinct", False):
431                self.unsupported("INTERSECT with All is not supported in Snowflake")
432            return super().intersect_op(expression)
433
434        def describe_sql(self, expression: exp.Describe) -> str:
435            # Default to table if kind is unknown
436            kind_value = expression.args.get("kind") or "TABLE"
437            kind = f" {kind_value}" if kind_value else ""
438            this = f" {self.sql(expression, 'this')}"
439            return f"DESCRIBE{kind}{this}"
440
441        def generatedasidentitycolumnconstraint_sql(
442            self, expression: exp.GeneratedAsIdentityColumnConstraint
443        ) -> str:
444            start = expression.args.get("start")
445            start = f" START {start}" if start else ""
446            increment = expression.args.get("increment")
447            increment = f" INCREMENT {increment}" if increment else ""
448            return f"AUTOINCREMENT{start}{increment}"