From f4a36d5347ec7d8a74ebbabcde15d617a44ce104 Mon Sep 17 00:00:00 2001 From: msepga Date: Tue, 15 Oct 2024 19:01:39 -0400 Subject: [PATCH] Mock `plpgsql_build_datatype_arrayof` with special type handling --- scripts/extract_source.rb | 39 +++++++++++- src/postgres/src_pl_plpgsql_src_pl_comp.c | 39 +++++++++++- test/sql/plpgsql_regress/plpgsql_array.sql | 71 ++++++++++++++++++++++ 3 files changed, 147 insertions(+), 2 deletions(-) diff --git a/scripts/extract_source.rb b/scripts/extract_source.rb index 7c4f3e30..ed6b9925 100644 --- 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 index dc7746c9..149ab679 100644 --- 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 index 4c3f26be..4b9ff515 100644 --- 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; +$$;