Merge pull request taozhi8833998#173 from taozhi8833998/feature-pg-returning

taozhi8833998 · web-flow · commit 792cf2da5a59 · 2019-12-24T11:35:35.000+08:00
Feature pg returning
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "node-sql-parser",
-  "version": "1.6.16",
+  "version": "1.6.17",
   "description": "simple node sql parser",
   "main": "index.js",
   "types": "index.d.ts",
diff --git a/pegjs/mysql.pegjs b/pegjs/mysql.pegjs
@@ -2102,7 +2102,6 @@ var_decl
 without_prefix_var_decl
   = name:ident_name m:mem_chain {
     //push for analysis
-    console.log('name===', name)
     varList.push(name);
     return {
       type: 'var',
diff --git a/pegjs/postgresql.pegjs b/pegjs/postgresql.pegjs
@@ -1074,7 +1074,8 @@ update_stmt
     t:table_ref_list __
     KW_SET       __
     l:set_list   __
-    w:where_clause? {
+    w:where_clause? __
+    r:returning_stmt? {
       if (t) t.forEach(tableInfo => {
         const { db, as, table } = tableInfo
         tableList.add(`update::${db}::${table}`)
@@ -1089,7 +1090,8 @@ update_stmt
           type: 'update',
           table: t,
           set: l,
-          where: w
+          where: w,
+          returning: r,
         }
       };
     }
@@ -1137,7 +1139,13 @@ set_item
   = tbl:(ident __ DOT)? __ c:column __ '=' __ v:additive_expr {
       return { column: c, value: v, table: tbl && tbl[0] };
     }
-
+returning_stmt
+  = k:KW_RETURNING __ c:column_ref_list {
+    return {
+      type: k && k.toLowerCase() || 'returning',
+      columns: c
+    }
+  }
 insert_value_clause
   = value_clause
   / select_stmt_nake
@@ -1155,7 +1163,8 @@ replace_insert_stmt
     KW_INTO?                 __
     t:table_name  __
     p:insert_partition? __ LPAREN __ c:column_list  __ RPAREN __
-    v:insert_value_clause {
+    v:insert_value_clause __
+    r:returning_stmt? {
       if (t) {
         tableList.add(`insert::${t.db}::${t.table}`)
         t.as = null
@@ -1173,6 +1182,7 @@ replace_insert_stmt
           columns: c,
           values: v,
           partition: p,
+          returning: r,
         }
       };
     }
@@ -1182,7 +1192,8 @@ insert_no_columns_stmt
     KW_INTO                 __
     t:table_name  __
     p:insert_partition? __
-    v:insert_value_clause {
+    v:insert_value_clause __
+    r:returning_stmt? {
       if (t) {
         tableList.add(`insert::${t.db}::${t.table}`)
         columnList.add(`insert::${t.table}::(.*)`);
@@ -1197,6 +1208,7 @@ insert_no_columns_stmt
           columns: null,
           values: v,
           partition: p,
+          returning: r,
         }
       };
     }
@@ -1783,6 +1795,7 @@ KW_DELETE   = "DELETE"i     !ident_start
 KW_INSERT   = "INSERT"i     !ident_start
 KW_RECURSIVE= "RECURSIVE"   !ident_start
 KW_REPLACE  = "REPLACE"i    !ident_start
+KW_RETURNING  = "RETURNING"i    !ident_start { return 'RETURNING' }
 KW_RENAME   = "RENAME"i     !ident_start
 KW_IGNORE   = "IGNORE"i     !ident_start
 KW_EXPLAIN  = "EXPLAIN"i    !ident_start
diff --git a/src/insert.js b/src/insert.js
@@ -1,6 +1,6 @@
 import { tablesToSQL } from './tables'
 import { exprToSQL } from './expr'
-import { identifierToSql, commonOptionConnector, hasVal, toUpper } from './util'
+import { identifierToSql, commonOptionConnector, hasVal, toUpper, returningToSQL } from './util'
 import { selectToSQL } from './select'
 
 /**
@@ -34,11 +34,13 @@ function insertToSQL(stmt) {
     values,
     where,
     partition,
+    returning,
   } = stmt
   const clauses = ['INSERT', toUpper(prefix), tablesToSQL(table), partitionToSQL(partition)]
   if (Array.isArray(columns)) clauses.push(`(${columns.map(identifierToSql).join(', ')})`)
   clauses.push(commonOptionConnector(Array.isArray(values) ? 'VALUES' : '', valuesToSQL, values))
   clauses.push(commonOptionConnector('WHERE', exprToSQL, where))
+  clauses.push(returningToSQL(returning))
   return clauses.filter(hasVal).join(' ')
 }
 
diff --git a/src/update.js b/src/update.js
@@ -1,6 +1,6 @@
 import { tablesToSQL } from './tables'
 import { exprToSQL } from './expr'
-import { hasVal, identifierToSql, commonOptionConnector } from './util'
+import { hasVal, identifierToSql, commonOptionConnector, returningToSQL } from './util'
 
 /**
  * @param {Array} sets
@@ -20,12 +20,13 @@ function setToSQL(sets) {
 }
 
 function updateToSQL(stmt) {
-  const { table, set, where } = stmt
+  const { table, set, where, returning } = stmt
   const clauses = [
     'UPDATE',
     tablesToSQL(table),
     commonOptionConnector('SET', setToSQL, set),
     commonOptionConnector('WHERE', exprToSQL, where),
+    returningToSQL(returning),
   ]
   return clauses.filter(hasVal).join(' ')
 }
diff --git a/src/util.js b/src/util.js
@@ -209,6 +209,12 @@ function commentToSQL(comment) {
   return result.join(' ')
 }
 
+function returningToSQL(returning) {
+  if (!returning) return ''
+  const { columns } = returning
+  return ['RETURNING', columns.map(columnRefToSQL).filter(hasVal).join(', ')].join(' ')
+}
+
 export {
   commonOptionConnector,
   connector,
@@ -222,6 +228,7 @@ export {
   literalToSQL,
   identifierToSql,
   replaceParams,
+  returningToSQL,
   hasVal,
   setParserOpt,
   toUpper,
diff --git a/test/insert.spec.js b/test/insert.spec.js
@@ -126,4 +126,11 @@ describe('insert', () => {
       const backSQL = parser.sqlify(ast)
       expect(backSQL).to.be.equal("INSERT OVERWRITE TABLE `account` PARTITION(`date`, `id`) SELECT * FROM `tmp`")
     })
+
+    it('should support parse pg insert returning', () => {
+      const sql = 'INSERT into account (date, id) values("2019-12-23", 123) returning id'
+      const ast = parser.astify(sql, { database: 'postgresql' })
+      const backSQL = parser.sqlify(ast)
+      expect(backSQL).to.be.equal("INSERT INTO `account` (`date`, `id`) VALUES ('2019-12-23',123) RETURNING `id`")
+    })
 });
diff --git a/test/update.spec.js b/test/update.spec.js
@@ -193,4 +193,11 @@ describe('update', () => {
       }
     });
   })
+
+  it('should support parse pg update returning', () => {
+    const sql = 'update account set id = 1 where name = "abc" returning id'
+    const ast = parser.astify(sql, { database: 'postgresql' })
+    const backSQL = parser.sqlify(ast)
+    expect(backSQL).to.be.equal("UPDATE `account` SET `id` = 1 WHERE `name` = 'abc' RETURNING `id`")
+  })
 });

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`{`
`2`	`2`	`"name": "node-sql-parser",`
`3`		`- "version": "1.6.16",`
	`3`	`+ "version": "1.6.17",`
`4`	`4`	`"description": "simple node sql parser",`
`5`	`5`	`"main": "index.js",`
`6`	`6`	`"types": "index.d.ts",`