From a5be0533458f90ca62a30d6bd72047724477f44a Mon Sep 17 00:00:00 2001
From: taozhi8833998 <taozhi8833998@163.com>
Date: Wed, 17 Jan 2024 20:13:46 +0800
Subject: [PATCH 1/2] feat: support set to value in pg

---
 pegjs/postgresql.pegjs | 4 ++--
 test/postgres.spec.js  | 7 +++++++
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/pegjs/postgresql.pegjs b/pegjs/postgresql.pegjs
index 6d3e7b8b..ffa4fb79 100644
--- a/pegjs/postgresql.pegjs
+++ b/pegjs/postgresql.pegjs
@@ -5033,12 +5033,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/test/postgres.spec.js b/test/postgres.spec.js
index 2c180a32..d166e33d 100644
--- a/test/postgres.spec.js
+++ b/test/postgres.spec.js
@@ -1378,6 +1378,13 @@ describe('Postgres', () => {
         'SELECT "item" FROM "public"."orders" WHERE "ID" = 1 INTERSECT SELECT "sku" FROM "public"."inventory" WHERE "ID" = 1'
       ]
     },
+    {
+      title: 'set to',
+      sql: [
+        'SET search_path TO ht_hyt',
+        'SET search_path TO ht_hyt'
+      ]
+    }
   ]
   function neatlyNestTestedSQL(sqlList){
     sqlList.forEach(sqlInfo => {

From f416e213dfef14cbd62ffc6b468ab593e4a1ed73 Mon Sep 17 00:00:00 2001
From: taozhi8833998 <taozhi8833998@163.com>
Date: Thu, 18 Jan 2024 09:13:06 +0800
Subject: [PATCH 2/2] feat: support transactions in pg

---
 pegjs/postgresql.pegjs |  5 ++++-
 src/sql.js             |  5 +++--
 test/postgres.spec.js  | 12 ++++++++----
 3 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/pegjs/postgresql.pegjs b/pegjs/postgresql.pegjs
index ffa4fb79..244e0baf 100644
--- 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
diff --git a/src/sql.js b/src/sql.js
index b066e53d..6e850e03 100644
--- 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
index d166e33d..3b566138 100644
--- a/test/postgres.spec.js
+++ b/test/postgres.spec.js
@@ -1379,10 +1379,14 @@ describe('Postgres', () => {
       ]
     },
     {
-      title: 'set to',
-      sql: [
-        'SET search_path TO ht_hyt',
-        'SET search_path TO ht_hyt'
+      title: 'set to in transactions',
+      sql: [
+        `BEGIN;
+        SET search_path TO ht_hyt;
+        COMMIT;`,
+        `BEGIN;
+SET search_path TO ht_hyt;
+COMMIT;`,
       ]
     }
   ]