import os import logging import psycopg import yaml logger = logging.getLogger(__name__) # Set the logger name, to the name of the module class IllegalInstructionException(Exception): pass with open(os.path.join('configs', 'database.yml'), 'r') as file: db_con_params = yaml.safe_load(file.read()) with psycopg.connect(**db_con_params) as conn: with conn.cursor() as cur: cur.execute(""" SELECT class FROM poi GROUP BY class HAVING COUNT(*) > 1 ORDER BY class ; """) classes = cur.fetchall() categories = [category[0] for category in classes] cur.execute(""" SELECT class, subclass FROM poi WHERE subclass NOT LIKE '%;%' GROUP BY class, subclass HAVING COUNT(*) > 1 ORDER BY class, subclass ; """) sub_cat = cur.fetchall() categories += [(f"{category[0]}:{category[1]}") for category in sub_cat] print(f"Loaded: {len(categories)} categories") cur.execute(""" SELECT country FROM poi WHERE country is not null GROUP BY country HAVING COUNT(*) > 1 ;""") countries = cur.fetchall() countries = [country[0] for country in countries] print(f"Loaded countries: {countries}") def get_chains() -> (list[dict]): with psycopg.connect(**db_con_params, row_factory=psycopg.rows.dict_row) as conn: with conn.cursor() as cur: cur.execute(""" SELECT name, color FROM chain """) chains = cur.fetchall() return chains def get_all(country: str, category: str) -> (list[dict]): if category not in categories: raise IllegalInstructionException("Category not found") if country not in countries: raise IllegalInstructionException("Country not found") query = """ WITH filtered AS ( SELECT osm_id, name, brand, geom, class, subclass FROM poi WHERE class = %(class)s """ params = {'country': country} if ':' in category: query += " AND subclass = %(subclass)s " class_, sub_class = category.split(':') params['class'] = class_ params['subclass'] = sub_class else: params['class'] = category query += """AND country = %(country)s ) SELECT filtered.osm_id, filtered.name, filtered.brand, ST_Y(ST_Transform(filtered.geom, 4326)) AS lat, ST_X(ST_Transform(filtered.geom, 4326)) AS long, ST_AsGeoJSON(ST_Transform(polygon.geom, 4326)) as polygon FROM filtered JOIN ( SELECT (ST_DUMP(ST_VoronoiPolygons(ST_Collect(geom)))).geom as geom FROM filtered ) polygon ON ST_Contains(polygon.geom, filtered.geom) ; """ with psycopg.connect(**db_con_params, row_factory=psycopg.rows.dict_row) as conn: with conn.cursor() as cur: cur.execute(query, params) all = cur.fetchall() return all