broadcast-areas: refactor to use repo.py

share functionality between repo used by create script and by package

Signed-off-by: Toby Lorne <toby.lornewelch-richards@digital.cabinet-office.gov.uk>

notifications_utils/broadcast_areas/__init__.py

@@ -9,59 +9,7 @@ from notifications_utils.formatters import formatted_list
 from notifications_utils.serialised_model import SerialisedModelCollection
 from notifications_utils.safe_string import make_string_safe_for_id
 
-
-class BroadcastAreasRepository(object):
-    def __init__(self):
-        self.database = Path(__file__).resolve().parent / 'broadcast-areas.sql'
-
-    def query(self, sql, *args):
-        package_path = Path(__file__).resolve().parent
-        db_filepath = package_path / "broadcast-areas.sqlite3"
-        with sqlite3.connect(str(db_filepath)) as conn:
-            cursor = conn.cursor()
-            cursor.execute(sql, (*args,))
-            return cursor.fetchall()
-
-    def get_datasets(self):
-        q = "SELECT DISTINCT(dataset) AS dataset FROM broadcast_areas"
-        datasets = [row[0] for row in self.query(q)]
-        return sorted(datasets)
-
-    def get_dataset_description(self, dataset):
-        q = """
-        WITH
-        areas AS (SELECT * FROM broadcast_areas WHERE dataset = ?),
-        area_count AS (SELECT COUNT(*) AS c FROM areas),
-        subset_area_count AS (SELECT c - 4 FROM area_count),
-        some_area_names  AS (SELECT name FROM areas LIMIT 100),
-        some_shuffled_area_names AS (
-            SELECT name FROM some_area_names ORDER BY RANDOM()
-        ),
-        description_area_names AS (
-            SELECT name FROM some_shuffled_area_names LIMIT 4
-        ),
-        description_areas_joined AS (
-            SELECT GROUP_CONCAT(name, ", ") FROM description_area_names
-        )
-        SELECT
-        CASE (SELECT * FROM subset_area_count)
-        WHEN 0 THEN
-            (SELECT * FROM description_areas_joined)
-        ELSE
-            (SELECT * FROM description_areas_joined)
-            || ", "
-            || (SELECT * FROM subset_area_count)
-            || " more…"
-        END
-        """
-        description = self.query(q, dataset)[0][0]
-        return description
-
-    def get_areas(self):
-        pass
-
-    def get_areas_from_list(self, areas):
-        pass
+from .repo import BroadcastAreasRepository
 
 
 @lru_cache(maxsize=128)

notifications_utils/broadcast_areas/create-broadcast-areas-db.py

@@ -3,23 +3,14 @@
 import geojson
 from pathlib import Path
 
+from repo import BroadcastAreasRepository
+
 package_path = Path(__file__).resolve().parent
 
-db_filepath = package_path / "broadcast-areas.sqlite3"
-os.remove(str(db_filepath))
-conn = sqlite3.connect(str(db_filepath))
-conn.execute("""
-CREATE TABLE broadcast_areas (
-    id TEXT PRIMARY KEY,
-    name TEXT NOT NULL,
-    dataset TEXT NOT NULL,
-    feature_geojson TEXT NOT NULL
-)""")
-conn.execute("""
-CREATE INDEX broadcast_areas_dataset
-ON broadcast_areas (dataset);
-""")
-conn.commit()
+repo = BroadcastAreasRepository()
+
+repo.delete_db()
+repo.create_tables()
 
 simple_datasets = [
     ("Countries", "ctry19nm"),
@@ -30,17 +21,11 @@ for dataset_name, name_field in simple_datasets:
     filepath = package_path / "{}.geojson".format(dataset_name)
     dataset_geojson = geojson.loads(filepath.read_text())
 
-    q = """
-    INSERT INTO broadcast_areas (id, name, dataset, feature_geojson)
-    VALUES (?, ?, ?, ?)
-    """
+    repo.insert_broadcast_area_library(dataset_name)
 
     for feature in dataset_geojson["features"]:
         f_name = feature["properties"][name_field]
-        f_id = make_string_safe_for_id(f_name)
-        conn.execute(q, (f_id, f_name, dataset_name, geojson.dumps(feature)))
-
-    conn.commit()
+        repo.insert_broadcast_areas([[f_name, dataset_name, feature]])
 
 # https://geoportal.statistics.gov.uk/datasets/wards-may-2020-boundaries-uk-bgc
 # Converted to geojson manually from SHP because of GeoJSON download limits
@@ -54,8 +39,12 @@ ward_code_to_la_mapping = {
     for f in geojson.loads(las_filepath.read_text())["features"]
 }
 
+dataset_name = "Electoral Wards of the United Kingdom"
+repo.insert_broadcast_area_library(dataset_name)
+
+areas_to_add = []
+
 for f in geojson.loads(wards_filepath.read_text())["features"]:
-    dataset_name = "Electoral Wards of the United Kingdom"
 
     ward_code = f["properties"]["wd20cd"]
     ward_name = f["properties"]["wd20nm"]
@@ -64,15 +53,9 @@ for f in geojson.loads(wards_filepath.read_text())["features"]:
         la_name = ward_code_to_la_mapping[ward_code]
 
         f_name = "{} - {}".format(la_name, ward_name)
-        f_id = make_string_safe_for_id(f_name)
-
-        q = """
-        INSERT INTO broadcast_areas (id, name, dataset, feature_geojson)
-        VALUES (?, ?, ?, ?)
-        """
-        conn.execute(q, (f_id, f_name, dataset_name, geojson.dumps(f)))
+        areas_to_add.append([f_name, dataset_name, f])
 
     except KeyError:
         print("Skipping", ward_code, ward_name)  # noqa: T001
 
-conn.commit()
+repo.insert_broadcast_areas(areas_to_add)

notifications_utils/broadcast_areas/repo.py

@@ -0,0 +1,145 @@
+import geojson
+import os
+from pathlib import Path
+import sqlite3
+
+from notifications_utils.safe_string import make_string_safe_for_id
+
+
+class BroadcastAreasRepository(object):
+    def __init__(self):
+        self.database = Path(__file__).resolve().parent / 'broadcast-areas.sqlite3'
+
+    def conn(self):
+        return sqlite3.connect(str(self.database))
+
+    def delete_db(self):
+        os.remove(str(self.database))
+
+    def create_tables(self):
+        with self.conn() as conn:
+            conn.execute("""
+            CREATE TABLE broadcast_area_libraries (
+                id TEXT PRIMARY KEY,
+                name TEXT NOT NULL
+            )""")
+
+            conn.execute("""
+            CREATE TABLE broadcast_areas (
+                id TEXT PRIMARY KEY,
+                name TEXT NOT NULL,
+                broadcast_area_library_id TEXT NOT NULL,
+                feature_geojson TEXT NOT NULL,
+
+                FOREIGN KEY (broadcast_area_library_id)
+                    REFERENCES broadcast_area_libraries(id)
+            )""")
+
+            conn.execute("""
+            CREATE INDEX broadcast_areas_broadcast_area_library_id
+            ON broadcast_areas (broadcast_area_library_id);
+            """)
+
+    def insert_broadcast_area_library(self, broadcast_area_name):
+        broadcast_area_id = make_string_safe_for_id(broadcast_area_name)
+
+        q = """
+        INSERT INTO broadcast_area_libraries (id, name)
+        VALUES (?, ?)
+        """
+
+        with self.conn() as conn:
+            conn.execute(q, (broadcast_area_id, broadcast_area_name))
+
+    def insert_broadcast_areas(self, areas):
+
+        q = """
+        INSERT INTO broadcast_areas (
+            id, name,
+            broadcast_area_library_id, feature_geojson
+        )
+        VALUES (?, ?, ?, ?)
+        """
+
+        with self.conn() as conn:
+            for name, area_name, feature in areas:
+                id = make_string_safe_for_id(name)
+                area_id = make_string_safe_for_id(area_name)
+
+                conn.execute(q, (id, name, area_id, geojson.dumps(feature)))
+
+    def query(self, sql, *args):
+        with self.conn() as conn:
+            cursor = conn.cursor()
+            cursor.execute(sql, (*args,))
+            return cursor.fetchall()
+
+    def get_libraries(self):
+        q = "SELECT id, name FROM broadcast_area_libraries"
+        results = self.query(q)
+        libraries = [row[1] for row in results]
+        return sorted(libraries)
+
+    def get_library_description(self, library_id):
+        q = """
+        WITH
+        areas AS (SELECT * FROM broadcast_areas
+                  WHERE broadcast_area_library_id = ?),
+        area_count AS (SELECT COUNT(*) AS c FROM areas),
+        subset_area_count AS (SELECT c - 4 FROM area_count),
+        some_area_names  AS (SELECT name FROM areas LIMIT 100),
+        some_shuffled_area_names AS (
+            SELECT name FROM some_area_names ORDER BY RANDOM()
+        ),
+        description_area_names AS (
+            SELECT name FROM some_shuffled_area_names LIMIT 4
+        ),
+        description_areas_joined AS (
+            SELECT GROUP_CONCAT(name, ", ") FROM description_area_names
+        )
+        SELECT
+        CASE (SELECT * FROM subset_area_count)
+        WHEN 0 THEN
+            (SELECT * FROM description_areas_joined)
+        ELSE
+            (SELECT * FROM description_areas_joined)
+            || ", "
+            || (SELECT * FROM subset_area_count)
+            || " more…"
+        END
+        """
+        description = self.query(q, library_id)[0][0]
+        return description
+
+    def get_areas(self, *area_ids):
+        with self.conn() as conn:
+            cursor = conn.cursor()
+
+            q = """
+            SELECT id, name, feature_geojson FROM broadcast_areas
+            WHERE id IN ({})
+            """.format(("?," * len(*area_ids))[:-1])
+            cursor.execute(q, *area_ids)
+            results = cursor.fetchall()
+
+            areas = [
+                (row[0], row[1], geojson.loads(row[2]))
+                for row in results
+            ]
+
+            return areas
+
+    def get_all_areas_for_library(self, library_id):
+        q = """
+        SELECT id, name, feature_geojson FROM broadcast_areas
+        WHERE broadcast_area_library_id = ?
+        """
+
+        results = self.query(q, library_id)
+
+        areas = [
+            (row[0], row[1], geojson.loads(row[2]))
+            for row in results
+        ]
+
+        return areas