Support SET ROLE

Author: nene

src/cst/Role.ts

@@ -5,7 +5,11 @@ import { NullLiteral, NumberLiteral, StringLiteral } from "./Literal";
 
 export type AllRoleNodes = AllRoleStatements | RoleOption | InDatabaseClause;
 
-export type AllRoleStatements = CreateRoleStmt | AlterRoleStmt | DropRoleStmt;
+export type AllRoleStatements =
+  | CreateRoleStmt
+  | AlterRoleStmt
+  | DropRoleStmt
+  | SetRoleStmt;
 
 export type RoleSpecification = Identifier | FuncCall;
 
@@ -113,3 +117,12 @@ export interface DropRoleStmt extends BaseNode {
   ifExistsKw?: [Keyword<"IF">, Keyword<"EXISTS">];
   names: ListExpr<RoleSpecification>;
 }
+
+// SET ROLE
+export interface SetRoleStmt extends BaseNode {
+  type: "set_role_stmt";
+  setKw: Keyword<"SET">;
+  scopeKw?: Keyword<"SESSION" | "LOCAL">;
+  roleKw: Keyword<"ROLE">;
+  name: Identifier | StringLiteral | Keyword<"NONE">;
+}

src/parser.pegjs

@@ -161,6 +161,7 @@ ddl_statement_postgres
   / create_role_stmt
   / alter_role_stmt
   / drop_role_stmt
+  / set_role_stmt
 
 dml_statement
   = compound_select_stmt
@@ -4556,6 +4557,17 @@ drop_role_stmt
     });
   }
 
+set_role_stmt
+  = kw:(SET __) scopeKw:((SESSION / LOCAL) __)? roleKw:(ROLE __) name:(NONE / ident / string_literal) {
+    return loc({
+      type: "set_role_stmt",
+      setKw: read(kw),
+      scopeKw: read(scopeKw),
+      roleKw: read(roleKw),
+      name,
+    });
+  }
+
 /**
  * ------------------------------------------------------------------------------------ *
  *                                                                                      *

src/showNode/role.ts

@@ -25,4 +25,8 @@ export const roleMap: FullTransformMap<string, AllRoleNodes> = {
   // DROP ROLE
   drop_role_stmt: (node) =>
     show([node.dropRoleKw, node.ifExistsKw, node.names]),
+
+  // SET ROLE
+  set_role_stmt: (node) =>
+    show([node.setKw, node.scopeKw, node.roleKw, node.name]),
 };

test/ddl/role.test.ts

@@ -176,6 +176,26 @@ describe("role", () => {
         testWc("DROP GROUP IF EXISTS mary, john");
       });
     });
+
+    describe("SET ROLE", () => {
+      it("supports SET ROLE name", () => {
+        testWc("SET ROLE paul");
+        testWc("SET SESSION ROLE paul");
+        testWc("SET LOCAL ROLE paul");
+      });
+
+      it("supports SET ROLE 'name'", () => {
+        testWc("SET ROLE 'paul'");
+        testWc("SET SESSION ROLE 'paul'");
+        testWc("SET LOCAL ROLE 'paul'");
+      });
+
+      it("supports SET ROLE NONE", () => {
+        testWc("SET ROLE NONE");
+        testWc("SET SESSION ROLE NONE");
+        testWc("SET LOCAL ROLE NONE");
+      });
+    });
   });
 
   notDialect("postgresql", () => {