Support GROUP BY GROUPING SETS in PostgreSQL

Author: nene

src/cst/Select.ts

@@ -37,6 +37,7 @@ export type AllSelectNodes =
   | GroupByClause
   | GroupByRollup
   | GroupByCube
+  | GroupByGroupingSets
   | HavingClause
   | WindowClause
   | QualifyClause
@@ -266,10 +267,12 @@ export interface GroupByClause extends BaseNode {
   type: "group_by_clause";
   groupByKw: [Keyword<"GROUP">, Keyword<"BY">];
   distinctKw?: Keyword<"ALL" | "DISTINCT">; // PostgreSQL
-  columns: ListExpr<Expr | GroupByRollup | GroupByCube>;
+  columns: ListExpr<GroupingElement>;
   withRollupKw?: [Keyword<"WITH">, Keyword<"ROLLUP">]; // MySQL, MariaDB
 }
 
+type GroupingElement = Expr | GroupByRollup | GroupByCube | GroupByGroupingSets;
+
 // BigQuery, PostgreSQL
 export interface GroupByRollup extends BaseNode {
   type: "group_by_rollup";
@@ -284,6 +287,13 @@ export interface GroupByCube extends BaseNode {
   columns: ParenExpr<ListExpr<Expr>>;
 }
 
+// PostgreSQL
+export interface GroupByGroupingSets extends BaseNode {
+  type: "group_by_grouping_sets";
+  groupingSetsKw: Keyword<"GROUPING SETS">;
+  columns: ParenExpr<ListExpr<GroupingElement>>;
+}
+
 export interface HavingClause extends BaseNode {
   type: "having_clause";
   havingKw: Keyword<"HAVING">;

src/parser.pegjs

@@ -802,6 +802,7 @@ group_by_distinct
 grouping_element
   = group_by_rollup
   / group_by_cube
+  / group_by_grouping_sets
   / paren$empty_list
   / expr
 
@@ -823,6 +824,15 @@ group_by_cube
     });
   }
 
+group_by_grouping_sets
+  = kw:(GROUPING __ SETS __) columns:paren$list$grouping_element &postgres {
+    return loc({
+      type: "group_by_grouping_sets",
+      groupingSetsKw: read(kw),
+      columns,
+    });
+  }
+
 /**
  * SELECT .. HAVING
  * --------------------------------------------------------------------------------------
@@ -4951,6 +4961,7 @@ paren$list$equals_expr = .
 paren$list$expr = .
 paren$list$expr_or_default = .
 paren$list$func_param = .
+paren$list$grouping_element = .
 paren$list$ident = .
 paren$list$literal = .
 paren$list$procedure_param = .
@@ -6035,6 +6046,7 @@ GRANT               = kw:"GRANT"i               !ident_part { return loc(createK
 GRANTS              = kw:"GRANTS"i              !ident_part { return loc(createKeyword(kw)); }
 GROUP               = kw:"GROUP"i               !ident_part { return loc(createKeyword(kw)); }
 GROUP_CONCAT        = kw:"GROUP_CONCAT"i        !ident_part { return loc(createKeyword(kw)); }
+GROUPING            = kw:"GROUPING"i            !ident_part { return loc(createKeyword(kw)); }
 GROUPS              = kw:"GROUPS"i              !ident_part { return loc(createKeyword(kw)); }
 HASH                = kw:"HASH"i                !ident_part { return loc(createKeyword(kw)); }
 HAVING              = kw:"HAVING"i              !ident_part { return loc(createKeyword(kw)); }
@@ -6240,6 +6252,7 @@ SELECT              = kw:"SELECT"i              !ident_part { return loc(createK
 SESSION             = kw:"SESSION"i             !ident_part { return loc(createKeyword(kw)); }
 SESSION_USER        = kw:"SESSION_USER"i        !ident_part { return loc(createKeyword(kw)); }
 SET                 = kw:"SET"i                 !ident_part { return loc(createKeyword(kw)); }
+SETS                = kw:"SETS"i                !ident_part { return loc(createKeyword(kw)); }
 SHARE               = kw:"SHARE"i               !ident_part { return loc(createKeyword(kw)); }
 SHARED              = kw:"SHARED"i              !ident_part { return loc(createKeyword(kw)); }
 SHOW                = kw:"SHOW"i                !ident_part { return loc(createKeyword(kw)); }

src/showNode/select.ts

@@ -103,6 +103,7 @@ export const selectMap: FullTransformMap<string, AllSelectNodes> = {
     show([node.groupByKw, node.distinctKw, node.columns, node.withRollupKw]),
   group_by_rollup: (node) => show([node.rollupKw, node.columns]),
   group_by_cube: (node) => show([node.cubeKw, node.columns]),
+  group_by_grouping_sets: (node) => show([node.groupingSetsKw, node.columns]),
   having_clause: (node) => show([node.havingKw, node.expr]),
   qualify_clause: (node) => show([node.qualifyKw, node.expr]),
   order_by_clause: (node) =>

test/select/group_by.test.ts

@@ -47,6 +47,24 @@ describe("select GROUP BY", () => {
     });
   });
 
+  dialect("postgresql", () => {
+    it("supports GROUPING SETS", () => {
+      testClauseWc("GROUP BY GROUPING SETS ( (id, name), age )");
+    });
+
+    it("supports multiple GROUPING SETS", () => {
+      testClauseWc("GROUP BY GROUPING SETS (id), GROUPING SETS (name)");
+    });
+
+    it("supports ROLLUP & CUBE inside GROUPING SETS", () => {
+      testClauseWc("GROUP BY GROUPING SETS ( ROLLUP(id), CUBE(age) )");
+    });
+
+    it("supports GROUPING SETS inside GROUPING SETS", () => {
+      testClauseWc("GROUP BY GROUPING SETS ( GROUPING SETS (foo, bar), GROUPING SETS (age) )");
+    });
+  });
+
   dialect(["mysql", "mariadb"], () => {
     it("supports GROUP BY .. WITH ROLLUP", () => {
       testClauseWc("GROUP BY id, name WITH ROLLUP");