From 5b7b81c2932c9b708dd7e80eaa3cbebf9e9d514b Mon Sep 17 00:00:00 2001
From: Harry Brundage <harry.brundage@gmail.com>
Date: Mon, 8 Jan 2024 10:19:24 -0500
Subject: [PATCH] Add support for SET LOCAL and SET SESSION statements (#147)

---
 src/syntax/ast.ts                    |  1 +
 src/syntax/base.ne                   |  1 +
 src/syntax/simple-statements.ne      |  7 ++++---
 src/syntax/simple-statements.spec.ts | 20 ++++++++++++++++++++
 src/to-sql.ts                        |  6 +++++-
 5 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/src/syntax/ast.ts b/src/syntax/ast.ts
index b307ce8..e405263 100644
--- a/src/syntax/ast.ts
+++ b/src/syntax/ast.ts
@@ -1002,6 +1002,7 @@ export interface ExprWhen extends PGNode {
 export interface SetGlobalStatement extends PGNode {
     type: 'set';
     variable: Name;
+    scope?: string;
     set: SetGlobalValue;
 }
 export interface SetTimezone extends PGNode {
diff --git a/src/syntax/base.ne b/src/syntax/base.ne
index 5125418..68c50fa 100644
--- a/src/syntax/base.ne
+++ b/src/syntax/base.ne
@@ -235,6 +235,7 @@ kw_interval -> %word {% notReservedKw('interval')  %}
 kw_hour -> %word {% notReservedKw('hour')  %}
 kw_minute -> %word {% notReservedKw('minute')  %}
 kw_local -> %word {% notReservedKw('local')  %}
+kw_session -> %word {% notReservedKw('session')  %}
 kw_prepare -> %word {% notReservedKw('prepare')  %}
 kw_deallocate -> %word {% notReservedKw('deallocate')  %}
 kw_raise -> %word {% notReservedKw('raise')  %}
diff --git a/src/syntax/simple-statements.ne b/src/syntax/simple-statements.ne
index 7078f64..7677bbf 100644
--- a/src/syntax/simple-statements.ne
+++ b/src/syntax/simple-statements.ne
@@ -49,10 +49,11 @@ simplestatements_set_names -> kw_names simplestatements_set_names_val {% x => tr
 simplestatements_set_names_val
     -> (string) {% x => track(x, { type: 'value', value: unwrap(x[0]) }) %}
 
-simplestatements_set_simple -> ident (%op_eq | %kw_to) simplestatements_set_val {% x  => track(x, {
+simplestatements_set_simple -> (kw_local | kw_session):? ident (%op_eq | %kw_to) simplestatements_set_val {% x  => track(x, {
         type: 'set',
-        variable: asName(x[0]),
-        set: unbox(x[2]),
+        variable: asName(x[1]),
+        scope: unwrap(x[0])?.toLowerCase(),
+        set: unbox(x[3]),
     }) %}
 
 simplestatements_set_val
diff --git a/src/syntax/simple-statements.spec.ts b/src/syntax/simple-statements.spec.ts
index 46b5ccf..25dfcd8 100644
--- a/src/syntax/simple-statements.spec.ts
+++ b/src/syntax/simple-statements.spec.ts
@@ -67,6 +67,26 @@ describe('Simple statements', () => {
         }
     })
 
+    checkStatement(`SET LOCAL lock_timeout = '50ms'`, {
+        type: 'set',
+        variable: { name: 'lock_timeout' },
+        scope: 'local',
+        set: {
+            type: 'value',
+            value: '50ms',
+        }
+    })
+
+    checkStatement(`SET SESSION lock_timeout = '50ms'`, {
+        type: 'set',
+        variable: { name: 'lock_timeout' },
+        scope: 'session',
+        set: {
+            type: 'value',
+            value: '50ms',
+        }
+    })
+
     checkStatement(`SET TIME ZONE INTERVAL '+00:00' HOUR TO MINUTE`, {
         type: 'set timezone',
         to: {
diff --git a/src/to-sql.ts b/src/to-sql.ts
index 89acb14..158e7a4 100644
--- a/src/to-sql.ts
+++ b/src/to-sql.ts
@@ -809,7 +809,11 @@ const visitor = astVisitor<IAstFullVisitor>(m => ({
 
 
     setGlobal: g => {
-        ret.push('SET ', name(g.variable), ' = ');
+        ret.push('SET ')
+        if (g.scope) {
+            ret.push(g.scope.toUpperCase() + ' ');
+        }
+        ret.push(name(g.variable), ' = ');
         visitSetVal(g.set);
     },