Repository: tobymao/sqlglot

Issue:
Failed to parse SET statement in Postgres CREATE FUNCTION
**Fully reproducible code snippet**

```py
from sqlglot import parse

parse(
    """
CREATE FUNCTION add(integer, integer) RETURNS integer
    SET search_path TO 'public'
    AS 'select $1 + $2;'
    LANGUAGE SQL
    IMMUTABLE;
""",
    dialect="postgres",
)
```

Raises

```
Traceback (most recent call last):
  File "<...>/create-function-repro.py", line 3, in <module>
    parse(
  File "<...>/python3.9/site-packages/sqlglot/__init__.py", line 86, in parse
    return Dialect.get_or_raise(read or dialect).parse(sql, **opts)
  File "<...>/python3.9/site-packages/sqlglot/dialects/dialect.py", line 442, in parse
    return self.parser(**opts).parse(self.tokenize(sql), sql)
  File "<...>/python3.9/site-packages/sqlglot/parser.py", line 1026, in parse
    return self._parse(
  File "<...>/python3.9/site-packages/sqlglot/parser.py", line 1095, in _parse
    self.raise_error("Invalid expression / Unexpected token")
  File "<...>/python3.9/site-packages/sqlglot/parser.py", line 1136, in raise_error
    raise error
sqlglot.errors.ParseError: Invalid expression / Unexpected token. Line 3, Col: 28.
  
CREATE FUNCTION add(integer, integer) RETURNS integer
    SET search_path TO 'public'
    AS 'select $1 + $2;'
    LANGUAGE SQL
    IMMUTABLE;
```

**Official Documentation**

- https://www.postgresql.org/docs/current/sql-createfunction.html

Notice the allowed statement `SET configuration_parameter { TO value | = value | FROM CURRENT }` in `CREATE FUNCTION`.

Unrelated, but `STRICT/RETURNS NULL ON NULL INPUT` also fails, thus I omitted it from the example above.


Discussion:

Patch:
diff --git a/sqlglot/dialects/postgres.py b/sqlglot/dialects/postgres.py
index 1ca0a781..6ee3bdec 100644
--- a/sqlglot/dialects/postgres.py
+++ b/sqlglot/dialects/postgres.py
@@ -282,6 +282,11 @@ class Postgres(Dialect):
         VAR_SINGLE_TOKENS = {"$"}
 
     class Parser(parser.Parser):
+        PROPERTY_PARSERS = {
+            **parser.Parser.PROPERTY_PARSERS,
+            "SET": lambda self: self.expression(exp.SetConfigProperty, this=self._parse_set()),
+        }
+
         FUNCTIONS = {
             **parser.Parser.FUNCTIONS,
             "DATE_TRUNC": parse_timestamp_trunc,
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 0e6608e2..98114fbd 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -2281,6 +2281,10 @@ class SetProperty(Property):
     arg_types = {"multi": True}
 
 
+class SetConfigProperty(Property):
+    arg_types = {"this": True}
+
+
 class SettingsProperty(Property):
     arg_types = {"expressions": True}
 
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index aa78aef2..75a61c22 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -97,6 +97,7 @@ class Generator:
         exp.ReturnsProperty: lambda self, e: self.naked_property(e),
         exp.SampleProperty: lambda self, e: f"SAMPLE BY {self.sql(e, 'this')}",
         exp.SetProperty: lambda self, e: f"{'MULTI' if e.args.get('multi') else ''}SET",
+        exp.SetConfigProperty: lambda self, e: self.sql(e, "this"),
         exp.SettingsProperty: lambda self, e: f"SETTINGS{self.seg('')}{(self.expressions(e))}",
         exp.SqlReadWriteProperty: lambda self, e: e.name,
         exp.SqlSecurityProperty: lambda self, e: f"SQL SECURITY {'DEFINER' if e.args.get('definer') else 'INVOKER'}",
@@ -355,6 +356,7 @@ class Generator:
         exp.Set: exp.Properties.Location.POST_SCHEMA,
         exp.SettingsProperty: exp.Properties.Location.POST_SCHEMA,
         exp.SetProperty: exp.Properties.Location.POST_CREATE,
+        exp.SetConfigProperty: exp.Properties.Location.POST_SCHEMA,
         exp.SortKeyProperty: exp.Properties.Location.POST_SCHEMA,
         exp.SqlReadWriteProperty: exp.Properties.Location.POST_SCHEMA,
         exp.SqlSecurityProperty: exp.Properties.Location.POST_CREATE,


Test patch:
diff --git a/tests/dialects/test_postgres.py b/tests/dialects/test_postgres.py
index 7aafa37b..da3a2065 100644
--- a/tests/dialects/test_postgres.py
+++ b/tests/dialects/test_postgres.py
@@ -597,6 +597,8 @@ class TestPostgres(Validator):
         self.validate_identity("CREATE TABLE cities_partdef PARTITION OF cities DEFAULT")
         self.validate_identity("CREATE TABLE t (c CHAR(2) UNIQUE NOT NULL) INHERITS (t1)")
         self.validate_identity("CREATE TABLE s.t (c CHAR(2) UNIQUE NOT NULL) INHERITS (s.t1, s.t2)")
+        self.validate_identity("CREATE FUNCTION x(INT) RETURNS INT SET search_path = 'public'")
+        self.validate_identity("CREATE FUNCTION x(INT) RETURNS INT SET foo FROM CURRENT")
         self.validate_identity(
             "CREATE CONSTRAINT TRIGGER my_trigger AFTER INSERT OR DELETE OR UPDATE OF col_a, col_b ON public.my_table DEFERRABLE INITIALLY DEFERRED FOR EACH ROW EXECUTE FUNCTION do_sth()"
         )
@@ -642,16 +644,23 @@ class TestPostgres(Validator):
         self.validate_identity(
             "DELETE FROM event USING sales AS s WHERE event.eventid = s.eventid RETURNING a"
         )
-        self.validate_identity(
-            "CREATE TABLE test (x TIMESTAMP WITHOUT TIME ZONE[][])",
-            "CREATE TABLE test (x TIMESTAMP[][])",
-        )
         self.validate_identity(
             "CREATE UNLOGGED TABLE foo AS WITH t(c) AS (SELECT 1) SELECT * FROM (SELECT c AS c FROM t) AS temp"
         )
         self.validate_identity(
             "WITH t(c) AS (SELECT 1) SELECT * INTO UNLOGGED foo FROM (SELECT c AS c FROM t) AS temp"
         )
+        self.validate_identity(
+            "CREATE FUNCTION add(INT, INT) RETURNS INT SET search_path TO 'public' AS 'select $1 + $2;' LANGUAGE SQL IMMUTABLE"
+        )
+        self.validate_identity(
+            "CREATE FUNCTION x(INT) RETURNS INT SET search_path TO 'public'",
+            "CREATE FUNCTION x(INT) RETURNS INT SET search_path = 'public'",
+        )
+        self.validate_identity(
+            "CREATE TABLE test (x TIMESTAMP WITHOUT TIME ZONE[][])",
+            "CREATE TABLE test (x TIMESTAMP[][])",
+        )
 
         self.validate_all(
             "CREATE OR REPLACE FUNCTION function_name (input_a character varying DEFAULT NULL::character varying)",