Deparser: Add support for multi-statement CREATE PROCEDURE definitions (

#197)
pganalyze · Jul 8, 2023 · 1b508f8 · 1b508f8
1 parent f1b475d
commit 1b508f8
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 2 deletions.
diff --git a/src/postgres_deparse.c b/src/postgres_deparse.c
@@ -204,6 +204,7 @@ static void deparseExplainableStmt(StringInfo str, Node *node);
 static void deparseStmt(StringInfo str, Node *node);
 static void deparseValue(StringInfo str, union ValUnion *value, DeparseNodeContext context);
 
+
 // "any_name" in gram.y
 static void deparseAnyName(StringInfo str, List *parts)
 {
@@ -785,7 +786,7 @@ static void deparseCreateGenericOptions(StringInfo str, List *options)
 		if (lnext(options, lc))
 			appendStringInfoString(str, ", ");
 	}
-	appendStringInfoString(str, ") ");
+	appendStringInfoString(str, ")");
 }
 
 // "common_func_opt_item" in gram.y
@@ -4902,7 +4903,24 @@ static void deparseCreateFunctionStmt(StringInfo str, CreateFunctionStmt *create
 		else
 		{
 			appendStringInfoString(str, "BEGIN ATOMIC ");
-			deparseExprList(str, castNode(List, create_function_stmt->sql_body));
+			if (linitial(create_function_stmt->sql_body) != NULL)
+			{
+				List *body_stmt_list = castNode(List, linitial(create_function_stmt->sql_body));
+				foreach(lc, body_stmt_list)
+				{
+					if (IsA(lfirst(lc), ReturnStmt))
+					{
+						deparseReturnStmt(str, lfirst_node(ReturnStmt, lc));
+						appendStringInfoString(str, "; ");
+					}
+					else
+					{
+						deparseStmt(str, lfirst(lc));
+						appendStringInfoString(str, "; ");
+					}
+				}
+			}
+
 			appendStringInfoString(str, "END ");
 		}
 	}

diff --git a/test/deparse_tests.c b/test/deparse_tests.c
@@ -393,6 +393,10 @@ const char* tests[] = {
   "CREATE TABLE distributors (did int, name varchar(40), UNIQUE (name) WITH (fillfactor=70)) WITH (fillfactor=70)",
   "SHOW ALL",
   "ALTER TABLE ONLY public.\"Test 123\" ADD CONSTRAINT \"Test 123_pkey\" PRIMARY KEY (c1)",
+  "CREATE PROCEDURE insert_data(a int, b int) LANGUAGE sql BEGIN ATOMIC INSERT INTO tbl VALUES (a); INSERT INTO tbl VALUES (b); END",
+  "CREATE PROCEDURE do_nothing() LANGUAGE sql BEGIN ATOMIC END",
+  "CREATE PROCEDURE returns_one() LANGUAGE sql BEGIN ATOMIC RETURN 1; END",
+  "CREATE PROCEDURE updates_and_returns_one() LANGUAGE sql BEGIN ATOMIC UPDATE tbl SET a = 1; RETURN 1; END",
 };
 
 size_t testsLength = __LINE__ - 4;