Mock plpgsql_build_datatype_arrayof with special type handling

pganalyze · Oct 15, 2024 · f4a36d5 · f4a36d5
1 parent 97ce3a3
commit f4a36d5
Show file tree

Hide file tree

Showing 3 changed files with 147 additions and 2 deletions.
diff --git a/scripts/extract_source.rb b/scripts/extract_source.rb
@@ -576,7 +576,44 @@ def write_out
 	return typ;
 }
 ))
-runner.mock('plpgsql_build_datatype_arrayof', 'PLpgSQL_type * plpgsql_build_datatype_arrayof(PLpgSQL_type *dtype) { PLpgSQL_type *typ; typ = (PLpgSQL_type *) palloc0(sizeof(PLpgSQL_type)); typ->typname = pstrdup("UNKNOWN"); typ->ttype = PLPGSQL_TTYPE_SCALAR; return typ;  }')
+runner.mock('plpgsql_build_datatype_arrayof', %(
+PLpgSQL_type * plpgsql_build_datatype_arrayof(PLpgSQL_type *dtype)
+{
+	if (dtype->typisarray)
+		return dtype;
+
+	PLpgSQL_type *array_type;
+	array_type = (PLpgSQL_type *) palloc0(sizeof(PLpgSQL_type));
+
+	array_type->ttype = PLPGSQL_TTYPE_REC;
+	array_type->atttypmod = dtype->atttypmod;
+	array_type->collation = dtype->collation;
+
+	array_type->typisarray = true;
+
+	switch(dtype->typoid)
+	{
+		case BOOLOID:
+			array_type->typoid = BOOLARRAYOID;
+			array_type->typname = pstrdup("boolean[]");
+			break;
+		case INT4OID:
+			array_type->typoid = INT4ARRAYOID;
+			array_type->typname = pstrdup("integer[]");
+			break;
+		case TEXTOID:
+			array_type->typoid = TEXTARRAYOID;
+			array_type->typname = pstrdup("text[]");
+			break;
+		default:
+			array_type->typname = pstrdup("UNKNOWN");
+			break;
+	}
+	array_type->typoid = dtype->typoid;
+
+	return array_type;
+}
+))
 runner.mock('parse_datatype', 'static PLpgSQL_type * parse_datatype(const char *string, int location) { PLpgSQL_type *typ; typ = (PLpgSQL_type *) palloc0(sizeof(PLpgSQL_type)); typ->typname = pstrdup(string); typ->ttype = strcmp(string, "RECORD") == 0 ? PLPGSQL_TTYPE_REC : PLPGSQL_TTYPE_SCALAR; return typ; }')
 runner.mock('get_collation_oid', 'Oid get_collation_oid(List *name, bool missing_ok) { return -1; }')
 runner.mock('plpgsql_parse_wordtype', 'PLpgSQL_type * plpgsql_parse_wordtype(char *ident) { return NULL; }')

diff --git a/src/postgres/src_pl_plpgsql_src_pl_comp.c b/src/postgres/src_pl_plpgsql_src_pl_comp.c
@@ -940,7 +940,44 @@ PLpgSQL_type * plpgsql_build_datatype(Oid typeOid, int32 typmod, Oid collation,
 /*
  * Build an array type for the element type specified as argument.
  */
-PLpgSQL_type * plpgsql_build_datatype_arrayof(PLpgSQL_type *dtype) { PLpgSQL_type *typ; typ = (PLpgSQL_type *) palloc0(sizeof(PLpgSQL_type)); typ->typname = pstrdup("UNKNOWN"); typ->ttype = PLPGSQL_TTYPE_SCALAR; return typ;  }
+
+PLpgSQL_type * plpgsql_build_datatype_arrayof(PLpgSQL_type *dtype)
+{
+	if (dtype->typisarray)
+		return dtype;
+
+	PLpgSQL_type *array_type;
+	array_type = (PLpgSQL_type *) palloc0(sizeof(PLpgSQL_type));
+
+	array_type->ttype = PLPGSQL_TTYPE_REC;
+	array_type->atttypmod = dtype->atttypmod;
+	array_type->collation = dtype->collation;
+
+	array_type->typisarray = true;
+
+	switch(dtype->typoid)
+	{
+		case BOOLOID:
+			array_type->typoid = BOOLARRAYOID;
+			array_type->typname = pstrdup("boolean[]");
+			break;
+		case INT4OID:
+			array_type->typoid = INT4ARRAYOID;
+			array_type->typname = pstrdup("integer[]");
+			break;
+		case TEXTOID:
+			array_type->typoid = TEXTARRAYOID;
+			array_type->typname = pstrdup("text[]");
+			break;
+		default:
+			array_type->typname = pstrdup("UNKNOWN");
+			break;
+	}
+	array_type->typoid = dtype->typoid;
+
+	return array_type;
+}
+
 
 
 /*

diff --git a/test/sql/plpgsql_regress/plpgsql_array.sql b/test/sql/plpgsql_regress/plpgsql_array.sql
@@ -77,3 +77,74 @@ begin a[1] := 2; raise notice 'a = %', a; end$$;
 
 do $$ declare a complex;
 begin a.r[1] := 2; raise notice 'a = %', a; end$$;
+
+--
+-- 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;
+$$;
+
+-- some types don't support arrays
+do $$
+declare
+  v pg_node_tree;
+  v1 v%type[];
+begin
+end;
+$$;
+
+-- 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;
+$$;
+
+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;
+$$;