parser: fix some set related issues (#657)

Author: sbdchd

crates/squawk_parser/src/grammar.rs

@@ -12986,15 +12986,21 @@ fn set(p: &mut Parser<'_>) -> CompletedMarker {
         if !p.eat(LOCAL_KW) && !config_value(p) {
             p.error(format!("expected config value, got {:?}", p.current()));
         }
+    } else if p.eat(CATALOG_KW) || p.eat(SCHEMA_KW) {
+        string_literal(p);
     // configuration_parameter { TO | = } { value | 'value' | DEFAULT }
     } else {
         // configuration_parameter
         path_name_ref(p);
-        // { TO | = }
-        let _ = p.eat(TO_KW) || p.expect(EQ);
-        // { value | 'value' | DEFAULT }
-        if !config_value(p) {
-            p.error(format!("expected config value, got {:?}", p.current()));
+        if p.eat(FROM_KW) {
+            p.expect(CURRENT_KW);
+        } else {
+            // { TO | = }
+            let _ = p.eat(TO_KW) || p.expect(EQ);
+            // { value | 'value' | DEFAULT }
+            if !config_value(p) {
+                p.error(format!("expected config value, got {:?}", p.current()));
+            }
         }
     }
     m.complete(p, SET)

crates/squawk_parser/tests/data/ok/schemas.sql

@@ -44,12 +44,23 @@ set search_path to myschema;
 
 set foo = bar;
 
+set schema 'foo';
+
+set xml option document;
+set xml option content;
+
+set role foo;
+
 set time zone 'America/Los_Angeles';
 set time zone default;
 set time zone local;
 
+set foo from current;
+set foo.bar from current;
+
 set foo = default;
 set foo to a, 10.0, 1, 'foo', true, false;
+set foo to default;
 
 -- operator
 -- binary

crates/squawk_parser/tests/data/ok/schemas_pg17.sql

@@ -0,0 +1 @@
+set catalog 'foo';

crates/squawk_parser/tests/snapshots/tests__schemas_ok.snap

@@ -415,6 +415,43 @@ SOURCE_FILE
     IDENT "bar"
   SEMICOLON ";"
   WHITESPACE "\n\n"
+  SET
+    SET_KW "set"
+    WHITESPACE " "
+    SCHEMA_KW "schema"
+    WHITESPACE " "
+    LITERAL
+      STRING "'foo'"
+  SEMICOLON ";"
+  WHITESPACE "\n\n"
+  SET
+    SET_KW "set"
+    WHITESPACE " "
+    XML_KW "xml"
+    WHITESPACE " "
+    OPTION_KW "option"
+    WHITESPACE " "
+    DOCUMENT_KW "document"
+  SEMICOLON ";"
+  WHITESPACE "\n"
+  SET
+    SET_KW "set"
+    WHITESPACE " "
+    XML_KW "xml"
+    WHITESPACE " "
+    OPTION_KW "option"
+    WHITESPACE " "
+    CONTENT_KW "content"
+  SEMICOLON ";"
+  WHITESPACE "\n\n"
+  SET_ROLE
+    SET_KW "set"
+    WHITESPACE " "
+    ROLE_KW "role"
+    WHITESPACE " "
+    IDENT "foo"
+  SEMICOLON ";"
+  WHITESPACE "\n\n"
   SET
     SET_KW "set"
     WHITESPACE " "
@@ -446,6 +483,37 @@ SOURCE_FILE
     LOCAL_KW "local"
   SEMICOLON ";"
   WHITESPACE "\n\n"
+  SET
+    SET_KW "set"
+    WHITESPACE " "
+    PATH
+      PATH_SEGMENT
+        NAME_REF
+          IDENT "foo"
+    WHITESPACE " "
+    FROM_KW "from"
+    WHITESPACE " "
+    CURRENT_KW "current"
+  SEMICOLON ";"
+  WHITESPACE "\n"
+  SET
+    SET_KW "set"
+    WHITESPACE " "
+    PATH
+      PATH
+        PATH_SEGMENT
+          NAME_REF
+            IDENT "foo"
+      DOT "."
+      PATH_SEGMENT
+        NAME_REF
+          IDENT "bar"
+    WHITESPACE " "
+    FROM_KW "from"
+    WHITESPACE " "
+    CURRENT_KW "current"
+  SEMICOLON ";"
+  WHITESPACE "\n\n"
   SET
     SET_KW "set"
     WHITESPACE " "
@@ -491,6 +559,19 @@ SOURCE_FILE
     LITERAL
       FALSE_KW "false"
   SEMICOLON ";"
+  WHITESPACE "\n"
+  SET
+    SET_KW "set"
+    WHITESPACE " "
+    PATH
+      PATH_SEGMENT
+        NAME_REF
+          IDENT "foo"
+    WHITESPACE " "
+    TO_KW "to"
+    WHITESPACE " "
+    DEFAULT_KW "default"
+  SEMICOLON ";"
   WHITESPACE "\n\n"
   COMMENT "-- operator"
   WHITESPACE "\n"

crates/squawk_parser/tests/snapshots/tests__schemas_pg17_ok.snap

@@ -0,0 +1,14 @@
+---
+source: crates/squawk_parser/tests/tests.rs
+input_file: crates/squawk_parser/tests/data/ok/schemas_pg17.sql
+---
+SOURCE_FILE
+  SET
+    SET_KW "set"
+    WHITESPACE " "
+    CATALOG_KW "catalog"
+    WHITESPACE " "
+    LITERAL
+      STRING "'foo'"
+  SEMICOLON ";"
+  WHITESPACE "\n"