sqlglot.parsers.trino

  1from __future__ import annotations
  2
  3import re
  4import typing as t
  5
  6from sqlglot import exp, parser
  7from sqlglot.helper import ensure_list
  8from sqlglot.parsers.presto import PrestoParser, _build_approx_percentile
  9from sqlglot.tokens import TokenType
 10
 11FRACTIONAL_SECONDS_RE = re.compile(r":\d+\.(\d+)")
 12
 13
 14class TrinoParser(PrestoParser):
 15    TYPE_LITERAL_PARSERS: t.ClassVar = {
 16        **PrestoParser.TYPE_LITERAL_PARSERS,
 17        **dict.fromkeys(
 18            (
 19                exp.DType.TIMESTAMP,
 20                exp.DType.TIMESTAMPTZ,
 21                exp.DType.TIME,
 22                exp.DType.TIMETZ,
 23            ),
 24            lambda self, this, to: self._parse_temporal_literal(this, to),
 25        ),
 26    }
 27    NO_PAREN_FUNCTIONS = {
 28        **PrestoParser.NO_PAREN_FUNCTIONS,
 29        TokenType.CURRENT_CATALOG: exp.CurrentCatalog,
 30    }
 31
 32    FUNCTIONS = {
 33        **PrestoParser.FUNCTIONS,
 34        "APPROX_PERCENTILE": lambda args: (
 35            exp.ApproxQuantile(this=args[0], weight=args[1], quantile=args[2])
 36            if len(args) == 3
 37            else _build_approx_percentile(args)
 38        ),
 39        "VERSION": exp.CurrentVersion.from_arg_list,
 40    }
 41
 42    FUNCTION_PARSERS = {
 43        **PrestoParser.FUNCTION_PARSERS,
 44        "TRIM": lambda self: self._parse_trim(),
 45        "JSON_QUERY": lambda self: self._parse_json_query(),
 46        "JSON_VALUE": lambda self: self._parse_json_value(),
 47        "LISTAGG": lambda self: self._parse_string_agg(),
 48    }
 49
 50    JSON_QUERY_OPTIONS: t.ClassVar[parser.OPTIONS_TYPE] = {
 51        **dict.fromkeys(
 52            ("WITH", "WITHOUT"),
 53            (
 54                ("WRAPPER"),
 55                ("ARRAY", "WRAPPER"),
 56                ("CONDITIONAL", "WRAPPER"),
 57                ("CONDITIONAL", "ARRAY", "WRAPPED"),
 58                ("UNCONDITIONAL", "WRAPPER"),
 59                ("UNCONDITIONAL", "ARRAY", "WRAPPER"),
 60            ),
 61        ),
 62    }
 63
 64    def _parse_json_query_quote(self) -> exp.JSONExtractQuote | None:
 65        if not (self._match_text_seq("KEEP", "QUOTES") or self._match_text_seq("OMIT", "QUOTES")):
 66            return None
 67
 68        return self.expression(
 69            exp.JSONExtractQuote(
 70                option=self._tokens[self._index - 2].text.upper(),
 71                scalar=self._match_text_seq("ON", "SCALAR", "STRING"),
 72            )
 73        )
 74
 75    def _parse_json_query(self) -> exp.JSONExtract:
 76        return self.expression(
 77            exp.JSONExtract(
 78                this=self._parse_bitwise(),
 79                expression=self._match(TokenType.COMMA) and self._parse_bitwise(),
 80                option=self._parse_var_from_options(self.JSON_QUERY_OPTIONS, raise_unmatched=False),
 81                json_query=True,
 82                quote=self._parse_json_query_quote(),
 83                on_condition=self._parse_on_condition(),
 84            )
 85        )
 86
 87    def _parse_cte(self) -> exp.CTE | exp.FunctionSpecification | None:
 88        # A `WITH` clause entry that starts with `FUNCTION <name>` is an inline SQL UDF
 89        # specification (https://trino.io/docs/current/udf/sql.html), as opposed to a
 90        # CTE that happens to be named "function", e.g. `WITH function AS (SELECT 1)`
 91        if (
 92            self._match(TokenType.FUNCTION, advance=False)
 93            and self._next
 94            and self._next.token_type in self.ID_VAR_TOKENS
 95        ):
 96            self._advance()
 97            return self._parse_function_specification()
 98
 99        return super()._parse_cte()
100
101    def _parse_function_specification(self) -> exp.FunctionSpecification:
102        this = self._parse_user_defined_function(kind=TokenType.FUNCTION)
103
104        # Collected separately (rather than one _parse_properties() call) so the
105        # generator can place `WITH (...)` in its own bracketed clause at the end,
106        # instead of it blending into the bare characteristics list.
107        characteristics = []
108        properties = []
109
110        while True:
111            if self._match(TokenType.WITH):
112                properties.extend(self._parse_wrapped_csv(self._parse_key_value_property))
113                continue
114
115            characteristic = self._parse_property()
116            if not characteristic:
117                break
118
119            characteristics.extend(ensure_list(characteristic))
120
121        return self.expression(
122            exp.FunctionSpecification(
123                this=this,
124                characteristics=self.expression(exp.Properties(expressions=characteristics))
125                if characteristics
126                else None,
127                properties=self.expression(exp.Properties(expressions=properties))
128                if properties
129                else None,
130                expression=self._parse_routine_statement(),
131            )
132        )
133
134    def _parse_routine_statements(self, *terminators: str) -> list[exp.Expr]:
135        # Unlike _parse_block(), stops on any of `terminators` even when tokens
136        # follow (Trino's own END is always followed by the enclosing query, left
137        # for the caller), matching them as text rather than just TokenType.END, so
138        # this same chunk-continuation loop also serves IF/ELSEIF/ELSE.
139        statements: list[exp.Expr] = []
140
141        while not self._match_texts(terminators):
142            if not self._curr:
143                if self._chunk_index >= len(self._chunks):
144                    self.raise_error("Unexpected end of routine body")
145                    break
146
147                self._advance_chunk()
148            elif not self._match(TokenType.SEMICOLON):
149                statement = self._parse_routine_statement()
150                if not statement:
151                    break
152
153                statements.append(statement)
154
155        return statements
156
157    def _parse_routine_block(self) -> exp.Block:
158        self._match(TokenType.BEGIN)
159        statements = self._parse_routine_statements("END")
160        statements.append(exp.EndStatement())
161
162        return self.expression(exp.Block(expressions=statements, begin=True))
163
164    def _parse_routine_if(self) -> exp.IfBlock:
165        # https://trino.io/docs/current/udf/sql/if.html
166        # ELSEIF chains nest into `false` rather than a flat list, reusing the
167        # existing binary IfBlock instead of a new N-way expression; each loop
168        # iteration links the next branch into the previous one's `false` slot.
169        def parse_branch() -> exp.IfBlock:
170            condition = self._parse_disjunction()
171            self._match_text_seq("THEN")
172            true = self.expression(
173                exp.Block(expressions=self._parse_routine_statements("ELSEIF", "ELSE", "END"))
174            )
175            return self.expression(exp.IfBlock(this=condition, true=true))
176
177        this = tail = parse_branch()
178        while self._prev.text.upper() == "ELSEIF":
179            node = parse_branch()
180            tail.set("false", node)
181            tail = node
182
183        if self._prev.text.upper() == "ELSE":
184            tail.set(
185                "false",
186                self.expression(exp.Block(expressions=self._parse_routine_statements("END"))),
187            )
188
189        self._match_text_seq("IF")
190        return this
191
192    def _parse_routine_case(self) -> exp.CaseStatement:
193        # https://trino.io/docs/current/udf/sql/case.html
194        this = self._parse_disjunction()
195
196        def parse_branch() -> exp.If:
197            condition = self._parse_disjunction()
198            self._match_text_seq("THEN")
199            true = self.expression(
200                exp.Block(expressions=self._parse_routine_statements("WHEN", "ELSE", "END"))
201            )
202            return self.expression(exp.If(this=condition, true=true))
203
204        ifs = []
205        self._match_text_seq("WHEN")
206        while self._prev.text.upper() == "WHEN":
207            ifs.append(parse_branch())
208
209        default = None
210        if self._prev.text.upper() == "ELSE":
211            default = self.expression(exp.Block(expressions=self._parse_routine_statements("END")))
212
213        self._match_text_seq("CASE")
214        return self.expression(exp.CaseStatement(this=this, ifs=ifs, default=default))
215
216    def _parse_routine_while(self, label: exp.Expr | None = None) -> exp.WhileBlock:
217        # https://trino.io/docs/current/udf/sql/while.html
218        condition = self._parse_disjunction()
219        self._match_text_seq("DO")
220        body = self.expression(exp.Block(expressions=self._parse_routine_statements("END")))
221        self._match_text_seq("WHILE")
222        return self.expression(exp.WhileBlock(this=condition, body=body, label=label))
223
224    def _parse_routine_loop(self, label: exp.Expr | None = None) -> exp.LoopBlock:
225        # https://trino.io/docs/current/udf/sql/loop.html
226        body = self.expression(exp.Block(expressions=self._parse_routine_statements("END")))
227        self._match_text_seq("LOOP")
228        return self.expression(exp.LoopBlock(body=body, label=label))
229
230    def _parse_routine_repeat(self, label: exp.Expr | None = None) -> exp.RepeatBlock:
231        # https://trino.io/docs/current/udf/sql/repeat.html - unlike WHILE/LOOP,
232        # the body is evaluated at least once, with UNTIL <condition> tested after
233        body = self.expression(exp.Block(expressions=self._parse_routine_statements("UNTIL")))
234        until = self._parse_disjunction()
235        self._match_text_seq("END", "REPEAT")
236        return self.expression(exp.RepeatBlock(body=body, until=until, label=label))
237
238    def _parse_routine_statement(self) -> exp.Expr | None:
239        # An optional `label :` can precede WHILE, LOOP, or REPEAT to name the block for ITERATE/LEAVE.
240        # Any non-reserved keyword (SET, IF, ITERATE etc) is a valid label name, so the colon lookahead
241        # has to run first, before the statement keywords are matched.
242        label = None
243        if self._next and self._next.token_type == TokenType.COLON:
244            label = self._parse_id_var()
245            self._match(TokenType.COLON)
246
247        if self._match(TokenType.BEGIN, advance=False):
248            return self._parse_routine_block()
249
250        if self._match_text_seq("RETURN"):
251            return self.expression(exp.Return(this=self._parse_disjunction()))
252
253        if self._match_text_seq("IF"):
254            return self._parse_routine_if()
255
256        if self._match_text_seq("CASE"):
257            return self._parse_routine_case()
258
259        if self._match(TokenType.DECLARE):
260            return self._parse_declare()
261
262        if self._match(TokenType.SET):
263            return self._parse_set()
264
265        if self._match_text_seq("ITERATE"):
266            return self.expression(exp.Iterate(this=self._parse_id_var()))
267
268        if self._match_text_seq("LEAVE"):
269            return self.expression(exp.Leave(this=self._parse_id_var()))
270
271        if self._match_text_seq("WHILE"):
272            return self._parse_routine_while(label=label)
273
274        if self._match_text_seq("LOOP"):
275            return self._parse_routine_loop(label=label)
276
277        if self._match_text_seq("REPEAT"):
278            return self._parse_routine_repeat(label=label)
279
280        self.raise_error("Expected routine statement")
281        return None
282
283    def _parse_temporal_literal(self, this: exp.Literal, data_type: exp.DataType) -> exp.Cast:
284        if not data_type.expressions:
285            match = FRACTIONAL_SECONDS_RE.search(this.name)
286            data_type.set(
287                "expressions",
288                [exp.DataTypeParam(this=exp.Literal.number(len(match.group(1)) if match else 0))],
289            )
290
291        return self.expression(exp.Cast(this=this, to=data_type))