SQL-standard function body

This adds support for writing CREATE FUNCTION and CREATE PROCEDURE
statements for language SQL with a function body that conforms to the
SQL standard and is portable to other implementations.

Instead of the PostgreSQL-specific AS $$ string literal $$ syntax,
this allows writing out the SQL statements making up the body
unquoted, either as a single statement:

    CREATE FUNCTION add(a integer, b integer) RETURNS integer
        LANGUAGE SQL
        RETURN a + b;

or as a block

    CREATE PROCEDURE insert_data(a integer, b integer)
    LANGUAGE SQL
    BEGIN ATOMIC
      INSERT INTO tbl VALUES (a);
      INSERT INTO tbl VALUES (b);
    END;

The function body is parsed at function definition time and stored as
expression nodes in a new pg_proc column prosqlbody.  So at run time,
no further parsing is required.

However, this form does not support polymorphic arguments, because
there is no more parse analysis done at call time.

Dependencies between the function and the objects it uses are fully
tracked.

A new RETURN statement is introduced.  This can only be used inside
function bodies.  Internally, it is treated much like a SELECT
statement.

psql needs some new intelligence to keep track of function body
boundaries so that it doesn't send off statements when it sees
semicolons that are inside a function body.

Tested-by: Jaime Casanova <jcasanov@systemguards.com.ec>
Reviewed-by: Julien Rouhaud <rjuju123@gmail.com>
Discussion: https://www.postgresql.org/message-id/flat/1c11f1eb-f00c-43b7-799d-2d44132c02d7@2ndquadrant.com

10 files changed, 46 insertions, 2 deletions

diff --git a/src/include/catalog/catversion.h b/src/include/catalog/catversion.h
index 53fd6c08b33..abd4bffff5f 100644
--- a/src/include/catalog/catversion.h
+++ b/src/include/catalog/catversion.h
@@ -53,6 +53,6 @@
  */
 
 /*							yyyymmddN */
-#define CATALOG_VERSION_NO	202104071
+#define CATALOG_VERSION_NO	202104072
 
 #endif
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 97d8238a997..6feaaa44597 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -3718,6 +3718,10 @@
   proname => 'pg_get_function_arg_default', provolatile => 's',
   prorettype => 'text', proargtypes => 'oid int4',
   prosrc => 'pg_get_function_arg_default' },
+{ oid => '9704', descr => 'function SQL body',
+  proname => 'pg_get_function_sqlbody', provolatile => 's',
+  prorettype => 'text', proargtypes => 'oid',
+  prosrc => 'pg_get_function_sqlbody' },
 
 { oid => '1686', descr => 'list of SQL keywords',
   proname => 'pg_get_keywords', procost => '10', prorows => '500',
diff --git a/src/include/catalog/pg_proc.h b/src/include/catalog/pg_proc.h
index 78f230894bd..8d58067d034 100644
--- a/src/include/catalog/pg_proc.h
+++ b/src/include/catalog/pg_proc.h
@@ -112,11 +112,14 @@ CATALOG(pg_proc,1255,ProcedureRelationId) BKI_BOOTSTRAP BKI_ROWTYPE_OID(81,Proce
 	Oid			protrftypes[1] BKI_DEFAULT(_null_) BKI_LOOKUP(pg_type);
 
 	/* procedure source text */
-	text		prosrc BKI_FORCE_NOT_NULL;
+	text		prosrc;
 
 	/* secondary procedure info (can be NULL) */
 	text		probin BKI_DEFAULT(_null_);
 
+	/* pre-parsed SQL function body */
+	pg_node_tree prosqlbody BKI_DEFAULT(_null_);
+
 	/* procedure-local GUC settings */
 	text		proconfig[1] BKI_DEFAULT(_null_);
 
@@ -194,6 +197,7 @@ extern ObjectAddress ProcedureCreate(const char *procedureName,
 									 Oid languageValidator,
 									 const char *prosrc,
 									 const char *probin,
+									 Node *prosqlbody,
 									 char prokind,
 									 bool security_definer,
 									 bool isLeakProof,
diff --git a/src/include/commands/defrem.h b/src/include/commands/defrem.h
index 339f29f4c85..6bce4d76fe5 100644
--- a/src/include/commands/defrem.h
+++ b/src/include/commands/defrem.h
@@ -65,9 +65,11 @@ extern void interpret_function_parameter_list(ParseState *pstate,
 											  Oid languageOid,
 											  ObjectType objtype,
 											  oidvector **parameterTypes,
+											  List **parameterTypes_list,
 											  ArrayType **allParameterTypes,
 											  ArrayType **parameterModes,
 											  ArrayType **parameterNames,
+											  List **inParameterNames_list,
 											  List **parameterDefaults,
 											  Oid *variadicArgType,
 											  Oid *requiredResultType);
diff --git a/src/include/executor/functions.h b/src/include/executor/functions.h
index c975a93661e..dcb8e18437f 100644
--- a/src/include/executor/functions.h
+++ b/src/include/executor/functions.h
@@ -20,6 +20,21 @@
 /* This struct is known only within executor/functions.c */
 typedef struct SQLFunctionParseInfo *SQLFunctionParseInfoPtr;
 
+/*
+ * Data structure needed by the parser callback hooks to resolve parameter
+ * references during parsing of a SQL function's body.  This is separate from
+ * SQLFunctionCache since we sometimes do parsing separately from execution.
+ */
+typedef struct SQLFunctionParseInfo
+{
+	char	   *fname;			/* function's name */
+	int			nargs;			/* number of input arguments */
+	Oid		   *argtypes;		/* resolved types of input arguments */
+	char	  **argnames;		/* names of input arguments; NULL if none */
+	/* Note that argnames[i] can be NULL, if some args are unnamed */
+	Oid			collation;		/* function's input collation, if known */
+}			SQLFunctionParseInfo;
+
 extern Datum fmgr_sql(PG_FUNCTION_ARGS);
 
 extern SQLFunctionParseInfoPtr prepare_sql_fn_parse_info(HeapTuple procedureTuple,
diff --git a/src/include/fe_utils/psqlscan_int.h b/src/include/fe_utils/psqlscan_int.h
index 2bc3cc4ae73..91d7d4d5c6c 100644
--- a/src/include/fe_utils/psqlscan_int.h
+++ b/src/include/fe_utils/psqlscan_int.h
@@ -114,6 +114,8 @@ typedef struct PsqlScanStateData
 	int			paren_depth;	/* depth of nesting in parentheses */
 	int			xcdepth;		/* depth of nesting in slash-star comments */
 	char	   *dolqstart;		/* current $foo$ quote start string */
+	int			identifier_count;	/* identifiers since start of statement */
+	int			begin_depth;	/* depth of begin/end routine body blocks */
 
 	/*
 	 * Callback functions provided by the program making use of the lexer,
diff --git a/src/include/nodes/nodes.h b/src/include/nodes/nodes.h
index 2051abbbf92..d9e417bcd70 100644
--- a/src/include/nodes/nodes.h
+++ b/src/include/nodes/nodes.h
@@ -321,6 +321,7 @@ typedef enum NodeTag
 	T_DeleteStmt,
 	T_UpdateStmt,
 	T_SelectStmt,
+	T_ReturnStmt,
 	T_PLAssignStmt,
 	T_AlterTableStmt,
 	T_AlterTableCmd,
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index acc093d6e0f..7a44bccdd3b 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -140,6 +140,8 @@ typedef struct Query
 	bool		hasForUpdate;	/* FOR [KEY] UPDATE/SHARE was specified */
 	bool		hasRowSecurity; /* rewriter has applied some RLS policy */
 
+	bool		isReturn;		/* is a RETURN statement */
+
 	List	   *cteList;		/* WITH list (of CommonTableExpr's) */
 
 	List	   *rtable;			/* list of range table entries */
@@ -1721,6 +1723,16 @@ typedef struct SetOperationStmt
 } SetOperationStmt;
 
 
+/*
+ * RETURN statement (inside SQL function body)
+ */
+typedef struct ReturnStmt
+{
+	NodeTag		type;
+	Node	   *returnval;
+} ReturnStmt;
+
+
 /* ----------------------
  *		PL/pgSQL Assignment Statement
  *
@@ -2924,6 +2936,7 @@ typedef struct CreateFunctionStmt
 	List	   *parameters;		/* a list of FunctionParameter */
 	TypeName   *returnType;		/* the return type */
 	List	   *options;		/* a list of DefElem */
+	Node	   *sql_body;
 } CreateFunctionStmt;
 
 typedef enum FunctionParameterMode
diff --git a/src/include/parser/kwlist.h b/src/include/parser/kwlist.h
index d3ecd72e72e..f836acf876a 100644
--- a/src/include/parser/kwlist.h
+++ b/src/include/parser/kwlist.h
@@ -49,6 +49,7 @@ PG_KEYWORD("assertion", ASSERTION, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("assignment", ASSIGNMENT, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("asymmetric", ASYMMETRIC, RESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("at", AT, UNRESERVED_KEYWORD, BARE_LABEL)
+PG_KEYWORD("atomic", ATOMIC, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("attach", ATTACH, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("attribute", ATTRIBUTE, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("authorization", AUTHORIZATION, TYPE_FUNC_NAME_KEYWORD, BARE_LABEL)
@@ -349,6 +350,7 @@ PG_KEYWORD("replica", REPLICA, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("reset", RESET, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("restart", RESTART, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("restrict", RESTRICT, UNRESERVED_KEYWORD, BARE_LABEL)
+PG_KEYWORD("return", RETURN, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("returning", RETURNING, RESERVED_KEYWORD, AS_LABEL)
 PG_KEYWORD("returns", RETURNS, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("revoke", REVOKE, UNRESERVED_KEYWORD, BARE_LABEL)
diff --git a/src/include/tcop/tcopprot.h b/src/include/tcop/tcopprot.h
index 241e7c99614..968345404e5 100644
--- a/src/include/tcop/tcopprot.h
+++ b/src/include/tcop/tcopprot.h
@@ -44,6 +44,7 @@ typedef enum
 extern PGDLLIMPORT int log_statement;
 
 extern List *pg_parse_query(const char *query_string);
+extern List *pg_rewrite_query(Query *query);
 extern List *pg_analyze_and_rewrite(RawStmt *parsetree,
 									const char *query_string,
 									Oid *paramTypes, int numParams,