From 1a3f9cd25975a97a93eac691b962ef3b55c84260 Mon Sep 17 00:00:00 2001
From: taozhi8833998 <taozhi8833998@163.com>
Date: Mon, 18 Dec 2023 10:48:17 +0800
Subject: [PATCH] feat: support json type, fix array column vistior in bigquery

---
 pegjs/bigquery.pegjs  | 32 ++++++++++++++++++++++----------
 src/column.js         |  3 ++-
 src/expr.js           |  2 ++
 src/json.js           | 12 ++++++++++++
 test/bigquery.spec.js | 24 ++++++++++++++++++++++++
 5 files changed, 62 insertions(+), 11 deletions(-)
 create mode 100644 src/json.js

diff --git a/pegjs/bigquery.pegjs b/pegjs/bigquery.pegjs
index c911af3e..a14c9a90 100644
--- a/pegjs/bigquery.pegjs
+++ b/pegjs/bigquery.pegjs
@@ -1650,17 +1650,18 @@ columns_list
       return createList(head, tail);
     }
 
-column_offset_expr
-  = n:expr __ LBRAKE __ l:literal_numeric __ RBRAKE {
-    return {
-      expr: n,
-      offset: `[${l.value}]`
-    }
+column_offset_expr_list
+  = l:(LBRAKE __ (literal_numeric / literal_string) __ RBRAKE)+ {
+    return l.map(item => ({ value: item[2] }))
+  }
+  / l:(LBRAKE __ (KW_OFFSET / KW_ORDINAL / KW_SAFE_OFFSET / KW_SAFE_ORDINAL) __ LPAREN __ (literal_numeric / literal_string) __ RPAREN __ RBRAKE)+ {
+    return l.map(item => ({ name: item[2], value: item[6] }))
   }
-  / n:expr __ LBRAKE __ t:(KW_OFFSET / KW_ORDINAL / KW_SAFE_OFFSET / KW_SAFE_ORDINAL) __ LPAREN __ l:literal_numeric __ RPAREN __ RBRAKE {
+column_offset_expr
+  = n:expr __ l:column_offset_expr_list {
     return {
       expr: n,
-      offset: `[${t}(${l.value})]`
+      offset: l
     }
   }
 
@@ -1706,8 +1707,8 @@ column_list_item
         as: null
       }
     }
-  / c:column_offset_expr __ as:alias_clause? {
-    columnList.add(`select::null::${c}`)
+  / c:column_offset_expr __ s:(DOT __ column_without_kw)? __ as:alias_clause? {
+    if (s) c.suffix = `.${s[2]}`
     return {
         expr: {
           type: 'column_ref',
@@ -2014,6 +2015,7 @@ expr_list
 
 _expr
   = struct_expr
+  / json_expr
   / logic_operator_expr // support concatenation operator || and &&
   / or_expr
   / unary_expr
@@ -2062,6 +2064,15 @@ array_expr
     }
   }
 
+json_expr
+  = KW_JSON __ l:literal_list {
+    return {
+      type: 'json',
+      keyword: 'json',
+      expr_list: l
+    }
+  }
+
 struct_expr
   = s:(struct_type / KW_STRUCT) __ LPAREN __ c:column_clause __ RPAREN {
     return {
@@ -2240,6 +2251,7 @@ multiplicative_operator
 primary
   = array_expr
   / struct_expr
+  / json_expr
   / cast_expr
   / literal
   / aggr_func
diff --git a/src/column.js b/src/column.js
index 09fb63a0..28b03159 100644
--- a/src/column.js
+++ b/src/column.js
@@ -18,7 +18,8 @@ import {
 function columnOffsetToSQL(column, isDual) {
   if (typeof column === 'string') return identifierToSql(column, isDual)
   const { expr, offset, suffix } = column
-  return [exprToSQL(expr), offset, suffix].filter(hasVal).join('')
+  const offsetExpr = offset && offset.map(offsetItem => ['[', offsetItem.name, `${offsetItem.name ? '(' : ''}`, literalToSQL(offsetItem.value), `${offsetItem.name ? ')' : ''}`, ']'].filter(hasVal).join('')).join('')
+  return [exprToSQL(expr), offsetExpr, suffix].filter(hasVal).join('')
 }
 function columnRefToSQL(expr) {
   const {
diff --git a/src/expr.js b/src/expr.js
index 75ecb857..840680e7 100644
--- a/src/expr.js
+++ b/src/expr.js
@@ -7,6 +7,7 @@ import { caseToSQL } from './case'
 import { columnDefinitionToSQL, columnRefToSQL, fullTextSearchToSQL } from './column'
 import { anyValueFuncToSQL, castToSQL, extractFunToSQL, flattenFunToSQL, funcToSQL, tablefuncFunToSQL } from './func'
 import { intervalToSQL } from './interval'
+import { jsonExprToSQL } from './json'
 import { selectToSQL } from './select'
 import { showToSQL } from './show'
 import { arrayStructExprToSQL } from './array-struct'
@@ -33,6 +34,7 @@ const exprToSQLConvertFn = {
   function          : funcToSQL,
   insert            : unionToSQL,
   interval          : intervalToSQL,
+  json              : jsonExprToSQL,
   show              : showToSQL,
   struct            : arrayStructExprToSQL,
   tablefunc         : tablefuncFunToSQL,
diff --git a/src/json.js b/src/json.js
new file mode 100644
index 00000000..78e2eee5
--- /dev/null
+++ b/src/json.js
@@ -0,0 +1,12 @@
+import { exprToSQL } from './expr'
+import { toUpper } from './util'
+
+function jsonExprToSQL(expr) {
+  const { keyword, expr_list: exprList } = expr
+  const result = [toUpper(keyword), exprList.map(exprItem => exprToSQL(exprItem)).join(', ')].join(' ')
+  return result
+}
+
+export {
+  jsonExprToSQL,
+}
diff --git a/test/bigquery.spec.js b/test/bigquery.spec.js
index c0df6c11..d5ad23a8 100644
--- a/test/bigquery.spec.js
+++ b/test/bigquery.spec.js
@@ -823,6 +823,30 @@ describe('BigQuery', () => {
         'SELECT if(((a)), b, NULL)'
       ]
     },
+    {
+      title: 'offset column with dot',
+      sql: [
+        `WITH your_table AS (
+          SELECT [STRUCT(1 AS id, 'John' AS name), STRUCT(2 AS id, 'Jane' AS name)] AS some_array_column
+        )
+        SELECT some_array_column[SAFE_OFFSET(0)].id from your_table`,
+        "WITH your_table AS (SELECT [STRUCT(1 AS id, 'John' AS name), STRUCT(2 AS id, 'Jane' AS name)] AS some_array_column) SELECT some_array_column[SAFE_OFFSET(0)].id FROM your_table"
+      ]
+    },
+    {
+      title: 'json expr',
+      sql: [
+        `SELECT json_value.class.students[0]['name'] AS first_student
+        FROM
+          UNNEST(
+            [
+              JSON '{"class" : {"students" : [{"name" : "Jane"}]}}',
+              JSON '{"class" : {"students" : []}}',
+              JSON '{"class" : {"students" : [{"name" : "John"}, {"name": "Jamie"}]}}'])
+        AS json_value`,
+        `SELECT json_value.class.students[0]['name'] AS first_student FROM UNNEST([JSON '{"class" : {"students" : [{"name" : "Jane"}]}}', JSON '{"class" : {"students" : []}}', JSON '{"class" : {"students" : [{"name" : "John"}, {"name": "Jamie"}]}}']) AS json_value`
+      ]
+    }
   ]
 
   SQL_LIST.forEach(sqlInfo => {