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) -> 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 = {
 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) -> exp.Expr | None:
 59        eq = self._parse_comparison()
 60        while self._match_set(self.EQUALITY):
 61            comments = self._prev_comments
 62            eq = self.expression(
 63                self.EQUALITY[self._prev.token_type](this=eq, expression=self._parse_comparison()),
 64                comments=comments,
 65            )
 66        if not isinstance(eq, (exp.EQ, exp.NEQ)):
 67            return eq
 68
 69        # https://prql-lang.org/book/reference/spec/null.html
 70        if isinstance(eq.expression, exp.Null):
 71            is_exp = exp.Is(this=eq.this, expression=eq.expression)
 72            return is_exp if isinstance(eq, exp.EQ) else exp.Not(this=is_exp)
 73        if isinstance(eq.this, exp.Null):
 74            is_exp = exp.Is(this=eq.expression, expression=eq.this)
 75            return is_exp if isinstance(eq, exp.EQ) else exp.Not(this=is_exp)
 76        return eq
 77
 78    def _parse_statement(self) -> exp.Expr | None:
 79        expression = self._parse_expression()
 80        expression = expression if expression else self._parse_query()
 81        return expression
 82
 83    def _parse_query(self) -> exp.Query | None:
 84        from_ = self._parse_from()
 85
 86        if not from_:
 87            return None
 88
 89        query: exp.Query = exp.select("*").from_(from_, copy=False)
 90
 91        while self._match_texts(self.TRANSFORM_PARSERS):
 92            query = self.TRANSFORM_PARSERS[self._prev.text.upper()](self, query)
 93
 94        return query
 95
 96    def _parse_selection(
 97        self,
 98        query: exp.Query,
 99        parse_method: t.Callable | None = None,
100        append: bool = True,
101    ) -> exp.Query:
102        parse_method = parse_method if parse_method else self._parse_expression
103        if self._match(TokenType.L_BRACE):
104            selects = self._parse_csv(parse_method)
105
106            if not self._match(TokenType.R_BRACE, expression=query):
107                self.raise_error("Expecting }")
108        else:
109            expression = parse_method()
110            selects = [expression] if expression else []
111
112        projections = {
113            select.alias_or_name: select.this if isinstance(select, exp.Alias) else select
114            for select in query.selects
115        }
116
117        resolved = [
118            select.transform(_resolve_projection, projections=projections, copy=False)
119            for select in selects
120        ]
121
122        return query.select(*resolved, append=append, copy=False)
123
124    def _parse_take(self, query: exp.Query) -> exp.Query | None:
125        num = self._parse_number()  # TODO: TAKE for ranges a..b
126        return query.limit(num) if num else None
127
128    def _parse_ordered(self, parse_method: t.Callable | None = None) -> exp.Ordered | None:
129        asc = self._match(TokenType.PLUS)
130        desc = self._match(TokenType.DASH) or (asc and False)
131        term = term = super()._parse_ordered(parse_method=parse_method)
132        if term and desc:
133            term.set("desc", True)
134            term.set("nulls_first", False)
135        return term
136
137    def _parse_order_by(self, query: exp.Select) -> exp.Query | None:
138        l_brace = self._match(TokenType.L_BRACE)
139        expressions = self._parse_csv(self._parse_ordered)
140        if l_brace and not self._match(TokenType.R_BRACE):
141            self.raise_error("Expecting }")
142        return query.order_by(self.expression(exp.Order(expressions=expressions)), copy=False)
143
144    def _parse_aggregate(self) -> exp.Expr | None:
145        alias = None
146        if self._next.token_type == TokenType.ALIAS:
147            alias = self._parse_id_var(any_token=True)
148            self._match(TokenType.ALIAS)
149
150        name = self._curr.text.upper()
151        func_builder = self.FUNCTIONS.get(name)
152        if func_builder:
153            self._advance()
154            args = self._parse_column()
155            func = func_builder([args])
156        else:
157            self.raise_error(f"Unsupported aggregation function {name}")
158        if alias:
159            return self.expression(exp.Alias(this=func, alias=alias))
160        return func
161
162    def _parse_expression(self) -> exp.Expr | None:
163        if self._next.token_type == TokenType.ALIAS:
164            alias = self._parse_id_var(True)
165            self._match(TokenType.ALIAS)
166            return self.expression(exp.Alias(this=self._parse_assignment(), alias=alias))
167        return self._parse_assignment()
168
169    def _parse_table(
170        self,
171        schema: bool = False,
172        joins: bool = False,
173        alias_tokens: Collection[TokenType] | None = None,
174        parse_bracket: bool = False,
175        is_db_reference: bool = False,
176        parse_partition: bool = False,
177        consume_pipe: bool = False,
178    ) -> exp.Expr | None:
179        return self._parse_table_parts()
180
181    def _parse_from(
182        self,
183        joins: bool = False,
184        skip_from_token: bool = False,
185        consume_pipe: bool = False,
186    ) -> exp.From | None:
187        if not skip_from_token and not self._match(TokenType.FROM):
188            return None
189
190        comments = self._prev_comments
191        return self.expression(
192            exp.From(this=self._parse_table(joins=joins)),
193            comments=comments,
194        )