1 files changed, 34 insertions, 9 deletions
diff --git a/src/db_handler.py b/src/db_handler.py
index fd68c39..335e9e4 100644
--- a/src/db_handler.py
+++ b/src/db_handler.py
@@ -15,6 +15,18 @@ 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 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
@@ -24,8 +36,8 @@ HAVING COUNT(*) > 1
 ORDER BY class, subclass
 ;
         """)
-        categories = cur.fetchall()
-        categories = [(f"{category[0]}:{category[1]}") for category in categories]
+        sub_cat = cur.fetchall()
+        categories += [(f"{category[0]}:{category[1]}") for category in sub_cat]
         print(f"Loaded: {len(categories)} categories")
 
         cur.execute("""
@@ -55,16 +67,25 @@ def get_all(country: str, category: str) -> (list[dict]):
         raise IllegalInstructionException("Category not found")
     if country not in countries:
         raise IllegalInstructionException("Country not found")
-    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("""
+
+    query = """
 WITH filtered AS (
 SELECT osm_id, name, brand, geom, class, subclass
 FROM poi
 WHERE class = %(class)s
-AND subclass = %(subclass)s
-AND country = %(country)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
@@ -81,7 +102,11 @@ SELECT
    FROM filtered
  ) polygon ON ST_Contains(polygon.geom, filtered.geom)
 ;
-            """, {'class': class_, 'subclass': sub_class, 'country': country})
+    """
+
+    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