diff --git a/notifications_utils/broadcast_areas/__init__.py b/notifications_utils/broadcast_areas/__init__.py index 5cf3854b3..9c81a6c18 100644 --- a/notifications_utils/broadcast_areas/__init__.py +++ b/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) diff --git a/notifications_utils/broadcast_areas/create-broadcast-areas-db.py b/notifications_utils/broadcast_areas/create-broadcast-areas-db.py index bf20df9d9..f9fd1e575 100755 --- a/notifications_utils/broadcast_areas/create-broadcast-areas-db.py +++ b/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) diff --git a/notifications_utils/broadcast_areas/repo.py b/notifications_utils/broadcast_areas/repo.py new file mode 100644 index 000000000..61eb24ea4 --- /dev/null +++ b/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