Merge pull request #2228 from taozhi8833998/fix-key-athena

fix: key as column in athena
taozhi8833998 · Nov 20, 2024 · dbef224 · dbef224
2 parents 125b7c9 + befa43c
commit dbef224
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 2 deletions.
diff --git a/pegjs/athena.pegjs b/pegjs/athena.pegjs
@@ -47,7 +47,7 @@
     'JOIN': true,
     'JSON': true,
 
-    'KEY': true,
+    // 'KEY': true,
 
     'LEFT': true,
     'LIKE': true,
@@ -1845,7 +1845,16 @@ unary_operator
   = '!' / '-' / '+' / '~'
 
 column_ref
-  = tbl:ident __ DOT __ col:column {
+  = schema:ident tbl:(__ DOT __ ident) col:(__ DOT __ column) {
+      columnList.add(`select::${schema}.${tbl[3]}::${col[3].value}`);
+      return {
+        type: 'column_ref',
+        schema: schema,
+        table: tbl[3],
+        column: col[3]
+      };
+    }
+  / tbl:ident __ DOT __ col:column {
       columnList.add(`select::${tbl}::${col}`);
       return {
         type: 'column_ref',

diff --git a/test/athena.spec.js b/test/athena.spec.js
@@ -304,4 +304,33 @@ describe('athena', () => {
       ) = 0;`
       expect(getParsedSql(sql)).to.be.equal('SELECT `id`, CAST(CURRENT_TIMESTAMP AS TIMESTAMP(6)) AS `dbt_insert_time` FROM `some_table` WHERE cardinality(FILTER(map_values(`note`), VALUE -> `VALUE` IS NOT NULL)) = 0')
   })
+  it('should support key as column name', () => {
+    let sql = `WITH CTE AS (
+      SELECT * FROM test_cte
+    )
+    SELECT
+        organization,
+        date,
+        author_email,
+        t.key AS tag,
+        t.value AS count
+    FROM CTE
+    CROSS JOIN UNNEST(tags_counts) AS t(key, value)
+    ORDER BY 1, 2`;
+    expect(getParsedSql(sql)).to.be.equal('WITH `CTE` AS (SELECT * FROM `test_cte`) SELECT `organization`, DATE , `author_email`, `t`.`key` AS `tag`, `t`.`value` AS `count` FROM `CTE` CROSS JOIN UNNEST(`tags_counts`) AS t(`key`, `value`) ORDER BY 1 ASC, 2 ASC')
+    sql = `SELECT
+        j.id,
+        h.created AS change_time,
+        i.fromstring AS from_status,
+        i.tostring AS to_status
+    FROM
+        "bronze_prod"."jira_issues" j
+    CROSS JOIN
+        UNNEST(j.changelog.histories) AS T (h)
+    CROSS JOIN
+        UNNEST(h.items) AS T (i)
+    WHERE
+        i.field = 'status'`
+    expect(getParsedSql(sql)).to.be.equal("SELECT `j`.`id`, `h`.`created` AS `change_time`, `i`.`fromstring` AS `from_status`, `i`.`tostring` AS `to_status` FROM `bronze_prod`.`jira_issues` AS `j` CROSS JOIN UNNEST(`j`.`changelog`.`histories`) AS T(`h`) CROSS JOIN UNNEST(`h`.`items`) AS T(`i`) WHERE `i`.`field` = 'status'")
+  })
 })