From 1b508f86667ffa322e7bc695a2a9e8a58a0c30b2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mehmet=20Emin=20Karaka=C5=9F?= <emin100@gmail.com>
Date: Fri, 7 Jul 2023 17:56:31 -0700
Subject: [PATCH] Deparser: Add support for multi-statement CREATE PROCEDURE
 definitions (#197)

---
 src/postgres_deparse.c | 22 ++++++++++++++++++++--
 test/deparse_tests.c   |  4 ++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/src/postgres_deparse.c b/src/postgres_deparse.c
index c24533cc..054500ed 100644
--- 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
index 106f2349..b94c94da 100644
--- 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;