sqlglot.dialects.prql

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

19    class Tokenizer(tokens.Tokenizer):
20        IDENTIFIERS = ["`"]
21        QUOTES = ["'", '"']
22
23        SINGLE_TOKENS = {
24            **tokens.Tokenizer.SINGLE_TOKENS,
25            "=": TokenType.ALIAS,
26            "'": TokenType.QUOTE,
27            '"': TokenType.QUOTE,
28            "`": TokenType.IDENTIFIER,
29            "#": TokenType.COMMENT,
30        }
31
32        KEYWORDS = {
33            **tokens.Tokenizer.KEYWORDS,
34        }

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