sqlglot.parsers.prql

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp, parser
  6from sqlglot.helper import seq_get
  7from sqlglot.tokens import TokenType
  8from collections.abc import Collection
  9
 10
 11def _select_all(table: exp.Expr | None) -> exp.Select | None:
 12    return exp.select("*").from_(table, copy=False) if table else None
 13
 14
 15def _resolve_projection(s: exp.Expr, projections: dict[str, exp.Expr]) -> exp.Expr:
 16    if isinstance(s, exp.Column) and s.name in projections:
 17        return projections[s.name].copy()
 18    return s
 19
 20
 21class PRQLParser(parser.Parser):
 22    CONJUNCTION = {
 23        **parser.Parser.CONJUNCTION,
 24        TokenType.DAMP: exp.And,
 25    }
 26
 27    DISJUNCTION = {
 28        **parser.Parser.DISJUNCTION,
 29        TokenType.DPIPE: exp.Or,
 30    }
 31
 32    TRANSFORM_PARSERS: t.ClassVar = {
 33        "DERIVE": lambda self, query: self._parse_selection(query),
 34        "SELECT": lambda self, query: self._parse_selection(query, append=False),
 35        "TAKE": lambda self, query: self._parse_take(query),
 36        "FILTER": lambda self, query: query.where(self._parse_disjunction()),
 37        "APPEND": lambda self, query: query.union(
 38            _select_all(self._parse_table()), distinct=False, copy=False
 39        ),
 40        "REMOVE": lambda self, query: query.except_(
 41            _select_all(self._parse_table()), distinct=False, copy=False
 42        ),
 43        "INTERSECT": lambda self, query: query.intersect(
 44            _select_all(self._parse_table()), distinct=False, copy=False
 45        ),
 46        "SORT": lambda self, query: self._parse_order_by(query),
 47        "AGGREGATE": lambda self, query: self._parse_selection(
 48            query, parse_method=self._parse_aggregate, append=False
 49        ),
 50    }
 51
 52    FUNCTIONS = {
 53        **parser.Parser.FUNCTIONS,
 54        "AVERAGE": exp.Avg.from_arg_list,
 55        "SUM": lambda args: exp.func("COALESCE", exp.Sum(this=seq_get(args, 0)), 0),
 56    }
 57
 58    def _parse_equality(self, parse_is: bool = True) -> exp.Expr | None:
 59        eq = super()._parse_equality(parse_is=parse_is)
 60        if not isinstance(eq, (exp.EQ, exp.NEQ)):
 61            return eq
 62
 63        # https://prql-lang.org/book/reference/spec/null.html
 64        if isinstance(eq.expression, exp.Null):
 65            is_exp = exp.Is(this=eq.this, expression=eq.expression)
 66            return is_exp if isinstance(eq, exp.EQ) else exp.Not(this=is_exp)
 67        if isinstance(eq.this, exp.Null):
 68            is_exp = exp.Is(this=eq.expression, expression=eq.this)
 69            return is_exp if isinstance(eq, exp.EQ) else exp.Not(this=is_exp)
 70        return eq
 71
 72    def _parse_statement(self) -> exp.Expr | None:
 73        expression = self._parse_expression()
 74        expression = expression if expression else self._parse_query()
 75        return expression
 76
 77    def _parse_query(self) -> exp.Query | None:
 78        from_ = self._parse_from()
 79
 80        if not from_:
 81            return None
 82
 83        query: exp.Query = exp.select("*").from_(from_, copy=False)
 84
 85        while self._match_texts(self.TRANSFORM_PARSERS):
 86            query = self.TRANSFORM_PARSERS[self._prev.text.upper()](self, query)
 87
 88        return query
 89
 90    def _parse_selection(
 91        self,
 92        query: exp.Query,
 93        parse_method: t.Callable[..., exp.Expr | None] | None = None,
 94        append: bool = True,
 95    ) -> exp.Query:
 96        parse_method = parse_method if parse_method else self._parse_expression
 97        if self._match(TokenType.L_BRACE):
 98            selects = self._parse_csv(parse_method)
 99
100            if not self._match(TokenType.R_BRACE, expression=query):
101                self.raise_error("Expecting }")
102        else:
103            expression = parse_method()
104            selects = [expression] if expression else []
105
106        projections: dict[str, exp.Expr] = {
107            select.alias_or_name: select.this if isinstance(select, exp.Alias) else select
108            for select in query.selects
109        }
110
111        resolved = [
112            select.transform(_resolve_projection, projections=projections, copy=False)
113            for select in selects
114        ]
115
116        return query.select(*resolved, append=append, copy=False)
117
118    def _parse_take(self, query: exp.Query) -> exp.Query | None:
119        num = self._parse_number()  # TODO: TAKE for ranges a..b
120        return query.limit(num) if num else None
121
122    def _parse_ordered(self, parse_method: t.Callable | None = None) -> exp.Ordered | None:
123        asc = self._match(TokenType.PLUS)
124        desc = self._match(TokenType.DASH) or (asc and False)
125        term = term = super()._parse_ordered(parse_method=parse_method)
126        if term and desc:
127            term.set("desc", True)
128            term.set("nulls_first", False)
129        return term
130
131    def _parse_order_by(self, query: exp.Select) -> exp.Query | None:
132        l_brace = self._match(TokenType.L_BRACE)
133        expressions = self._parse_csv(self._parse_ordered)
134        if l_brace and not self._match(TokenType.R_BRACE):
135            self.raise_error("Expecting }")
136        return query.order_by(self.expression(exp.Order(expressions=expressions)), copy=False)
137
138    def _parse_aggregate(self) -> exp.Expr | None:
139        alias = None
140        if self._next.token_type == TokenType.ALIAS:
141            alias = self._parse_id_var(any_token=True)
142            self._match(TokenType.ALIAS)
143
144        name = self._curr.text.upper()
145        func_builder = self.FUNCTIONS.get(name)
146        if func_builder:
147            self._advance()
148            args = self._parse_column()
149            func = func_builder([args])
150        else:
151            self.raise_error(f"Unsupported aggregation function {name}")
152        if alias:
153            return self.expression(exp.Alias(this=func, alias=alias))
154        return func
155
156    def _parse_expression(self) -> exp.Expr | None:
157        if self._next.token_type == TokenType.ALIAS:
158            alias = self._parse_id_var(True)
159            self._match(TokenType.ALIAS)
160            return self.expression(exp.Alias(this=self._parse_assignment(), alias=alias))
161        return self._parse_assignment()
162
163    def _parse_table(
164        self,
165        schema: bool = False,
166        joins: bool = False,
167        alias_tokens: Collection[TokenType] | None = None,
168        parse_bracket: bool = False,
169        is_db_reference: bool = False,
170        parse_partition: bool = False,
171        consume_pipe: bool = False,
172    ) -> exp.Expr | None:
173        return self._parse_table_parts()
174
175    def _parse_from(
176        self,
177        joins: bool = False,
178        skip_from_token: bool = False,
179        consume_pipe: bool = False,
180    ) -> exp.From | None:
181        if not skip_from_token and not self._match(TokenType.FROM):
182            return None
183
184        comments = self._prev_comments
185        return self.expression(
186            exp.From(this=self._parse_table(joins=joins)),
187            comments=comments,
188        )