Merge pull request #1977 from taozhi8833998/feat-top-union-snowflake

feat: support top clause and fix multiple union all in snowflake
taozhi8833998 · Jun 26, 2024 · 44bed0a · 44bed0a
2 parents f20bf65 + 144acc1
commit 44bed0a
Show file tree

Hide file tree

Showing 15 changed files with 53 additions and 20 deletions.
diff --git a/pegjs/athena.pegjs b/pegjs/athena.pegjs
@@ -1158,7 +1158,7 @@ table_join
       t.on   = expr;
       return t;
     }
-  / op:join_op __ LPAREN __ stmt:union_stmt __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
+  / op:(join_op / set_op) __ LPAREN __ stmt:union_stmt __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
     stmt.parentheses = true;
     return {
       expr: stmt,

diff --git a/pegjs/bigquery.pegjs b/pegjs/bigquery.pegjs
@@ -1817,7 +1817,7 @@ table_join
       t.on   = expr;
       return t;
     }
-  / op:join_op __ LPAREN __ stmt:union_stmt __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
+  / op:(join_op / set_op) __ LPAREN __ stmt:union_stmt __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
     stmt.parentheses = true;
     return {
       expr: stmt,

diff --git a/pegjs/db2.pegjs b/pegjs/db2.pegjs
@@ -1198,7 +1198,7 @@ table_join
       t.on   = expr;
       return t;
     }
-  / op:join_op __ LPAREN __ stmt:union_stmt __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
+  / op:(join_op / set_op) __ LPAREN __ stmt:union_stmt __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
     stmt.parentheses = true;
     return {
       expr: stmt,

diff --git a/pegjs/flinksql.pegjs b/pegjs/flinksql.pegjs
@@ -1871,11 +1871,11 @@ table_join
       t.on   = expr;
       return t;
     }
-  / op:join_op __ LPAREN __ stmt:union_stmt __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
+  / op:(join_op / set_op) __ LPAREN __ stmt:union_stmt __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
     /* => {
       expr: union_stmt & { parentheses: true; };
       as?: alias_clause;
-      join: join_op;
+      join: join_op | set_op;
       on?: on_clause;
     }*/
     stmt.parentheses = true;

diff --git a/pegjs/hive.pegjs b/pegjs/hive.pegjs
@@ -1158,7 +1158,7 @@ table_join
       t.on   = expr;
       return t;
     }
-  / op:join_op __ LPAREN __ stmt:union_stmt __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
+  / op:(join_op / set_op) __ LPAREN __ stmt:union_stmt __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
     stmt.parentheses = true;
     return {
       expr: stmt,

diff --git a/pegjs/mariadb.pegjs b/pegjs/mariadb.pegjs
@@ -2112,7 +2112,7 @@ table_join
       t.on   = expr;
       return t;
     }
-  / op:join_op __ LPAREN __ stmt:set_op_stmt __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
+  / op:(join_op / set_op) __ LPAREN __ stmt:set_op_stmt __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
     stmt.parentheses = true;
     return {
       expr: stmt,

diff --git a/pegjs/mysql.pegjs b/pegjs/mysql.pegjs
@@ -2366,7 +2366,7 @@ table_join
       t.on = expr;
       return t;
     }
-  / op:join_op __ LPAREN __ stmt:set_op_stmt __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
+  / op:(join_op / set_op) __ LPAREN __ stmt:set_op_stmt __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
     stmt.parentheses = true;
     return {
       expr: stmt,

diff --git a/pegjs/noql.pegjs b/pegjs/noql.pegjs
@@ -3010,7 +3010,7 @@ table_join
       t.on = expr;
       return t;
     }
-  / op:join_op __ LPAREN __ stmt:(union_stmt / table_ref_list) __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
+  / op:(join_op / set_op) __ LPAREN __ stmt:(union_stmt / table_ref_list) __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
     /* => {
       expr: (union_stmt || table_ref_list) & { parentheses: true; };
       as?: alias_clause;

diff --git a/pegjs/postgresql.pegjs b/pegjs/postgresql.pegjs
@@ -3211,11 +3211,11 @@ table_join
       t.on = expr;
       return t;
     }
-  / op:join_op __ LPAREN __ stmt:(union_stmt / table_ref_list) __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
+  / op:(join_op / set_op) __ LPAREN __ stmt:(union_stmt / table_ref_list) __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
     /* => {
       expr: (union_stmt | table_ref_list) & { parentheses: true; };
       as?: alias_clause;
-      join: join_op;
+      join: join_op | set_op;
       on?: on_clause;
     }*/
     if (Array.isArray(stmt)) stmt = { type: 'tables', expr: stmt }

diff --git a/pegjs/redshift.pegjs b/pegjs/redshift.pegjs
@@ -3044,11 +3044,11 @@ table_join
       t.on = expr;
       return t;
     }
-  / op:join_op __ LPAREN __ stmt:(union_stmt / table_ref_list) __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
+  / op:(join_op / set_op) __ LPAREN __ stmt:(union_stmt / table_ref_list) __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
     /* => {
       expr: (union_stmt || table_ref_list) & { parentheses: true; };
       as?: alias_clause;
-      join: join_op;
+      join: join_op | set_op;
       on?: on_clause;
     }*/
     if (Array.isArray(stmt)) stmt = { type: 'tables', expr: stmt }

diff --git a/pegjs/snowflake.pegjs b/pegjs/snowflake.pegjs
@@ -2061,6 +2061,7 @@ select_stmt_nake
   = __ cte:with_clause? __ KW_SELECT ___
     opts:option_clause? __
     d:distinct_on?      __
+    top:top_clause? __
     c:column_clause     __
     ci:into_clause?      __
     f:from_clause?      __
@@ -2085,6 +2086,7 @@ select_stmt_nake
           groupby?: group_by_clause;
           having?: having_clause;
           orderby?: order_by_clause;
+          top?: top_clause;
           limit?: limit_clause;
           window?: window_clause;
         }*/
@@ -2108,12 +2110,27 @@ select_stmt_nake
           having: h,
           qualify: q,
           orderby: o,
+          top,
           limit: l,
           window: win,
           ...getLocationObject()
       };
   }
 
+top_clause
+  = KW_TOP __ LPAREN __ n:number __ RPAREN __ p:('PERCENT'i)? {
+    return {
+      value: n,
+      percent: p && p.toLowerCase(),
+      parentheses: true,
+    }
+  }
+  / KW_TOP __ n:number __ p:('PERCENT'i)? {
+    return {
+      value: n,
+      percent: p && p.toLowerCase()
+    }
+  }
 // MySQL extensions to standard SQL
 option_clause
   = head:query_option tail:(__ query_option)* {
@@ -2389,11 +2406,11 @@ table_join
       t.on = expr;
       return t;
     }
-  / op:join_op __ LPAREN __ stmt:(union_stmt / table_ref_list) __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
+  / op:(join_op / set_op) __ LPAREN __ stmt:(union_stmt / table_ref_list) __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
     /* => {
       expr: (union_stmt || table_ref_list) & { parentheses: true; };
       as?: alias_clause;
-      join: join_op;
+      join: join_op | set_op;
       on?: on_clause;
     }*/
     if (Array.isArray(stmt)) stmt = { type: 'tables', expr: stmt }
@@ -2407,6 +2424,7 @@ table_join
     };
   }
 
+
 //NOTE that, the table assigned to `var` shouldn't write in `table_join`
 table_base
   = KW_DUAL {
@@ -4245,6 +4263,7 @@ KW_DEFAULT  = "DEFAULT"i    !ident_start
 KW_NOT_NULL = "NOT NULL"i   !ident_start
 KW_TRUE     = "TRUE"i       !ident_start
 KW_TO       = "TO"i         !ident_start
+KW_TOP      = "TOP"i        !ident_start
 KW_FALSE    = "FALSE"i      !ident_start
 
 KW_SHOW     = "SHOW"i       !ident_start

diff --git a/pegjs/sqlite.pegjs b/pegjs/sqlite.pegjs
@@ -1514,7 +1514,7 @@ table_join
       t.on   = expr;
       return t;
     }
-  / op:join_op __ LPAREN __ stmt:union_stmt __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
+  / op:(join_op / set_op) __ LPAREN __ stmt:union_stmt __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
     stmt.parentheses = true;
     return {
       expr: stmt,

diff --git a/pegjs/transactsql.pegjs b/pegjs/transactsql.pegjs
@@ -1410,7 +1410,7 @@ select_stmt_nake
   = __ cte:with_clause? __ KW_SELECT ___
     opts:option_clause? __
     d:KW_DISTINCT?      __
-    top: top_clause? __
+    top:top_clause? __
     c:column_clause     __
     f:from_clause?      __
     w:where_clause?     __
@@ -1727,7 +1727,7 @@ table_join
       t.on   = expr;
       return t;
     }
-  / op:join_op __ LPAREN __ stmt:union_stmt __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
+  / op:(join_op / set_op) __ LPAREN __ stmt:union_stmt __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
     stmt.parentheses = true;
     return {
       expr: stmt,

diff --git a/pegjs/trino.pegjs b/pegjs/trino.pegjs
@@ -2404,11 +2404,11 @@ table_join
       t.on = expr;
       return t;
     }
-  / op:join_op __ LPAREN __ stmt:(union_stmt / table_ref_list) __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
+  / op:(join_op / set_op) __ LPAREN __ stmt:(union_stmt / table_ref_list) __ RPAREN __ alias:alias_clause? __ expr:on_clause? {
     /* => {
       expr: (union_stmt || table_ref_list) & { parentheses: true; };
       as?: alias_clause;
-      join: join_op;
+      join: join_op | set_op;
       on?: on_clause;
     }*/
     if (Array.isArray(stmt)) stmt = { type: 'tables', expr: stmt }

diff --git a/test/snowflake.spec.js b/test/snowflake.spec.js
@@ -249,6 +249,20 @@ describe('snowflake', () => {
         'SELECT "A", "B", COUNT(*) FROM "T" GROUP BY ALL'
       ]
     },
+    {
+      title: 'top clause in select',
+      sql: [
+        'select top 10 * from cleansed_hosts',
+        'SELECT TOP 10 * FROM "cleansed_hosts"'
+      ]
+    },
+    {
+      title: 'multiple union all',
+      sql: [
+        'select * from ( ( ( select * from some_table ) union all ( select * from some_table ) )  union all ( select * from some_table )  )',
+        'SELECT * FROM (((SELECT * FROM "some_table") UNION ALL (SELECT * FROM "some_table")) UNION ALL (SELECT * FROM "some_table"))'
+      ]
+    },
   ]
   SQL_LIST.forEach(sqlInfo => {
     const { title, sql } = sqlInfo