In plpgsql, allow %TYPE and %ROWTYPE to be followed by array decoration.

This provides the useful ability to declare a variable that is an array
of the type of some other variable or some table column.

Quan Zongliang, Pavel Stehule

Discussion: https://postgr.es/m/ec4523e1-9e7e-f3ef-f9ce-bafd680ad6f6@yeah.net

Author: tglsfdc

doc/src/sgml/plpgsql.sgml

@@ -675,12 +675,14 @@ DECLARE
    <title>Copying Types</title>
 
 <synopsis>
-<replaceable>variable</replaceable>%TYPE
+<replaceable>name</replaceable> <replaceable>table</replaceable>.<replaceable>column</replaceable>%TYPE
+<replaceable>name</replaceable> <replaceable>variable</replaceable>%TYPE
 </synopsis>
 
    <para>
-    <literal>%TYPE</literal> provides the data type of a variable or
-    table column. You can use this to declare variables that will hold
+    <literal>%TYPE</literal> provides the data type of a table column
+    or a previously-declared <application>PL/pgSQL</application>
+    variable. You can use this to declare variables that will hold
     database values. For example, let's say you have a column named
     <literal>user_id</literal> in your <literal>users</literal>
     table. To declare a variable with the same data type as
@@ -690,6 +692,21 @@ user_id users.user_id%TYPE;
 </programlisting>
    </para>
 
+   <para>
+    It is also possible to write array decoration
+    after <literal>%TYPE</literal>, thereby creating a variable that holds
+    an array of the referenced type:
+<programlisting>
+user_ids users.user_id%TYPE[];
+user_ids users.user_id%TYPE ARRAY[4];  -- equivalent to the above
+</programlisting>
+    Just as when declaring table columns that are arrays, it doesn't
+    matter whether you write multiple bracket pairs or specific array
+    dimensions: <productname>PostgreSQL</productname> treats all arrays of
+    a given element type as the same type, regardless of dimensionality.
+    (See <xref linkend="arrays-declaration"/>.)
+   </para>
+
    <para>
     By using <literal>%TYPE</literal> you don't need to know the data
     type of the structure you are referencing, and most importantly,
@@ -739,6 +756,12 @@ user_id users.user_id%TYPE;
     <literal>%ROWTYPE</literal> is more portable.)
    </para>
 
+   <para>
+    As with <literal>%TYPE</literal>, <literal>%ROWTYPE</literal> can be
+    followed by array decoration to declare a variable that holds an array
+    of the referenced composite type.
+   </para>
+
    <para>
     Parameters to a function can be
     composite types (complete table rows). In that case, the

src/pl/plpgsql/src/expected/plpgsql_array.out

@@ -93,3 +93,82 @@ LINE 1: a.r[1] := 2
         ^
 QUERY:  a.r[1] := 2
 CONTEXT:  PL/pgSQL function inline_code_block line 2 at assignment
+--
+-- test of %type[] and %rowtype[] syntax
+--
+-- check supported syntax
+do $$
+declare
+  v int;
+  v1 v%type;
+  v2 v%type[];
+  v3 v%type[1];
+  v4 v%type[][];
+  v5 v%type[1][3];
+  v6 v%type array;
+  v7 v%type array[];
+  v8 v%type array[1];
+  v9 v%type array[1][1];
+  v10 pg_catalog.pg_class%rowtype[];
+begin
+  raise notice '%', pg_typeof(v1);
+  raise notice '%', pg_typeof(v2);
+  raise notice '%', pg_typeof(v3);
+  raise notice '%', pg_typeof(v4);
+  raise notice '%', pg_typeof(v5);
+  raise notice '%', pg_typeof(v6);
+  raise notice '%', pg_typeof(v7);
+  raise notice '%', pg_typeof(v8);
+  raise notice '%', pg_typeof(v9);
+  raise notice '%', pg_typeof(v10);
+end;
+$$;
+NOTICE:  integer
+NOTICE:  integer[]
+NOTICE:  integer[]
+NOTICE:  integer[]
+NOTICE:  integer[]
+NOTICE:  integer[]
+NOTICE:  integer[]
+NOTICE:  integer[]
+NOTICE:  integer[]
+NOTICE:  pg_class[]
+-- some types don't support arrays
+do $$
+declare
+  v pg_node_tree;
+  v1 v%type[];
+begin
+end;
+$$;
+ERROR:  could not find array type for data type pg_node_tree
+CONTEXT:  compilation of PL/pgSQL function "inline_code_block" near line 4
+-- check functionality
+do $$
+declare
+  v1 int;
+  v2 varchar;
+  a1 v1%type[];
+  a2 v2%type[];
+begin
+  v1 := 10;
+  v2 := 'Hi';
+  a1 := array[v1,v1];
+  a2 := array[v2,v2];
+  raise notice '% %', a1, a2;
+end;
+$$;
+NOTICE:  {10,10} {Hi,Hi}
+create table array_test_table(a int, b varchar);
+insert into array_test_table values(1, 'first'), (2, 'second');
+do $$
+declare tg array_test_table%rowtype[];
+begin
+  tg := array(select array_test_table from array_test_table);
+  raise notice '%', tg;
+  tg := array(select row(a,b) from array_test_table);
+  raise notice '%', tg;
+end;
+$$;
+NOTICE:  {"(1,first)","(2,second)"}
+NOTICE:  {"(1,first)","(2,second)"}

src/pl/plpgsql/src/pl_comp.c

@@ -2208,6 +2208,33 @@ build_datatype(HeapTuple typeTup, int32 typmod,
 	return typ;
 }
 
+/*
+ * Build an array type for the element type specified as argument.
+ */
+PLpgSQL_type *
+plpgsql_build_datatype_arrayof(PLpgSQL_type *dtype)
+{
+	Oid			array_typeid;
+
+	/*
+	 * If it's already an array type, use it as-is: Postgres doesn't do nested
+	 * arrays.
+	 */
+	if (dtype->typisarray)
+		return dtype;
+
+	array_typeid = get_array_type(dtype->typoid);
+	if (!OidIsValid(array_typeid))
+		ereport(ERROR,
+				(errcode(ERRCODE_UNDEFINED_OBJECT),
+				 errmsg("could not find array type for data type %s",
+						format_type_be(dtype->typoid))));
+
+	/* Note we inherit typmod and collation, if any, from the element type */
+	return plpgsql_build_datatype(array_typeid, dtype->atttypmod,
+								  dtype->collation, NULL);
+}
+
 /*
  *	plpgsql_recognize_err_condition
  *		Check condition name and translate it to SQLSTATE.

src/pl/plpgsql/src/pl_gram.y

@@ -757,8 +757,9 @@ decl_const		:
 decl_datatype	:
 					{
 						/*
-						 * If there's a lookahead token, read_datatype
-						 * should consume it.
+						 * If there's a lookahead token, read_datatype() will
+						 * consume it, and then we must tell bison to forget
+						 * it.
 						 */
 						$$ = read_datatype(yychar);
 						yyclearin;
@@ -2783,13 +2784,17 @@ read_sql_construct(int until,
 	return expr;
 }
 
+/*
+ * Read a datatype declaration, consuming the current lookahead token if any.
+ * Returns a PLpgSQL_type struct.
+ */
 static PLpgSQL_type *
 read_datatype(int tok)
 {
 	StringInfoData ds;
 	char	   *type_name;
 	int			startlocation;
-	PLpgSQL_type *result;
+	PLpgSQL_type *result = NULL;
 	int			parenlevel = 0;
 
 	/* Should only be called while parsing DECLARE sections */
@@ -2799,11 +2804,15 @@ read_datatype(int tok)
 	if (tok == YYEMPTY)
 		tok = yylex();
 
+	/* The current token is the start of what we'll pass to parse_datatype */
 	startlocation = yylloc;
 
 	/*
-	 * If we have a simple or composite identifier, check for %TYPE
-	 * and %ROWTYPE constructs.
+	 * If we have a simple or composite identifier, check for %TYPE and
+	 * %ROWTYPE constructs.  (Note that if plpgsql_parse_wordtype et al fail
+	 * to recognize the identifier, we'll fall through and pass the whole
+	 * string to parse_datatype, which will assuredly give an unhelpful
+	 * "syntax error".  Should we try to give a more specific error?)
 	 */
 	if (tok == T_WORD)
 	{
@@ -2815,18 +2824,10 @@ read_datatype(int tok)
 			tok = yylex();
 			if (tok_is_keyword(tok, &yylval,
 							   K_TYPE, "type"))
-			{
 				result = plpgsql_parse_wordtype(dtname);
-				if (result)
-					return result;
-			}
 			else if (tok_is_keyword(tok, &yylval,
 									K_ROWTYPE, "rowtype"))
-			{
 				result = plpgsql_parse_wordrowtype(dtname);
-				if (result)
-					return result;
-			}
 		}
 	}
 	else if (plpgsql_token_is_unreserved_keyword(tok))
@@ -2839,18 +2840,10 @@ read_datatype(int tok)
 			tok = yylex();
 			if (tok_is_keyword(tok, &yylval,
 							   K_TYPE, "type"))
-			{
 				result = plpgsql_parse_wordtype(dtname);
-				if (result)
-					return result;
-			}
 			else if (tok_is_keyword(tok, &yylval,
 									K_ROWTYPE, "rowtype"))
-			{
 				result = plpgsql_parse_wordrowtype(dtname);
-				if (result)
-					return result;
-			}
 		}
 	}
 	else if (tok == T_CWORD)
@@ -2863,21 +2856,56 @@ read_datatype(int tok)
 			tok = yylex();
 			if (tok_is_keyword(tok, &yylval,
 							   K_TYPE, "type"))
-			{
 				result = plpgsql_parse_cwordtype(dtnames);
-				if (result)
-					return result;
-			}
 			else if (tok_is_keyword(tok, &yylval,
 									K_ROWTYPE, "rowtype"))
-			{
 				result = plpgsql_parse_cwordrowtype(dtnames);
-				if (result)
-					return result;
-			}
 		}
 	}
 
+	/*
+	 * If we recognized a %TYPE or %ROWTYPE construct, see if it is followed
+	 * by array decoration: [ ARRAY ] [ '[' [ iconst ] ']' [ ... ] ]
+	 *
+	 * Like the core parser, we ignore the specific numbers and sizes of
+	 * dimensions; arrays of different dimensionality are still the same type
+	 * in Postgres.
+	 */
+	if (result)
+	{
+		bool		is_array = false;
+
+		tok = yylex();
+		if (tok_is_keyword(tok, &yylval,
+						   K_ARRAY, "array"))
+		{
+			is_array = true;
+			tok = yylex();
+		}
+		while (tok == '[')
+		{
+			is_array = true;
+			tok = yylex();
+			if (tok == ICONST)
+				tok = yylex();
+			if (tok != ']')
+				yyerror("syntax error, expected \"]\"");
+			tok = yylex();
+		}
+		plpgsql_push_back_token(tok);
+
+		if (is_array)
+			result = plpgsql_build_datatype_arrayof(result);
+
+		return result;
+	}
+
+	/*
+	 * Not %TYPE or %ROWTYPE, so scan to the end of the datatype declaration,
+	 * which could include typmod or array decoration.  We are not very picky
+	 * here, instead relying on parse_datatype to complain about garbage.  But
+	 * we must count parens to handle typmods within cursor_arg correctly.
+	 */
 	while (tok != ';')
 	{
 		if (tok == 0)

src/pl/plpgsql/src/plpgsql.h

@@ -1249,6 +1249,7 @@ extern PLpgSQL_type *plpgsql_parse_cwordrowtype(List *idents);
 extern PGDLLEXPORT PLpgSQL_type *plpgsql_build_datatype(Oid typeOid, int32 typmod,
 														Oid collation,
 														TypeName *origtypname);
+extern PLpgSQL_type *plpgsql_build_datatype_arrayof(PLpgSQL_type *dtype);
 extern PLpgSQL_variable *plpgsql_build_variable(const char *refname, int lineno,
 												PLpgSQL_type *dtype,
 												bool add2namespace);