user input categories, refactor, remove zombie code

1 files changed, 45 insertions, 71 deletions

diff --git a/src/db_handler.py b/src/db_handler.py
index 5a2ab73..b438b6c 100644
--- a/src/db_handler.py
+++ b/src/db_handler.py
@@ -3,9 +3,27 @@ import os
 import psycopg
 import yaml
 
+
+class IllegalCategoryException(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 subclass
+        FROM poi
+        WHERE subclass NOT LIKE '%;%'
+        GROUP BY class, subclass
+        HAVING COUNT(*) > 1
+        ;
+        """)
+        categories = cur.fetchall()
+        categories = [category[0] for category in categories]
+
 
 def get_chains() -> (list[dict]):
     with psycopg.connect(**db_con_params, row_factory=psycopg.rows.dict_row) as conn:
@@ -16,79 +34,35 @@ def get_chains() -> (list[dict]):
             chains = cur.fetchall()
             return chains
 
-def get_all() -> (list[dict]):
-    with psycopg.connect(**db_con_params, row_factory=psycopg.rows.dict_row) as conn:
-        with conn.cursor() as cur:
-            cur.execute("""
- SELECT
- shop.osm_id,
- shop.name,
- shop.brand,
- ST_Y(ST_Transform(shop.geom, 4326)) AS lat,
- ST_X(ST_Transform(shop.geom, 4326)) AS long,
+
+def get_all(category: str) -> (list[dict]):
+    if category not in categories:
+        raise IllegalCategoryException()
+    else:
+        with psycopg.connect(**db_con_params, row_factory=psycopg.rows.dict_row) as conn:
+            with conn.cursor() as cur:
+                cur.execute("""
+WITH filtered AS (
+SELECT osm_id, name, brand, geom, class, subclass
+FROM poi
+WHERE subclass = %(subclass)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 shop
+ FROM filtered
  JOIN
  (
    SELECT (ST_DUMP(ST_VoronoiPolygons(ST_Collect(geom)))).geom as geom
-   FROM shop
-   WHERE subclass = 'supermarket'
- ) polygon ON ST_Contains(polygon.geom, shop.geom)
- WHERE subclass = 'supermarket'
- ;
-            """)
-
-            all = cur.fetchall()
-            return all
-
-
-def get_supermarkets() -> (list[dict], list[str]):
-    with psycopg.connect(**db_con_params, row_factory=psycopg.rows.dict_row) as conn:
-        with conn.cursor() as cur:
-            cur.execute("""
-            SELECT name, 
-            CASE WHEN brand ILIKE '%365%' THEN '365discount' ELSE BRAND END,
-            ST_Y(ST_Transform(geom, 4326)) AS lat, ST_X(ST_Transform(geom, 4326)) AS long 
-            FROM shop WHERE subclass = 'supermarket';
-            """)
-
-            supermarkets = cur.fetchall()
-            return supermarkets
-
-
-def get_stores() -> (list[dict], list[str]):
-    with psycopg.connect(**db_con_params) as conn:
-        with conn.cursor() as cur:
-            cur.execute("""
-            SELECT
-            ST_Y(ST_Transform(geom, 4326)) AS lat, ST_X(ST_Transform(geom, 4326)) AS long 
-            FROM shop;
-            """)
-
-            stores = cur.fetchall()
-            return stores
-
-
-def voronoi_polygons() -> (list[dict], list[str]):
-    with psycopg.connect(**db_con_params, row_factory=psycopg.rows.dict_row) as conn:
-        with conn.cursor() as cur:
-            """
-            SELECT ST_AsGeoJSON(ST_Transform(ST_VoronoiPolygons(ST_Collect(geom)), 4326)) as geojson
-            FROM shop 
-            WHERE subclass = 'supermarket';
-            """
-            """
-            SELECT *,
-            ST_ClusterDBSCAN(<geom>, 0, 1) OVER() AS _clst
-            FROM   <your_table>;
-            """
-            cur.execute("""
-                SELECT ARRAY_AGG(name) as names, ST_AsGeoJSON(ST_Transform(ST_VoronoiPolygons(ST_Collect(geom)), 4326)) as geojson
-                FROM shop
-                WHERE subclass = 'supermarket' AND osm_id IN (7688736659, 7682568726, 1085855265, 4517989992)
-                ;
-            """)
-
-            voronoi_polygons = cur.fetchone()
+   FROM filtered
+ ) polygon ON ST_Contains(polygon.geom, filtered.geom)
+;
+            """, {'subclass': category})
 
-            return voronoi_polygons
+                all = cur.fetchall()
+                return all