Merge pull request #1999 from taozhi8833998/feat-from-unnest-athena

feat: support from unnest in athena
taozhi8833998 · Jul 11, 2024 · 7bfe03b · 7bfe03b
2 parents 63673f2 + ebd6886
commit 7bfe03b
Show file tree

Hide file tree

Showing 4 changed files with 102 additions and 4 deletions.
diff --git a/pegjs/athena.pegjs b/pegjs/athena.pegjs
@@ -82,6 +82,7 @@
     'UNION': true,
     'UPDATE': true,
     'USING': true,
+    'UNNEST': true,
 
     'VALUES': true,
 
@@ -1088,9 +1089,28 @@ column_list_item
     }
 
 alias_clause
-  = KW_AS __ i:alias_ident { return i; }
+  = KW_AS __ i:(func_call / alias_ident) { return i; }
   / KW_AS? __ i:ident { return i; }
 
+with_offset
+  = KW_WITH __ KW_OFFSET __ alias:alias_clause? {
+    return {
+      keyword: 'with offset as',
+      as: alias
+    }
+  }
+
+from_unnest_item
+  = 'UNNEST'i __ LPAREN __ a:expr? __ RPAREN __ alias:alias_clause? __ wf:with_offset? {
+    return {
+      type: 'unnest',
+      expr: a,
+      parentheses: true,
+      as: alias,
+      with_offset: wf,
+    }
+  }
+
 from_clause
   = KW_FROM __ l:table_ref_list { return l; }
 
@@ -1186,7 +1206,8 @@ table_join
 
 //NOTE that, the table assigned to `var` shouldn't write in `table_join`
 table_base
-  = KW_DUAL {
+  = from_unnest_item
+  / KW_DUAL {
       return {
         type: 'dual'
       };
@@ -2142,6 +2163,17 @@ literal
   / literal_bool
   / literal_null
   / literal_datetime
+  / literal_array
+
+literal_array
+  = s:KW_ARRAY __ LBRAKE __ c:expr_list? __ RBRAKE {
+    return {
+      expr_list: c || { type: 'origin', value: '' },
+      type: 'array',
+      keyword: 'array',
+      brackets: true
+    }
+  }
 
 literal_list
   = head:literal tail:(__ COMMA __ literal)* {

diff --git a/src/tables.js b/src/tables.js
@@ -9,7 +9,7 @@ function unnestToSQL(unnestExpr) {
   const { type, as, expr, with_offset: withOffset } = unnestExpr
   const result = [
     `${toUpper(type)}(${expr && exprToSQL(expr) || ''})`,
-    commonOptionConnector('AS', identifierToSql, as),
+    commonOptionConnector('AS', typeof as === 'string' ? identifierToSql : exprToSQL, as),
     commonOptionConnector(
       toUpper(withOffset && withOffset.keyword),
       identifierToSql,

diff --git a/test/athena.spec.js b/test/athena.spec.js
diff --git a/test/util.spec.js b/test/util.spec.js
@@ -52,7 +52,7 @@ describe('util function test', () => {
     expect(columnIdentifierToSql("id")).to.be.equal('"id"')
   })
 
-  it.only('should support trim query option', () => {
+  it('should support trim query option', () => {
     const opt = {
       "database": "mysql",
       "parseOptions": {