plpgsql one input get all the related fields/columns to json failed - postgresql

Question

Try to answer this question then found out, I cannot solve it.
Basic idea: propagate twice, from country_id propagate to state_id, state_id propagate to city_id twice.Then the country_id need be joined twice. when we do array_agg on state level, we need explicitly join country_id, during city level we also need using join country_id.
Reference link: https://github.com/hettie-d/NORM/tree/master/sql
Basic idea, input one country_id, all the relevant country, state, city level information will be transformed to json format.
Prepare. I use country_id, state_id, city_id, since they are more descriptive.
begin;
create table public.country(country_id bigint primary key , name text, leader text);
create table public.states(state_id bigint primary key, name text, population bigint,country_id bigint REFERENCES public.country (country_id));
create table public.cities(city_id bigint,name text,state_id bigint REFERENCES public.states (state_id));
insert into public.country values ( 1, 'India', 'Narendra Modi');
insert into public.country values ( 2 , 'USA', 'Joe Biden');
insert into public.country values ( 3 , 'Australia', 'Scott Morrison');
insert into public.states values( 1 ,'California' , 39500000 , 2);
insert into public.states values( 2 , 'Washington' , 7610000 ,2 );
insert into public.states values( 4 , 'Karnataka' , 64100000,1);
insert into public.states values( 5 , 'Rajasthan' , 68900000,1 );
insert into public.states values( 6 , 'Maharashtra' , 125700000,1 );
insert into public.cities values( 1 , 'Mumbai' , 6 );
insert into public.cities values( 2 , 'Pune' , 6 );
insert into public.cities values( 3 , 'San Francisco' , 1 );
commit;
--- create composite types.
begin;
create type city_record as(city_name text);
create type state_record as (state_name text, population bigint,cities city_record[]);
create type country_record as (country_name text, leader text, states state_record[]);
commit;
array transport
create or replace
function array_transport (all_items anyarray) returns setof text
returns null on null input
language plpgsql as
$body$
declare
item record;
begin
foreach item in array all_items
loop
return next(to_json(item)::text);
end loop;
end;
$body$;
--the main function country_select_json
create or replace function country_select_json (_country_id bigint)
returns country_record[]
as
$$
declare
_result text;
begin
select array_agg(single_item)
from (select
array_agg(row(
co.name,
co.leader,
(select array_agg(row
(s.name,
s.population,
(select array_agg
(row
(c.name)::city_record)
from cities c
join states s using (state_id)
where s.country_id = co.country_id)
)::state_record) from states s where s.country_id = co.country_id
)
)::country_record)
as single_item
from country co
where co.country_id = _country_id)y into _result;
-- raise info 'state_record test: %', _result;
return (_result);
end
$$ language plpgsql;
run
select * from array_transport(country_select_json(1));
{"country_name":"India","leader":"Narendra Modi","states":[{"state_name":"Karnataka","population":64100000,"cities":[{"city_name":"Mumbai"},{"city_name":"Pune"}]},{"state_name":"Rajasthan","population":68900000,"cities":[{"city_name":"Mumbai"},{"city_name":"Pune"}]},{"state_name":"Maharashtra","population":125700000,"cities":[{"city_name":"Mumbai"},{"city_name":"Pune"}]}]}
(1 row)
country level, state level ok, but the city level is wrong. How to solve this problem.
Expected Result:
{"country_name":"India","leader":"Narendra Modi","states":[{"state_name":"Karnataka","population":64100000,"cities":[NULL]},{"state_name":"Rajasthan","population":68900000,"cities":[NULL]},{"state_name":"Maharashtra","population":125700000,"cities":[{"city_name":"Mumbai"},{"city_name":"Pune"}]}]}
update 2022-03-04.
(select array_agg
(c.name) as city_name
from cities c
join states s using (state_id)
where s.country_id = co.country_id)
Now I know the problem: because the propagate is first from city, then to state then country. Once function input the country_id then all the country related city name will be pulled together.

Answer 1

Does this query answer your needs ? (Result here)
with ci as (select cities.state_id,jsonb_agg(jsonb_build_object('city_name',cities.name)) as cities from cities group by state_id)
select jsonb_pretty(jsonb_build_object(
'country_name',c.name,
'leader',c.leader,
'states', jsonb_agg(jsonb_build_object(
'state_name',s.name,
'population',s.population,
'cities',ci.cities
))
))
from country c left join states s on s.country_id = c.country_id
left join ci on ci.state_id = s.state_id
where c.country_id = 1
group by c.name,c.leader
// Result
{
"country_name": "India",
"leader": "Narendra Modi",
"states": [
{
"state_name": "Maharashtra",
"population": 125700000,
"cities": [
{
"city_name": "Mumbai"
},
{
"city_name": "Pune"
}
]
},
{
"state_name": "Rajasthan",
"population": 68900000,
"cities": null
},
{
"state_name": "Karnataka",
"population": 64100000,
"cities": null
}
]