nixify running

needs databse, this is handled separately

1 files changed, 96 insertions, 96 deletions

diff --git a/src/db_handler.py b/src/db_handler.py
index 335e9e4..afc2852 100644
--- a/src/db_handler.py
+++ b/src/db_handler.py
@@ -9,105 +9,105 @@ logger = logging.getLogger(__name__) # Set the logger name, to the name of the m
 class IllegalInstructionException(Exception):
     pass
 
+class DBHandler:
+    def __init__(self, conninfo: str):
+        self.conninfo = conninfo
 
-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
+        with psycopg.connect(conninfo=self.conninfo) 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()
+                self.categories = [category[0] for category in classes]
+
+                cur.execute("""
+        SELECT class, subclass
         FROM poi
-        WHERE country is not null GROUP BY country
+        WHERE subclass NOT LIKE '%;%'
+        GROUP BY class, subclass
         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)
-;
+        ORDER BY class, subclass
+        ;
+                """)
+                sub_cat = cur.fetchall()
+                self.categories += [(f"{category[0]}:{category[1]}") for category in sub_cat]
+                print(f"Loaded: {len(self.categories)} categories")
+
+                cur.execute("""
+                SELECT country
+                FROM poi
+                WHERE country is not null GROUP BY country
+                HAVING COUNT(*) > 1
+                ;""")
+                self.countries = cur.fetchall()
+                self.countries = [country[0] for country in self.countries]
+
+                print(f"Loaded countries: {self.countries}")
+
+
+    def get_chains(self) -> (list[dict]):
+        with psycopg.connect(conninfo=self.conninfo, 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(self, country: str, category: str) -> (list[dict]):
+        if category not in self.categories:
+            raise IllegalInstructionException("Category not found")
+        if country not in self.countries:
+            raise IllegalInstructionException("Country not found")
+
+        query = """
+    WITH filtered AS (
+    SELECT osm_id, name, brand, geom, class, subclass
+    FROM poi
+    WHERE class = %(class)s
     """
 
-    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
+        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(conninfo=self.conninfo, row_factory=psycopg.rows.dict_row) as conn:
+            with conn.cursor() as cur:
+                cur.execute(query, params)
+
+                all = cur.fetchall()
+                return all