1 files changed, 22 insertions, 4 deletions
diff --git a/src/db_handler.py b/src/db_handler.py
index b438b6c..83cdad8 100644
--- a/src/db_handler.py
+++ b/src/db_handler.py
@@ -1,10 +1,12 @@
 import os
+import logging
 
 import psycopg
 import yaml
 
+logger = logging.getLogger(__name__) # Set the logger name, to the name of the module
 
-class IllegalCategoryException(Exception):
+class IllegalInstructionException(Exception):
     pass
 
 
@@ -23,6 +25,18 @@ with psycopg.connect(**db_con_params) as conn:
         """)
         categories = cur.fetchall()
         categories = [category[0] for category in categories]
+        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]):
@@ -35,9 +49,11 @@ def get_chains() -> (list[dict]):
             return chains
 
 
-def get_all(category: str) -> (list[dict]):
+def get_all(country: str, category: str) -> (list[dict]):
     if category not in categories:
-        raise IllegalCategoryException()
+        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:
@@ -46,6 +62,7 @@ WITH filtered AS (
 SELECT osm_id, name, brand, geom, class, subclass
 FROM poi
 WHERE subclass = %(subclass)s
+AND country = %(country)s
 )
 
 SELECT
@@ -62,7 +79,8 @@ SELECT
    FROM filtered
  ) polygon ON ST_Contains(polygon.geom, filtered.geom)
 ;
-            """, {'subclass': category})
+            """, {'subclass': category, 'country': country})
 
                 all = cur.fetchall()
                 return all
+