feat: support data type array in pg

taozhi8833998 · Jul 12, 2024 · f7ba2b9 · f7ba2b9
1 parent 79d4a87
commit f7ba2b9
Show file tree

Hide file tree

Showing 8 changed files with 60 additions and 21 deletions.
diff --git a/pegjs/noql.pegjs b/pegjs/noql.pegjs
@@ -5254,11 +5254,15 @@ data_type
 array_type
   = t:(numeric_type / character_string_type) __ LBRAKE __ RBRAKE __ LBRAKE __ RBRAKE {
     /* => data_type */
-    return { ...t, array: 'two' }
+    return { ...t, array: { dimension: 2 } }
   }
-  / t:(numeric_type / character_string_type) __ LBRAKE __ RBRAKE {
+  / t:(numeric_type / character_string_type) __ LBRAKE __ l:literal_numeric? __ RBRAKE {
     /* => data_type */
-    return { ...t, array: 'one' }
+    return { ...t, array: { dimension: 1,  length: [l] } }
+  }
+  / t:(numeric_type / character_string_type) __ KW_ARRAY {
+    /* => data_type */
+    return { ...t, array: { keyword: 'array' } }
   }
 
 boolean_type

diff --git a/pegjs/postgresql.pegjs b/pegjs/postgresql.pegjs
@@ -5528,12 +5528,17 @@ data_type
 array_type
   = t:(numeric_type / character_string_type) __ LBRAKE __ RBRAKE __ LBRAKE __ RBRAKE {
     /* => data_type */
-    return { ...t, array: 'two' }
+    return { ...t, array: { dimension: 2 } }
   }
-  / t:(numeric_type / character_string_type) __ LBRAKE __ RBRAKE {
+  / t:(numeric_type / character_string_type) __ LBRAKE __ l:literal_numeric? __ RBRAKE {
     /* => data_type */
-    return { ...t, array: 'one' }
+    return { ...t, array: { dimension: 1,  length: [l] } }
   }
+  / t:(numeric_type / character_string_type) __ KW_ARRAY {
+    /* => data_type */
+    return { ...t, array: { keyword: 'array' } }
+  }
+
 
 boolean_type
   = t:(KW_BOOL / KW_BOOLEAN) { /* => data_type */ return { dataType: t }}

diff --git a/pegjs/redshift.pegjs b/pegjs/redshift.pegjs
@@ -5329,11 +5329,15 @@ data_type
 array_type
   = t:(numeric_type / character_string_type) __ LBRAKE __ RBRAKE __ LBRAKE __ RBRAKE {
     /* => data_type */
-    return { ...t, array: 'two' }
+    return { ...t, array: { dimension: 2 } }
   }
-  / t:(numeric_type / character_string_type) __ LBRAKE __ RBRAKE {
+  / t:(numeric_type / character_string_type) __ LBRAKE __ l:literal_numeric? __ RBRAKE {
     /* => data_type */
-    return { ...t, array: 'one' }
+    return { ...t, array: { dimension: 1,  length: [l] } }
+  }
+  / t:(numeric_type / character_string_type) __ KW_ARRAY {
+    /* => data_type */
+    return { ...t, array: { keyword: 'array' } }
   }
 
 boolean_type

diff --git a/pegjs/snowflake.pegjs b/pegjs/snowflake.pegjs
@@ -4778,11 +4778,15 @@ data_type
 array_type
   = t:(numeric_type / character_string_type) __ LBRAKE __ RBRAKE __ LBRAKE __ RBRAKE {
     /* => data_type */
-    return { ...t, array: 'two' }
+    return { ...t, array: { dimension: 2 } }
   }
-  / t:(numeric_type / character_string_type) __ LBRAKE __ RBRAKE {
+  / t:(numeric_type / character_string_type) __ LBRAKE __ l:literal_numeric? __ RBRAKE {
     /* => data_type */
-    return { ...t, array: 'one' }
+    return { ...t, array: { dimension: 1,  length: [l] } }
+  }
+  / t:(numeric_type / character_string_type) __ KW_ARRAY {
+    /* => data_type */
+    return { ...t, array: { keyword: 'array' } }
   }
 
 boolean_type

diff --git a/pegjs/trino.pegjs b/pegjs/trino.pegjs
@@ -4779,11 +4779,15 @@ data_type
 array_type
   = t:(numeric_type / character_string_type) __ LBRAKE __ RBRAKE __ LBRAKE __ RBRAKE {
     /* => data_type */
-    return { ...t, array: 'two' }
+    return { ...t, array: { dimension: 2 } }
   }
-  / t:(numeric_type / character_string_type) __ LBRAKE __ RBRAKE {
+  / t:(numeric_type / character_string_type) __ LBRAKE __ l:literal_numeric? __ RBRAKE {
     /* => data_type */
-    return { ...t, array: 'one' }
+    return { ...t, array: { dimension: 1,  length: [l] } }
+  }
+  / t:(numeric_type / character_string_type) __ KW_ARRAY {
+    /* => data_type */
+    return { ...t, array: { keyword: 'array' } }
   }
 
 boolean_type

diff --git a/src/column.js b/src/column.js
@@ -1,6 +1,6 @@
 import { constraintDefinitionToSQL } from './constrain'
 import { exprToSQL } from './expr'
-import { castToSQL } from './func'
+import { arrayDimensionToSymbol, castToSQL } from './func'
 import { tablesToSQL } from './tables'
 import {
   autoIncrementToSQL,
@@ -61,11 +61,17 @@ function columnRefToSQL(expr) {
 }
 
 function columnDataType(definition) {
+  if (!definition) return
   const { dataType, length, suffix, scale, expr } = definition || {}
   let result = dataType
   if (length != null) result += `(${[length, scale].filter(val => val != null).join(', ')})`
   if (suffix && suffix.length) result += ` ${suffix.join(' ')}`
   if (expr) result += exprToSQL(expr)
+  if (definition.array) {
+    const arrayExpr = arrayDimensionToSymbol(definition)
+    const space = /^\[.*\]$/.test(arrayExpr) ? '' : ' '
+    result += [space, arrayExpr].join('')
+  }
   return result
 }
 

diff --git a/src/func.js b/src/func.js
@@ -15,12 +15,16 @@ function anyValueFuncToSQL(stmt) {
 
 function arrayDimensionToSymbol(target) {
   if (!target || !target.array) return ''
-  switch (target.array) {
-    case 'one':
-      return '[]'
-    case 'two':
-      return '[][]'
+  const { keyword } = target.array
+  if (keyword) return toUpper(keyword)
+  const { dimension, length } = target.array
+  const result = []
+  for (let i = 0; i < dimension; i++) {
+    result.push('[')
+    if (length && length[i]) result.push(literalToSQL(length[i]))
+    result.push(']')
   }
+  return result.join('')
 }
 
 function castToSQL(expr) {
@@ -107,6 +111,7 @@ function lambdaToSQL(stmt) {
 
 export {
   anyValueFuncToSQL,
+  arrayDimensionToSymbol,
   castToSQL,
   extractFunToSQL,
   flattenFunToSQL,

diff --git a/test/postgres.spec.js b/test/postgres.spec.js
@@ -1514,6 +1514,13 @@ describe('Postgres', () => {
         'CREATE TABLE "Books" (price DECIMAL(10, 2) CHECK (Price > 0))'
       ]
     },
+    {
+      title: 'array data type',
+      sql: [
+        `CREATE TABLE "table_0" ("hi" INTEGER ARRAY); CREATE TABLE "table_1" ("hi" INTEGER[3]);`,
+        `CREATE TABLE "table_0" ("hi" INTEGER ARRAY) ; CREATE TABLE "table_1" ("hi" INTEGER[3])`
+      ]
+    },
   ]
   function neatlyNestTestedSQL(sqlList){
     sqlList.forEach(sqlInfo => {