feat: support at time zone in pg

Author: taozhi8833998

pegjs/postgresql.pegjs

@@ -3292,13 +3292,39 @@ cast_data_type
     return t
   }
 
+column_item_suffix
+  = 'AT' __ KW_TIME __ 'ZONE'i __ t:quoted_ident_type {
+    // => [{ type: 'origin'; value: string; }, quoted_ident_type]
+    return [
+      {
+        type: 'origin',
+        value: 'at time zone'
+      },
+      t
+    ]
+  }
 column_list_item
   = c:string_constants_escape {
     // => { expr: expr; as: null; }
     return { expr: c, as: null }
   }
-  / e:(column_ref_quoted / expr_item) __ s:KW_DOUBLE_COLON __ t:cast_data_type __ tail:(__ (additive_operator / multiplicative_operator) __ expr_item)* __ alias:alias_clause? {
+  / e:column_ref __ s:KW_DOUBLE_COLON __ t:cast_data_type __ cs:column_item_suffix+ __ alias:alias_clause? {
+    // => { type: 'cast'; expr: expr; symbol: '::'; target: cast_data_type;  as?: null; }
+    t.suffix = cs.flat()
+    return {
+      as: alias,
+      type: 'cast',
+      expr: e,
+      symbol: '::',
+      target: t,
+      suffix: cs.flat(),
+    }
+  }
+  / e:(column_ref_quoted / expr_item) __ s:KW_DOUBLE_COLON __ t:cast_data_type __ tail:(__ (additive_operator / multiplicative_operator) __ expr_item)* __ cs:column_item_suffix* __ alias:alias_clause? {
     // => { type: 'cast'; expr: expr; symbol: '::'; target: cast_data_type;  as?: null; }
+    if (e.type === 'column_ref' && cs.length) {
+      e.column.options = { type: 'expr_list', value: cs.flat(), separator: ' ' }
+    }
     return {
       as: alias,
       type: 'cast',
@@ -3341,7 +3367,7 @@ column_list_item
         as: null
       };
     }
-  / e:expr_item  __ alias:alias_clause? {
+  / e:expr_item __ alias:alias_clause? {
     // => { type: 'expr'; expr: expr; as?: alias_clause; }
       return { type: 'expr', expr: e, as: alias };
     }
@@ -4074,7 +4100,7 @@ value_item
 
 expr_list
   = head:expr tail:(__ COMMA __ expr)* {
-    // => { type: 'expr_list'; value: expr[] }
+    // => { type: 'expr_list'; value: expr[]; parentheses?: boolean; separator?: string; }
       const el = { type: 'expr_list' };
       el.value = createList(head, tail);
       return el;
@@ -4404,9 +4430,10 @@ multiplicative_operator
   = "*" / "/" / "%" / "||"
 
 column_ref_array_index
-  = c:column_ref __ a:array_index_list? {
+  = c:column_ref __ a:array_index_list? __ cs:column_item_suffix* {
     // => column_ref
     if (a) c.array_index = a
+    if (cs.length) c.options = { type: 'expr_list', value: cs.flat(), separator: ' ' }
     return c
   }
 

src/column.js

@@ -35,7 +35,7 @@ function arrayIndexToSQL(arrayIndexList) {
 }
 function columnRefToSQL(expr) {
   const {
-    array_index, as, column, db, isDual, notations = [], schema, table, parentheses,
+    array_index, as, column, db, isDual, notations = [], options, schema, table, parentheses,
     suffix, order_by, subFields = [],
   } = expr
   let str = column === '*' ? '*' : columnOffsetToSQL(column, isDual)
@@ -51,6 +51,7 @@ function columnRefToSQL(expr) {
   str = [`${str}${arrayIndexToSQL(array_index)}`, ...subFields].join('.')
   const result = [
     str,
+    exprToSQL(options),
     commonOptionConnector('AS', exprToSQL, as),
   ]
   result.push(typeof suffix === 'string' ? toUpper(suffix) : exprToSQL(suffix))

src/expr.js

@@ -89,8 +89,12 @@ function getExprListSQL(exprList) {
 }
 
 exprToSQLConvertFn.expr_list = expr => {
-  const str = getExprListSQL(expr.value)
-  return expr.parentheses ? `(${str.join(', ')})` : str
+  const result = getExprListSQL(expr.value)
+  const { parentheses, separator } = expr
+  if (!parentheses && !separator) return result
+  const joinSymbol = separator || ', '
+  const str = result.join(joinSymbol)
+  return parentheses ? `(${str})` : str
 }
 
 exprToSQLConvertFn.select = expr => {

test/postgres.spec.js

@@ -1657,7 +1657,7 @@ describe('Postgres', () => {
       ]
     },
     {
-      title: 'substring',
+      title: 'substring function',
       sql: [
         `SELECT AVG(
             CASE
@@ -1672,6 +1672,20 @@ describe('Postgres', () => {
         `SELECT AVG(CASE WHEN "duration" LIKE '%min%' THEN CAST(SUBSTRING("duration" FROM '([0-9]+)') AS INTEGER) WHEN "duration" LIKE '%hr%' THEN CAST(SUBSTRING("duration" FROM '([0-9]+)') AS INTEGER) * 60 ELSE NULL END) AS "average_duration" FROM "netflix_shows" WHERE "listed_in" ILIKE '%Thriller%' LIMIT 100`
       ]
     },
+    {
+      title: 'column at time zone',
+      sql: [
+        "SELECT start_time AT TIME ZONE 'UTC' AS start_time FROM my_table",
+        `SELECT start_time AT TIME ZONE 'UTC' AS "start_time" FROM "my_table"`
+      ]
+    },
+    {
+      title: 'column cast with at time zone',
+      sql: [
+        "SELECT start_time::timestamp AT TIME ZONE 'UTC' AS start_time FROM my_table",
+        `SELECT start_time::TIMESTAMP AT TIME ZONE 'UTC' AS "start_time" FROM "my_table"`
+      ]
+    },
   ]
   function neatlyNestTestedSQL(sqlList){
     sqlList.forEach(sqlInfo => {

types.d.ts

@@ -103,6 +103,7 @@ export interface ColumnRefItem {
   type: "column_ref";
   table: string | null;
   column: string | { expr: ValueExpr };
+  options?: ExprList;
   loc?: LocationRange;
 }
 export interface ColumnRefExpr {
@@ -220,6 +221,8 @@ export type ExprList = {
   type: "expr_list";
   value: ExpressionValue[];
   loc?: LocationRange;
+  parentheses?: boolean;
+  separator?: string;
 };
 
 export type PartitionBy = {