Feat(prql): Add support for SORT (#3297)

fool1280 · web-flow · commit a72192306c8f · 2024-04-10T20:35:17.000+03:00
* orderby

* update

* sqlglot/dialects/prql.py
diff --git a/sqlglot/dialects/prql.py b/sqlglot/dialects/prql.py
@@ -52,6 +52,7 @@ class Parser(parser.Parser):
             "INTERSECT": lambda self, query: query.intersect(
                 _select_all(self._parse_table()), distinct=False, copy=False
             ),
+            "SORT": lambda self, query: self._parse_order_by(query),
         }
 
         def _parse_statement(self) -> t.Optional[exp.Expression]:
@@ -103,6 +104,13 @@ def _parse_take(self, query: exp.Query) -> t.Optional[exp.Query]:
             num = self._parse_number()  # TODO: TAKE for ranges a..b
             return query.limit(num) if num else None
 
+        def _parse_order_by(self, query: exp.Select) -> t.Optional[exp.Query]:
+            l_brace = self._match(TokenType.L_BRACE)
+            expressions = self._parse_csv(self._parse_ordered)
+            if l_brace and not self._match(TokenType.R_BRACE):
+                self.raise_error("Expecting }")
+            return query.order_by(self.expression(exp.Order, expressions=expressions), copy=False)
+
         def _parse_expression(self) -> t.Optional[exp.Expression]:
             if self._next and self._next.token_type == TokenType.ALIAS:
                 alias = self._parse_id_var(True)
diff --git a/tests/dialects/test_prql.py b/tests/dialects/test_prql.py
@@ -37,6 +37,14 @@ def test_prql(self):
             "FROM x FILTER (age > 25 || age < 22) FILTER age > 26 FILTER age < 27",
             "SELECT * FROM x WHERE ((age > 25 OR age < 22) AND age > 26) AND age < 27",
         )
+        self.validate_identity(
+            "FROM x SORT age",
+            "SELECT * FROM x ORDER BY age",
+        )
+        self.validate_identity(
+            "FROM x SORT {age, name}",
+            "SELECT * FROM x ORDER BY age, name",
+        )
         self.validate_identity("FROM x APPEND y", "SELECT * FROM x UNION ALL SELECT * FROM y")
         self.validate_identity("FROM x REMOVE y", "SELECT * FROM x EXCEPT ALL SELECT * FROM y")
         self.validate_identity(
