diff --git a/src/lib/data/locations.ts b/src/lib/data/locations.ts index 8cf2177e..8273d16b 100644 --- a/src/lib/data/locations.ts +++ b/src/lib/data/locations.ts @@ -1,5 +1,9 @@ import type { Database } from '../../db'; +import type { LocationId } from '../../db/models/location'; import type { PlanId } from '../../db/models/plan'; +import type { InstanceOfModel } from '../../db/util/types'; +import { getOrCreate } from '../../util'; +import { NotFoundError } from '../../util/error'; export const getPlanCountries = async ( database: Database, @@ -24,3 +28,55 @@ export const getPlanCountries = async ( }, }); }; + +export const getAllCountryLocations = async ( + database: Database, + countryId: LocationId +): Promise[]>> => { + const country = await database.location.findOne({ + where: { + id: countryId, + adminLevel: 0, + status: 'active', + parentId: { [database.Op.IS_NULL]: true }, + }, + }); + + if (!country) { + throw new NotFoundError( + `No country with ID ${countryId}. ` + + `Make sure location with ID ${countryId} is an active country (i.e. admin 0 location)` + ); + } + + let previousAdminLevelIds = [country.id]; + let nextAdminLevel = 1; + const countryLocationsByAdminLevel = new Map< + number, + Array> + >([[0, [country]]]); + + while (previousAdminLevelIds.length) { + const locations = await database.location.find({ + where: { + parentId: { [database.Op.IN]: previousAdminLevelIds }, + adminLevel: nextAdminLevel, + status: 'active', + }, + }); + + for (const location of locations) { + const countriesInAdminLevel = getOrCreate( + countryLocationsByAdminLevel, + nextAdminLevel, + () => [] + ); + countriesInAdminLevel.push(location); + } + + previousAdminLevelIds = locations.map((l) => l.id); + nextAdminLevel++; + } + + return countryLocationsByAdminLevel; +};