Merge PR #168

lelit · Feb 7, 2025 · 8147f7e · 8147f7e
2 parents eed764e + 405e43a
commit 8147f7e
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 4 deletions.
diff --git a/pglast/printers/ddl.py b/pglast/printers/ddl.py
@@ -3363,16 +3363,16 @@ def variable_set_stmt(node, output):
             output.print_list(node.args, ',')
         else:
             output.print_name(node.name.split('.'))
-            output.write(' TO ')
+            if node.kind == vsk.VAR_SET_CURRENT:
+                output.write(' FROM CURRENT')
+            else:
+                output.write(' TO ')
         if node.kind == vsk.VAR_SET_VALUE:
             output.print_list(node.args)
         elif node.kind == vsk.VAR_SET_DEFAULT:
             output.write('DEFAULT')
         elif node.kind == vsk.VAR_SET_MULTI:
             pass
-        else:
-            raise NotImplementedError("SET statement of kind %s not implemented yet"
-                                      % node.kind)
 
 
 @node_printer(ast.VariableShowStmt)

diff --git a/tests/test_printers_prettification/ddl/create_function.sql b/tests/test_printers_prettification/ddl/create_function.sql
@@ -79,3 +79,16 @@ BEGIN ATOMIC
   INSERT INTO bar
   VALUES (2, 3);
 END
+
+CREATE FUNCTION funcb(somearg text)
+RETURNS TABLE(c1 int, c2 text)
+AS $$
+SELECT 1, 'Label 1'
+$$ language sql SET search_path FROM CURRENT
+=
+CREATE FUNCTION funcb(somearg text)
+RETURNS TABLE (c1 integer, c2 text)
+AS $$
+SELECT 1, 'Label 1'
+$$
+LANGUAGE sql SET search_path FROM CURRENT
diff --git a/tests/test_printers_roundtrip/ddl/create_function.sql b/tests/test_printers_roundtrip/ddl/create_function.sql
@@ -26,3 +26,9 @@ $$ language sql
 CREATE FUNCTION func() RETURNS SETOF montype AS $$ $$ language sql
 
 CREATE PROCEDURE func() AS $$ $$ language sql
+
+CREATE FUNCTION funcb(somearg text)
+RETURNS TABLE(c1 int, c2 text)
+AS $$
+SELECT 1, 'Label 1'
+$$ language sql SET search_path FROM CURRENT