sqlglot.parsers.trino

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