Categories - Add category type prefix e.g. shop:supermarket

1 files changed, 17 insertions, 15 deletions

diff --git a/src/db_handler.py b/src/db_handler.py
index 83cdad8..fd68c39 100644
--- a/src/db_handler.py
+++ b/src/db_handler.py
@@ -16,15 +16,16 @@ with open(os.path.join('configs', 'database.yml'), 'r') as file:
 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
-        ;
+SELECT class, subclass
+FROM poi
+WHERE subclass NOT LIKE '%;%'
+GROUP BY class, subclass
+HAVING COUNT(*) > 1
+ORDER BY class, subclass
+;
         """)
         categories = cur.fetchall()
-        categories = [category[0] for category in categories]
+        categories = [(f"{category[0]}:{category[1]}") for category in categories]
         print(f"Loaded: {len(categories)} categories")
 
         cur.execute("""
@@ -54,14 +55,15 @@ def get_all(country: str, category: str) -> (list[dict]):
         raise IllegalInstructionException("Category not found")
     if country not in countries:
         raise IllegalInstructionException("Country not found")
-    else:
-        with psycopg.connect(**db_con_params, row_factory=psycopg.rows.dict_row) as conn:
-            with conn.cursor() as cur:
-                cur.execute("""
+    class_, sub_class = category.split(':')
+    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
+WHERE class = %(class)s
+AND subclass = %(subclass)s
 AND country = %(country)s
 )
 
@@ -79,8 +81,8 @@ SELECT
    FROM filtered
  ) polygon ON ST_Contains(polygon.geom, filtered.geom)
 ;
-            """, {'subclass': category, 'country': country})
+            """, {'class': class_, 'subclass': sub_class, 'country': country})
 
-                all = cur.fetchall()
-                return all
+            all = cur.fetchall()
+            return all