sqlglot.parsers.redshift

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp
  6from sqlglot.helper import seq_get
  7from sqlglot.parsers.postgres import PostgresParser
  8from sqlglot.parser import build_convert_timezone
  9from sqlglot.tokens import TokenType
 10from sqlglot.dialects.dialect import map_date_part
 11from builtins import type as Type
 12
 13if t.TYPE_CHECKING:
 14    from sqlglot._typing import E
 15    from collections.abc import Collection
 16
 17
 18def _build_date_delta(expr_type: Type[E]) -> t.Callable[[list], E]:
 19    def _builder(args: list) -> E:
 20        expr = expr_type(
 21            this=seq_get(args, 2),
 22            expression=seq_get(args, 1),
 23            unit=map_date_part(seq_get(args, 0)),
 24        )
 25        if expr_type is exp.TsOrDsAdd:
 26            expr.set("return_type", exp.DType.TIMESTAMP.into_expr())
 27
 28        return expr
 29
 30    return _builder
 31
 32
 33class RedshiftParser(PostgresParser):
 34    FUNCTIONS = {
 35        **{k: v for k, v in PostgresParser.FUNCTIONS.items() if k != "GET_BIT"},
 36        "ADD_MONTHS": lambda args: exp.TsOrDsAdd(
 37            this=seq_get(args, 0),
 38            expression=seq_get(args, 1),
 39            unit=exp.var("month"),
 40            return_type=exp.DType.TIMESTAMP.into_expr(),
 41        ),
 42        "CONVERT_TIMEZONE": lambda args: build_convert_timezone(args, "UTC"),
 43        "DATEADD": _build_date_delta(exp.TsOrDsAdd),
 44        "DATE_ADD": _build_date_delta(exp.TsOrDsAdd),
 45        "DATEDIFF": _build_date_delta(exp.TsOrDsDiff),
 46        "DATE_DIFF": _build_date_delta(exp.TsOrDsDiff),
 47        "GETDATE": exp.CurrentTimestamp.from_arg_list,
 48        "LISTAGG": exp.GroupConcat.from_arg_list,
 49        "REGEXP_SUBSTR": lambda args: exp.RegexpExtract(
 50            this=seq_get(args, 0),
 51            expression=seq_get(args, 1),
 52            position=seq_get(args, 2),
 53            occurrence=seq_get(args, 3),
 54            parameters=seq_get(args, 4),
 55        ),
 56        "SPLIT_TO_ARRAY": lambda args: exp.StringToArray(
 57            this=seq_get(args, 0), expression=seq_get(args, 1) or exp.Literal.string(",")
 58        ),
 59        "STRTOL": exp.FromBase.from_arg_list,
 60        "TEXTLEN": exp.Length.from_arg_list,
 61    }
 62
 63    NO_PAREN_FUNCTION_PARSERS = {
 64        **PostgresParser.NO_PAREN_FUNCTION_PARSERS,
 65        "APPROXIMATE": lambda self: self._parse_approximate_count(),
 66        "SYSDATE": lambda self: self.expression(exp.CurrentTimestamp(sysdate=True)),
 67    }
 68
 69    SUPPORTS_IMPLICIT_UNNEST = True
 70
 71    def _parse_table(
 72        self,
 73        schema: bool = False,
 74        joins: bool = False,
 75        alias_tokens: Collection[TokenType] | None = None,
 76        parse_bracket: bool = False,
 77        is_db_reference: bool = False,
 78        parse_partition: bool = False,
 79        consume_pipe: bool = False,
 80    ) -> exp.Expr | None:
 81        # Redshift supports UNPIVOTing SUPER objects, e.g. `UNPIVOT foo.obj[0] AS val AT attr`
 82        unpivot = self._match(TokenType.UNPIVOT)
 83        table = super()._parse_table(
 84            schema=schema,
 85            joins=joins,
 86            alias_tokens=alias_tokens,
 87            parse_bracket=parse_bracket,
 88            is_db_reference=is_db_reference,
 89        )
 90
 91        return self.expression(exp.Pivot(this=table, unpivot=True)) if unpivot else table
 92
 93    def _parse_convert(self, strict: bool, safe: bool | None = None) -> exp.Expr | None:
 94        to = self._parse_types()
 95        self._match(TokenType.COMMA)
 96        this = self._parse_bitwise()
 97        return self.expression(exp.TryCast(this=this, to=to, safe=safe))
 98
 99    def _parse_approximate_count(self) -> exp.ApproxDistinct | None:
100        index = self._index - 1
101        func = self._parse_function()
102
103        if isinstance(func, exp.Count) and isinstance(func.this, exp.Distinct):
104            return self.expression(exp.ApproxDistinct(this=seq_get(func.this.expressions, 0)))
105        self._retreat(index)
106        return None
107
108    def _parse_projections(self) -> tuple[list[exp.Expr], list[exp.Expr] | None]:
109        projections, _ = super()._parse_projections()
110        if self._prev.text.upper() == "EXCLUDE" and self._curr:
111            self._retreat(self._index - 1)
112
113        # EXCLUDE clause always comes at the end of the projection list and applies to it as a whole
114        exclude = (
115            self._parse_wrapped_csv(self._parse_expression, optional=True)
116            if self._match_text_seq("EXCLUDE")
117            else []
118        )
119
120        if (
121            exclude
122            and isinstance(expr := projections[-1], exp.Alias)
123            and expr.alias.upper() == "EXCLUDE"
124        ):
125            projections[-1] = expr.this.pop()
126
127        return projections, exclude