How to query 3 tables with one to many relationship along with json aggregation?

Question

I have 3 tables like shown below with some sample data:

countries

id	name	leader
1	India	Narendra Modi
2	USA	Joe Biden
3	Australia	Scott Morrison

states

id	name	population	country_id
1	California	39500000	2
2	Washington	7610000	2
3	Florida	21500000	2
4	Karnataka	64100000	1
5	Rajasthan	68900000	1
6	Maharastra	125700000	1

cities

id	name	state_id
1	Mumbai	6
2	Pune	6
3	San Francisco	1

I would like to get all the countries with their states and cities queried by population. The result should look like this:

[
  {
    "country_name": "USA",
    "leader": "Joe Biden",
    "states": [
      {
        "state_name": "California",
        "population": 39500000,
        "cities": [
          {
            "city_name": "San Fransico"
          }
        ]
      }
    ]
  },
  {
    "country_name": "India",
    "leader": "Narendra Modi",
    "states": [
      {
        "state_name": "Maharastra",
        "population": 125700000,
        "cities": [
          {
            "city_name": "Mumbai"
          },
          {
            "city_name": "Pune"
          }
        ]
      }
    ]
  }
]

I was able to achieve this somewhat with sub queries like this:

select *,
  (
          select
            array_to_json(array_agg(row_to_json(states.*))) as array_to_json
          from (
                select states.*, (
                  select
            array_to_json(array_agg(row_to_json(cities.*))) as array_to_json
            from(
              select cities.* from cities where state_id = states.id
            ) cities) as cities
                 from states
                         where country_id = countries.id and population=125700000 OR population = 39500000
               ) states ) as states
from countries;

But I'd like to use a neat query where I can achieve this using a single join query with JSON aggregators in Postgresql.

Answer 1

try this

WITH list AS
(
SELECT s.country_id
     , jsonb_build_object( 'state_name', s.name                         , 
                         , 'population', s.population
                         , 'cities', jsonb_agg( jsonb_build_object('city_name', c.name))
                         ) AS state_list
  FROM states AS s
 INNER JOIN cities AS c
    ON c.state_id = s.id
 GROUP BY s.country_id, s.name, s.population
)
SELECT jsonb_agg( jsonb_build_object( 'country_name', c.name
                                    , 'leader', c.leader
                                    , 'states', l.state_list
                                    )
                )
  FROM list AS l
 INNER JOIN countries AS c
    ON c.id = l.country_id