Merge pull request #1761 from taozhi8833998/feat-set-to-pg

feat: support set to value in pg
taozhi8833998 · Jan 19, 2024 · d2a31a8 · d2a31a8
2 parents b042c3d + f416e21
commit d2a31a8
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 5 deletions.
diff --git a/pegjs/postgresql.pegjs b/pegjs/postgresql.pegjs
@@ -202,10 +202,13 @@
 }
 
 start
-  = __ n:(create_function_stmt / multiple_stmt) {
+  = b:('begin'i __ SEMICOLON) __ n:(create_function_stmt / multiple_stmt) __ c:('commit'i __ SEMICOLON) {
     // => multiple_stmt
+    n.ast.transactions = true
     return n
   }
+  / create_function_stmt
+  / multiple_stmt
 
 cmd_stmt
   = drop_stmt
@@ -5033,12 +5036,12 @@ proc_stmt
     }
 
 assign_stmt
-  = va:(var_decl / without_prefix_var_decl) __ s: (KW_ASSIGN / KW_ASSIGIN_EQUAL) __ e:proc_expr {
+  = va:(var_decl / without_prefix_var_decl) __ s:(KW_ASSIGN / KW_ASSIGIN_EQUAL / KW_TO) __ e:proc_expr {
     // => { type: 'assign'; left: var_decl | without_prefix_var_decl; symbol: ':=' | '='; right: proc_expr; }
     return {
       type: 'assign',
       left: va,
-      symbol: s,
+      symbol: Array.isArray(s) ? s[0] : s,
       right: e
     };
   }

diff --git a/src/sql.js b/src/sql.js
@@ -24,6 +24,7 @@ function goToSQL(stmt) {
 }
 
 export default function astToSQL(ast) {
-  if (ast.go === 'go') return goToSQL(ast)
-  return toSQL(ast)
+  const sql = ast.go === 'go' ? goToSQL(ast) : toSQL(ast)
+  if (ast.transactions) return `BEGIN;\n${sql};\nCOMMIT;`
+  return sql
 }
diff --git a/test/postgres.spec.js b/test/postgres.spec.js
@@ -1378,6 +1378,17 @@ describe('Postgres', () => {
         'SELECT "item" FROM "public"."orders" WHERE "ID" = 1 INTERSECT SELECT "sku" FROM "public"."inventory" WHERE "ID" = 1'
       ]
     },
+    {
+      title: 'set to in transactions',
+      sql: [
+        `BEGIN;
+        SET search_path TO ht_hyt;
+        COMMIT;`,
+        `BEGIN;
+SET search_path TO ht_hyt;
+COMMIT;`,
+      ]
+    }
   ]
   function neatlyNestTestedSQL(sqlList){
     sqlList.forEach(sqlInfo => {