From 69fb2d2a2c3ac4b8b9dd2494a9968518e676454b Mon Sep 17 00:00:00 2001 From: mariogiampieri Date: Tue, 19 Nov 2024 00:06:46 -0500 Subject: [PATCH 1/2] Chart enhancements + stub of `places` and `districtingProblems` (#167) Co-authored-by: bailliekova Co-authored-by: Raphael Paul Laude Co-authored-by: nofurtherinformation --- .../app/components/sidebar/ColorPicker.tsx | 14 ++- .../sidebar/charts/HorizontalBarChart.tsx | 94 +++++++++++++--- app/src/app/store/mapStore.ts | 76 +++++++------ app/src/app/utils/api/apiHandlers.ts | 51 +++++---- app/src/app/utils/api/queries.ts | 41 ++++--- .../65a4fc0a727d_add_unshatter_udf.py | 1 - ...d35_add_available_stats_to_districtrmap.py | 6 +- backend/app/main.py | 19 +++- backend/tests/test_main.py | 30 +++-- pipelines/simple_elt/main.py | 103 +++++++++++++++++- 10 files changed, 325 insertions(+), 110 deletions(-) diff --git a/app/src/app/components/sidebar/ColorPicker.tsx b/app/src/app/components/sidebar/ColorPicker.tsx index 39f2b72f7..0d048e899 100644 --- a/app/src/app/components/sidebar/ColorPicker.tsx +++ b/app/src/app/components/sidebar/ColorPicker.tsx @@ -3,6 +3,7 @@ import {Button, Checkbox, CheckboxGroup} from '@radix-ui/themes'; import {styled} from '@stitches/react'; import * as RadioGroup from '@radix-ui/react-radio-group'; import {blackA} from '@radix-ui/colors'; +import {useMapStore} from '@/app/store/mapStore'; type ColorPickerProps = T extends true ? { @@ -27,6 +28,8 @@ export const ColorPicker = ({ colorArray, multiple, }: ColorPickerProps) => { + const mapDocument = useMapStore(state => state.mapDocument); + if (multiple) { return (
@@ -63,11 +66,12 @@ export const ColorPicker = ({ value={value !== undefined ? colorArray[value] : undefined} defaultValue={colorArray[defaultValue]} > - {colorArray.map((color, i) => ( - - - - ))} + {mapDocument && + colorArray.slice(0, mapDocument.num_districts ?? 0).map((color, i) => ( + + + + ))}
); diff --git a/app/src/app/components/sidebar/charts/HorizontalBarChart.tsx b/app/src/app/components/sidebar/charts/HorizontalBarChart.tsx index 74ed8b4da..8fa91d975 100644 --- a/app/src/app/components/sidebar/charts/HorizontalBarChart.tsx +++ b/app/src/app/components/sidebar/charts/HorizontalBarChart.tsx @@ -1,7 +1,18 @@ import {useMapStore} from '@/app/store/mapStore'; import {Card, Flex, Heading, Text} from '@radix-ui/themes'; -import {BarChart, Bar, ResponsiveContainer, Tooltip, XAxis, YAxis, Cell} from 'recharts'; +import { + BarChart, + Bar, + ResponsiveContainer, + Tooltip, + XAxis, + YAxis, + Cell, + ReferenceLine, + Label, +} from 'recharts'; import {colorScheme} from '@/app/constants/colors'; +import {useState, useMemo} from 'react'; type TooltipInput = { active?: boolean; @@ -24,6 +35,38 @@ const CustomTooltip = ({active, payload: items}: TooltipInput) => { export const HorizontalBar = () => { const mapMetrics = useMapStore(state => state.mapMetrics); + const summaryStats = useMapStore(state => state.summaryStats); + const numDistricts = useMapStore(state => state.mapDocument?.num_districts); + const idealPopulation = summaryStats?.idealpop?.data; + const maxNumberOrderedBars = 40; // max number of zones to consider while keeping blank spaces for missing zones + const [totalExpectedBars, setTotalExpectedBars] = useState< + Array<{zone: number; total_pop: number}> + >([]); + + const calculateChartObject = () => { + if ((numDistricts ?? 0) < maxNumberOrderedBars) { + return mapMetrics && mapMetrics.data && numDistricts + ? Array.from({length: numDistricts}, (_, i) => i + 1).reduce( + (acc, district) => { + const totalPop = mapMetrics.data.reduce((acc, entry) => { + return entry.zone === district ? acc + entry.total_pop : acc; + }, 0); + return [...acc, {zone: district, total_pop: totalPop}]; + }, + [] as Array<{zone: number; total_pop: number}> + ) + : []; + } else { + return mapMetrics?.data ?? []; + } + }; + + useMemo(() => { + if (mapMetrics) { + const chartObject = calculateChartObject(); + setTotalExpectedBars(chartObject); + } + }, [mapMetrics]); if (mapMetrics?.isPending) { return
Loading...
; @@ -44,30 +87,49 @@ export const HorizontalBar = () => { return ( - Population by Zone + Population by district - - + + + idealPopulation + ? Math.round(Math.max(idealPopulation * 2, dataMax + 1000)) + : dataMax, + ]} tickFormatter={value => numberFormat.format(value)} /> - + } /> - {mapMetrics.data - .sort((a, b) => a.zone - b.zone) - .map((entry, index) => ( - - ))} + {totalExpectedBars && + totalExpectedBars + .sort((a, b) => a.zone - b.zone) + .map((entry, index) => ( + + ))} + + diff --git a/app/src/app/store/mapStore.ts b/app/src/app/store/mapStore.ts index dccfb99d7..d1864251b 100644 --- a/app/src/app/store/mapStore.ts +++ b/app/src/app/store/mapStore.ts @@ -23,19 +23,19 @@ import { getFeaturesInBbox, resetZoneColors, setZones, -} from "../utils/helpers"; -import { getRenderSubscriptions } from "./mapRenderSubs"; -import { getSearchParamsObersver } from "../utils/api/queryParamsListener"; -import { getMapMetricsSubs } from "./metricsSubs"; -import { getMapEditSubs } from "./mapEditSubs"; -import { getQueriesResultsSubs } from "../utils/api/queries"; -import { persistOptions } from "./persistConfig"; +} from '../utils/helpers'; +import {getRenderSubscriptions} from './mapRenderSubs'; +import {getSearchParamsObersver} from '../utils/api/queryParamsListener'; +import {getMapMetricsSubs} from './metricsSubs'; +import {getMapEditSubs} from './mapEditSubs'; +import {getQueriesResultsSubs} from '../utils/api/queries'; +import {persistOptions} from './persistConfig'; import {patchReset, patchShatter, patchUnShatter} from '../utils/api/mutations'; import bbox from '@turf/bbox'; import {BLOCK_SOURCE_ID} from '../constants/layers'; import {DistrictrMapOptions} from './types'; -import { onlyUnique } from '../utils/arrays'; -import { queryClient } from '../utils/api/queryClient'; +import {onlyUnique} from '../utils/arrays'; +import {queryClient} from '../utils/api/queryClient'; const combineSetValues = (setRecord: Record>, keys?: string[]) => { const combinedSet = new Set(); // Create a new set to hold combined values @@ -80,13 +80,16 @@ export interface MapStore { setMapDocument: (mapDocument: DocumentObject) => void; summaryStats: { totpop?: { - data: P1TotPopSummaryStats - } - }, + data: P1TotPopSummaryStats; + }; + idealpop?: { + data: number; + }; + }; setSummaryStat: ( stat: T, value: MapStore['summaryStats'][T] - ) => void, + ) => void; // SHATTERING /** * A subset of IDs that a user is working on in a focused view. @@ -375,20 +378,20 @@ export const useMapStore = create( setFreshMap, resetZoneAssignments, upsertUserMap, - mapOptions + mapOptions, } = get(); if (currentMapDocument?.document_id === mapDocument.document_id) { return; } - setFreshMap(true) - resetZoneAssignments() - upsertUserMap({mapDocument}) - + setFreshMap(true); + resetZoneAssignments(); + upsertUserMap({mapDocument}); + set({ mapDocument: mapDocument, mapOptions: { ...mapOptions, - bounds: mapDocument.extent + bounds: mapDocument.extent, }, shatterIds: {parents: new Set(), children: new Set()}, }); @@ -398,9 +401,9 @@ export const useMapStore = create( set({ summaryStats: { ...get().summaryStats, - [stat]: value - } - }) + [stat]: value, + }, + }); }, // TODO: Refactor to something like this // featureStates: { @@ -581,13 +584,16 @@ export const useMapStore = create( delete shatterMappings[parent.parentId]; newShatterIds.parents.delete(parent.parentId); newZoneAssignments.set(parent.parentId, parent.zone!); - mapRef?.setFeatureState({ - source: BLOCK_SOURCE_ID, - id: parent.parentId, - sourceLayer: mapDocument?.parent_layer, - }, { - broken: false - }); + mapRef?.setFeatureState( + { + source: BLOCK_SOURCE_ID, + id: parent.parentId, + sourceLayer: mapDocument?.parent_layer, + }, + { + broken: false, + } + ); }); set({ @@ -633,8 +639,12 @@ export const useMapStore = create( userMaps.splice(i, 1, userMapData); // Replace the map at index i with the new data } else { const urlParams = new URL(window.location.href).searchParams; - urlParams.delete("document_id"); // Remove the document_id parameter - window.history.pushState({}, '', window.location.pathname + '?' + urlParams.toString()); // Update the URL without document_id + urlParams.delete('document_id'); // Remove the document_id parameter + window.history.pushState( + {}, + '', + window.location.pathname + '?' + urlParams.toString() + ); // Update the URL without document_id userMaps.splice(i, 1); } } @@ -783,8 +793,8 @@ export const useMapStore = create( selectedZone: 1, setSelectedZone: zone => set({selectedZone: zone}), zoneAssignments: new Map(), - assignmentsHash: "", - setAssignmentsHash: (hash) => set({ assignmentsHash: hash }), + assignmentsHash: '', + setAssignmentsHash: hash => set({assignmentsHash: hash}), accumulatedGeoids: new Set(), setAccumulatedGeoids: accumulatedGeoids => set({accumulatedGeoids}), setZoneAssignments: (zone, geoids) => { diff --git a/app/src/app/utils/api/apiHandlers.ts b/app/src/app/utils/api/apiHandlers.ts index e9886b4ae..3b475309c 100644 --- a/app/src/app/utils/api/apiHandlers.ts +++ b/app/src/app/utils/api/apiHandlers.ts @@ -1,7 +1,7 @@ import axios from 'axios'; import 'maplibre-gl'; import {useMapStore} from '@/app/store/mapStore'; -import { getEntryTotal } from '../summaryStats'; +import {getEntryTotal} from '../summaryStats'; export const FormatAssignments = () => { const assignments = Array.from(useMapStore.getState().zoneAssignments.entries()).map( @@ -172,7 +172,7 @@ export interface P1ZoneSummaryStats { black_pop: number; white_pop: number; } -export type P1TotPopSummaryStats = Omit +export type P1TotPopSummaryStats = Omit; export const P1ZoneSummaryStatsKeys = [ 'other_pop', @@ -180,8 +180,8 @@ export const P1ZoneSummaryStatsKeys = [ 'amin_pop', 'nhpi_pop', 'black_pop', - 'white_pop' -] as const + 'white_pop', +] as const; export const CleanedP1ZoneSummaryStatsKeys = [ ...P1ZoneSummaryStatsKeys, @@ -192,7 +192,7 @@ export const CleanedP1ZoneSummaryStatsKeys = [ 'nhpi_pop_pct', 'black_pop_pct', 'white_pop_pct', -] as const +] as const; export interface CleanedP1ZoneSummaryStats extends P1ZoneSummaryStats { total: number; @@ -214,23 +214,28 @@ export const getP1SummaryStats: ( ) => Promise> = async mapDocument => { if (mapDocument) { return await axios - .get>(`${process.env.NEXT_PUBLIC_API_URL}/api/document/${mapDocument.document_id}/P1`) + .get< + SummaryStatsResult + >(`${process.env.NEXT_PUBLIC_API_URL}/api/document/${mapDocument.document_id}/P1`) .then(res => { const results = res.data.results.map(row => { - const total = getEntryTotal(row) - return P1ZoneSummaryStatsKeys.reduce((acc, key) => { - acc[`${key}_pct`] = acc[key] / total; - return acc; - }, { - ...row, - total - }) as CleanedP1ZoneSummaryStats - }) - return { - ...res.data, - results - } - }) + const total = getEntryTotal(row); + return P1ZoneSummaryStatsKeys.reduce( + (acc, key) => { + acc[`${key}_pct`] = acc[key] / total; + return acc; + }, + { + ...row, + total, + } + ) as CleanedP1ZoneSummaryStats; + }); + return { + ...res.data, + results, + }; + }); } else { throw new Error('No document provided'); } @@ -246,8 +251,10 @@ export const getP1TotPopSummaryStats: ( ) => Promise> = async mapDocument => { if (mapDocument) { return await axios - .get>(`${process.env.NEXT_PUBLIC_API_URL}/api/districtrmap/summary_stats/P1/${mapDocument.parent_layer}`) - .then(res => res.data) + .get< + SummaryStatsResult + >(`${process.env.NEXT_PUBLIC_API_URL}/api/districtrmap/summary_stats/P1/${mapDocument.parent_layer}`) + .then(res => res.data); } else { throw new Error('No document provided'); } diff --git a/app/src/app/utils/api/queries.ts b/app/src/app/utils/api/queries.ts index 2549b3c0c..f4bbed4e6 100644 --- a/app/src/app/utils/api/queries.ts +++ b/app/src/app/utils/api/queries.ts @@ -13,6 +13,7 @@ import { getP1TotPopSummaryStats, P1TotPopSummaryStats, } from './apiHandlers'; +import {getEntryTotal} from '@/app/utils/summaryStats'; import {MapStore, useMapStore} from '@/app/store/mapStore'; const INITIAL_VIEW_LIMIT = 30; @@ -20,17 +21,20 @@ const INITIAL_VIEW_OFFSET = 0; /** * A utility function that returns a query function based on a nullable parameter. - * + * * @param callback - A function that takes a parameter of type ParamT and returns a Promise of type ResultT. * @param nullableParam - An optional parameter of type ParamT. If this parameter is not provided or is falsy, the function will return a function that returns null. - * - * @returns A function that, when called, will either return null (if nullableParam is not provided) + * + * @returns A function that, when called, will either return null (if nullableParam is not provided) * or call the callback function with the nullableParam and return its result. - * + * * @template ParamT - The type of the parameter that the callback function accepts. * @template ResultT - The type of the result that the callback function returns. */ -const getNullableParamQuery = (callback: (param: ParamT) => Promise, nullableParam?: ParamT) => { +const getNullableParamQuery = ( + callback: (param: ParamT) => Promise, + nullableParam?: ParamT +) => { if (!nullableParam) return () => null; return async () => await callback(nullableParam); }; @@ -71,9 +75,15 @@ export const getQueriesResultsSubs = (_useMapStore: typeof useMapStore) => { }); fetchTotPop.subscribe(response => { if (response?.data?.results) { - useMapStore.getState().setSummaryStat('totpop', { data: response.data.results}); + console.log(response?.data?.results); + useMapStore.getState().setSummaryStat('totpop', {data: response.data.results}); + useMapStore.getState().setSummaryStat('idealpop', { + data: + getEntryTotal(response.data.results) / + (useMapStore.getState().mapDocument?.num_districts ?? 1), + }); } else { - useMapStore.getState().setSummaryStat('totpop', undefined) + useMapStore.getState().setSummaryStat('totpop', undefined); } }); }; @@ -110,7 +120,7 @@ updateDocumentFromId.subscribe(mapDocument => { export const fetchAssignments = new QueryObserver(queryClient, { queryKey: ['assignments'], - queryFn: getNullableParamQuery(getAssignments) + queryFn: getNullableParamQuery(getAssignments), }); export const updateAssignments = (mapDocument: DocumentObject) => { @@ -126,10 +136,16 @@ fetchAssignments.subscribe(assignments => { } }); -export const fetchTotPop = new QueryObserver | null>(queryClient, { - queryKey: ['gerrydb_tot_pop'], - queryFn: getNullableParamQuery>(getP1TotPopSummaryStats), -}); +export const fetchTotPop = new QueryObserver | null>( + queryClient, + { + queryKey: ['gerrydb_tot_pop'], + queryFn: getNullableParamQuery< + MapStore['mapDocument'], + SummaryStatsResult + >(getP1TotPopSummaryStats), + } +); export const updateTotPop = (mapDocument: DocumentObject | null) => { fetchTotPop.setOptions({ @@ -137,4 +153,3 @@ export const updateTotPop = (mapDocument: DocumentObject | null) => { queryKey: ['gerrydb_tot_pop', mapDocument?.gerrydb_table], }); }; - diff --git a/backend/app/alembic/versions/65a4fc0a727d_add_unshatter_udf.py b/backend/app/alembic/versions/65a4fc0a727d_add_unshatter_udf.py index ab9041a44..6477c689e 100644 --- a/backend/app/alembic/versions/65a4fc0a727d_add_unshatter_udf.py +++ b/backend/app/alembic/versions/65a4fc0a727d_add_unshatter_udf.py @@ -19,7 +19,6 @@ revision: str = "65a4fc0a727d" down_revision: Union[str, None] = "dc391733e10a" branch_labels: Union[str, Sequence[str], None] = None -depends_on: Union[str, Sequence[str], None] = None def upgrade() -> None: diff --git a/backend/app/alembic/versions/c3541f016d35_add_available_stats_to_districtrmap.py b/backend/app/alembic/versions/c3541f016d35_add_available_stats_to_districtrmap.py index 091aa3bf3..89844bac1 100644 --- a/backend/app/alembic/versions/c3541f016d35_add_available_stats_to_districtrmap.py +++ b/backend/app/alembic/versions/c3541f016d35_add_available_stats_to_districtrmap.py @@ -26,7 +26,8 @@ def upgrade() -> None: ) op.execute( - sa.text(""" + sa.text( + """ UPDATE districtrmap d SET available_summary_stats = ( SELECT @@ -40,7 +41,8 @@ def upgrade() -> None: (SELECT ARRAY_AGG(summary_stat) FROM get_available_summary_stats(d.parent_layer)) END ) - """) + """ + ) ) diff --git a/backend/app/main.py b/backend/app/main.py index 104ab2be1..d207d0d66 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -94,6 +94,7 @@ async def create_document( {"gerrydb_table_name": data.gerrydb_table}, ) document_id = results.one()[0] # should be only one row, one column of results + stmt = ( select( Document.document_id, @@ -106,7 +107,9 @@ async def create_document( DistrictrMap.tiles_s3_path.label("tiles_s3_path"), # pyright: ignore DistrictrMap.num_districts.label("num_districts"), # pyright: ignore DistrictrMap.extent.label("extent"), # pyright: ignore - DistrictrMap.available_summary_stats.label("available_summary_stats"), # pyright: ignore + DistrictrMap.available_summary_stats.label( + "available_summary_stats" + ), # pyright: ignore ) .where(Document.document_id == document_id) .join( @@ -119,7 +122,7 @@ async def create_document( # Document id has a unique constraint so I'm not sure we need to hit the DB again here # more valuable would be to check that the assignments table doc = session.exec( - stmt + stmt, ).one() # again if we've got more than one, we have problems. if not doc.map_uuid: session.rollback() @@ -134,6 +137,7 @@ async def create_document( detail="Document creation failed", ) session.commit() + return doc @@ -218,10 +222,12 @@ async def reset_map(document_id: str, session: Session = Depends(get_session)): # Recreate the partition session.execute( - text(f""" + text( + f""" CREATE TABLE {partition_name} PARTITION OF document.assignments FOR VALUES IN ('{document_id}'); - """) + """ + ) ) session.commit() @@ -255,6 +261,7 @@ async def get_assignments(document_id: str, session: Session = Depends(get_sessi @app.get("/api/document/{document_id}", response_model=DocumentPublic) async def get_document(document_id: str, session: Session = Depends(get_session)): + stmt = ( select( Document.document_id, @@ -266,7 +273,9 @@ async def get_document(document_id: str, session: Session = Depends(get_session) DistrictrMap.tiles_s3_path.label("tiles_s3_path"), # pyright: ignore DistrictrMap.num_districts.label("num_districts"), # pyright: ignore DistrictrMap.extent.label("extent"), # pyright: ignore - DistrictrMap.available_summary_stats.label("available_summary_stats"), # pyright: ignore + DistrictrMap.available_summary_stats.label( + "available_summary_stats" + ), # pyright: ignore ) # pyright: ignore .where(Document.document_id == document_id) .join( diff --git a/backend/tests/test_main.py b/backend/tests/test_main.py index 724089d74..0e092126f 100644 --- a/backend/tests/test_main.py +++ b/backend/tests/test_main.py @@ -61,13 +61,15 @@ def ks_demo_view_census_blocks_fixture(session: Session): def ks_demo_view_census_blocks_districtrmap_fixture( session: Session, ks_demo_view_census_blocks_total_vap: None ): - upsert_query = text(""" + upsert_query = text( + """ INSERT INTO gerrydbtable (uuid, name, updated_at) VALUES (gen_random_uuid(), :name, now()) ON CONFLICT (name) DO UPDATE SET updated_at = now() - """) + """ + ) session.begin() session.execute(upsert_query, {"name": GERRY_DB_FIXTURE_NAME}) @@ -104,13 +106,15 @@ def ks_demo_view_census_blocks_total_vap_fixture(session: Session): def ks_demo_view_census_blocks_total_vap_districtrmap_fixture( session: Session, ks_demo_view_census_blocks_total_vap: None ): - upsert_query = text(""" + upsert_query = text( + """ INSERT INTO gerrydbtable (uuid, name, updated_at) VALUES (gen_random_uuid(), :name, now()) ON CONFLICT (name) DO UPDATE SET updated_at = now() - """) + """ + ) session.begin() session.execute(upsert_query, {"name": GERRY_DB_TOTAL_VAP_FIXTURE_NAME}) @@ -147,13 +151,15 @@ def ks_demo_view_census_blocks_no_pop_fixture(session: Session): def ks_demo_view_census_blocks_no_pop_districtrmap_fixture( session: Session, ks_demo_view_census_blocks_no_pop: None ): - upsert_query = text(""" + upsert_query = text( + """ INSERT INTO gerrydbtable (uuid, name, updated_at) VALUES (gen_random_uuid(), :name, now()) ON CONFLICT (name) DO UPDATE SET updated_at = now() - """) + """ + ) session.begin() session.execute(upsert_query, {"name": GERRY_DB_NO_POP_FIXTURE_NAME}) @@ -513,13 +519,15 @@ def ks_demo_view_census_blocks_summary_stats(session: Session): ], ) - upsert_query = text(""" + upsert_query = text( + """ INSERT INTO gerrydbtable (uuid, name, updated_at) VALUES (gen_random_uuid(), :name, now()) ON CONFLICT (name) DO UPDATE SET updated_at = now() - """) + """ + ) session.execute(upsert_query, {"name": layer}) @@ -598,13 +606,15 @@ def ks_demo_view_census_blocks_summary_stats_p4(session: Session): ], ) - upsert_query = text(""" + upsert_query = text( + """ INSERT INTO gerrydbtable (uuid, name, updated_at) VALUES (gen_random_uuid(), :name, now()) ON CONFLICT (name) DO UPDATE SET updated_at = now() - """) + """ + ) session.execute(upsert_query, {"name": layer}) diff --git a/pipelines/simple_elt/main.py b/pipelines/simple_elt/main.py index 268a24c27..8327edb28 100644 --- a/pipelines/simple_elt/main.py +++ b/pipelines/simple_elt/main.py @@ -5,10 +5,11 @@ import os import click import logging +from urllib.request import urlretrieve from urllib.parse import urlparse from subprocess import run from typing import Iterable - +import json from files import download_and_unzip_zipfile, exists_in_s3, download_file_from_s3 from settings import settings @@ -88,7 +89,8 @@ def create_county_tiles(replace: bool = False, upload: bool = False): LOGGER.info("Creating county label centroids") label_fgb = settings.OUT_SCRATCH / f"{file_name}_label.fgb" if replace or not label_fgb.exists(): - duckdb.execute(f""" + duckdb.execute( + f""" INSTALL SPATIAL; LOAD spatial; COPY ( SELECT @@ -98,7 +100,8 @@ def create_county_tiles(replace: bool = False, upload: bool = False): FROM st_read('{fgb}') ) TO '{label_fgb}' WITH (FORMAT GDAL, DRIVER 'FlatGeobuf', SRS 'EPSG:4326') - """) + """ + ) LOGGER.info("Creating county label tiles") label_tiles = settings.OUT_SCRATCH / f"{file_name}_label.pmtiles" @@ -311,5 +314,99 @@ def merge_gerrydb_tilesets( ) +@cli.command("load-districtr-v1-places") +@click.option("--replace", is_flag=True, help="Replace existing files", default=False) +def load_districtr_v1_places(replace: bool = False) -> None: + """ + Load data from districtr_v1 endpoint to the s3 bucket for later ingestion + """ + districtr_places = urlretrieve( + "https://districtr.org/assets/data/landing_pages.json?v=2", + settings.OUT_SCRATCH / "districtr_v1_places.json", + ) + + s3_client = settings.get_s3_client() + + key = f"{S3_PREFIX}/districtr_places/districtr_v1_places.json" + print(s3_client, settings.S3_BUCKET, key) + s3_client.upload_file(districtr_places, settings.S3_BUCKET, key) + + +def upsert_places_and_problems(): + """ + Upsert places and problems from districtr_v1. + WIP/not functional port of load_dv1_places_and_problems_problems in #167 + """ + + raise NotImplementedError + + s3_client = settings.get_s3_client() + + key = f"{S3_PREFIX}/districtr_places/districtr_v1_places.json" + districtr_places = download_file_from_s3(s3_client, urlparse(key)) + + if not districtr_places: + LOGGER.error("Failed to download districtr_v1_places.json") + return + + with open(districtr_places, "r") as file: + places = json.load(file) + + for place in places: + url = place["state"] + LOGGER.info(f"Downloading problems for {url}") + try: + problems = urlretrieve( + f"https://districtr.org/assets/data/modules/{place['state'].lower()}.json", + settings.OUT_SCRATCH / f"{place['state'].lower()}_problems.json", + ) + key = ( + f"{S3_PREFIX}/districtr_problems/{place['state'].lower()}_problems.json" + ) + s3_client.upload_file(problems, settings.S3_BUCKET, key) + except Exception as e: + LOGGER.error(f"Failed to download problems for {url}: {e}") + continue + + # load districtr_v1 places and problems + load_districtr_v1_places() + load_districtr_v1_problems() + + +@cli.command("load-districtr-v1-problems") +@click.option("--replace", is_flag=True, help="Replace existing files", default=False) +def load_districtr_v1_problems(replace: bool = False) -> None: + """ + load problems definition json file for states from districtr_v1 and store in s3 bucket + """ + s3_client = settings.get_s3_client() + + # check if the districtr_places object exists in s3; if not, download it using load_districtr_v1_places + key = f"{S3_PREFIX}/districtr_places/districtr_v1_places.json" + if not exists_in_s3(s3_client, settings.S3_BUCKET, key): + load_districtr_v1_places() + + districtr_places = download_file_from_s3(s3_client, urlparse(key), replace) + + with open(districtr_places, "r") as file: + places = json.load(file) + + for place in places: + url = place["state"] + LOGGER.info(f"Downloading problems for {url}") + try: + problems = urlretrieve( + f"https://districtr.org/assets/data/modules/{place['state'].lower()}.json", + settings.OUT_SCRATCH / f"{place['state'].lower()}_problems.json", + ) + key = ( + f"{S3_PREFIX}/districtr_problems/{place['state'].lower()}_problems.json" + ) + s3_client.upload_file(problems, settings.S3_BUCKET, key) + except Exception as e: + LOGGER.error(f"Failed to download problems for {url}: {e}") + continue + + if __name__ == "__main__": cli() From 8fc81fb7d330009436e52d67bb64e1bdbb21c23f Mon Sep 17 00:00:00 2001 From: Raphael Paul Laude Date: Tue, 19 Nov 2024 00:12:12 -0500 Subject: [PATCH 2/2] Add CLI command and utility function to update a DistrictrMap (#189) --- backend/app/models.py | 11 ++ backend/app/utils.py | 78 +++++++++++++- backend/cli.py | 198 ++++++++++++++++++++++-------------- backend/tests/test_utils.py | 32 ++++++ 4 files changed, 241 insertions(+), 78 deletions(-) diff --git a/backend/app/models.py b/backend/app/models.py index 16c2d07d7..b2f46d5c0 100644 --- a/backend/app/models.py +++ b/backend/app/models.py @@ -96,6 +96,17 @@ class DistrictrMapPublic(BaseModel): available_summary_stats: list[str] | None = None +class DistrictrMapUpdate(BaseModel): + gerrydb_table_name: str + name: str | None = None + parent_layer: str | None = None + child_layer: str | None = None + tiles_s3_path: str | None = None + num_districts: int | None = None + visible: bool | None = None + available_summary_stats: list[str] | None = None + + class GerryDBTable(TimeStampMixin, SQLModel, table=True): uuid: str = Field(sa_column=Column(UUIDType, unique=True, primary_key=True)) # Must correspond to the layer name in the tileset diff --git a/backend/app/utils.py b/backend/app/utils.py index 693a4a15f..f473c8153 100644 --- a/backend/app/utils.py +++ b/backend/app/utils.py @@ -1,11 +1,14 @@ -from sqlalchemy import text +from sqlalchemy import text, update from sqlalchemy import bindparam, Integer, String, Text from sqlalchemy.types import UUID from sqlmodel import Session, Float, Boolean import logging +from urllib.parse import ParseResult +import os +from app.core.config import settings -from app.models import SummaryStatisticType, UUIDType +from app.models import SummaryStatisticType, UUIDType, DistrictrMap, DistrictrMapUpdate logger = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO) @@ -74,6 +77,41 @@ def create_districtr_map( return inserted_map_uuid # pyright: ignore +def update_districtrmap( + session: Session, + gerrydb_table_name: str, + **kwargs, +): + """ + Update a districtr map. + + Args: + session: The database session. + gerrydb_table_name: The name of the gerrydb table. + **kwargs: The fields to update. + + Returns: + The updated districtr map. + """ + data = DistrictrMapUpdate(gerrydb_table_name=gerrydb_table_name, **kwargs) + update_districtrmap = data.model_dump( + exclude_unset=True, exclude={"gerrydb_table_name"}, exclude_none=True + ) + + if not update_districtrmap.keys(): + raise KeyError("No fields to update") + + stmt = ( + update(DistrictrMap) + .where(DistrictrMap.gerrydb_table_name == data.gerrydb_table_name) # pyright: ignore + .values(update_districtrmap) + .returning(DistrictrMap) + ) + (updated_districtrmap,) = session.execute(stmt).one() + + return updated_districtrmap + + def create_shatterable_gerrydb_view( session: Session, parent_layer_name: str, @@ -273,3 +311,39 @@ def add_available_summary_stats_to_districtrmap( f"Updated available summary stats for districtr map {districtr_map_uuid} to {available_summary_stats}" ) return available_summary_stats + + +def download_file_from_s3(s3, url: ParseResult, replace=False) -> str: + """ + Download a file from S3 to the local volume path. + + Args: + s3: S3 client + url (ParseResult): URL of the file to download + replace (bool): If True, replace the file if it already exists + + Returns the path to the downloaded file. + """ + if not s3: + raise ValueError("S3 client is not available") + + file_name = url.path.lstrip("/") + logger.info("File name: %s", file_name) + object_information = s3.head_object(Bucket=url.netloc, Key=file_name) + + if object_information["ResponseMetadata"]["HTTPStatusCode"] != 200: + raise ValueError( + f"GeoPackage file {file_name} not found in S3 bucket {url.netloc}" + ) + + logger.info("Downloading GerryDB view. Got response:\n%s", object_information) + + path = os.path.join(settings.VOLUME_PATH, file_name) + + if os.path.exists(path) and not replace: + logger.info("File already exists. Skipping download.") + else: + logger.info("Downloading file...") + s3.download_file(url.netloc, file_name, path) + + return path diff --git a/backend/cli.py b/backend/cli.py index 5aab2a6fa..47f469eb8 100644 --- a/backend/cli.py +++ b/backend/cli.py @@ -2,10 +2,10 @@ import click import logging -from app.main import get_session +from app.core.db import engine from app.core.config import settings import subprocess -from urllib.parse import urlparse, ParseResult +from urllib.parse import urlparse from sqlalchemy import text from app.constants import GERRY_DB_SCHEMA from app.utils import ( @@ -14,51 +14,53 @@ create_parent_child_edges as _create_parent_child_edges, add_extent_to_districtrmap as _add_extent_to_districtrmap, add_available_summary_stats_to_districtrmap as _add_available_summary_stats_to_districtrmap, + update_districtrmap as _update_districtrmap, + download_file_from_s3, ) +from functools import wraps +from contextlib import contextmanager +from sqlmodel import Session +from typing import Callable, TypeVar, Any logger = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO) -@click.group() -def cli(): - pass +T = TypeVar("T") -def download_file_from_s3(s3, url: ParseResult, replace=False) -> str: - """ - Download a file from S3 to the local volume path. +@contextmanager +def session_scope(): + """Provide a transactional scope around a series of operations.""" + session = Session(engine) + try: + yield session + session.commit() + except Exception: + session.rollback() + raise + finally: + session.close() - Args: - s3: S3 client - url (ParseResult): URL of the file to download - replace (bool): If True, replace the file if it already exists - Returns the path to the downloaded file. +def with_session(f: Callable[..., T]) -> Callable[..., T]: + """ + Decorator that handles database session creation and cleanup. + Compatible with Click commands. """ - if not s3: - raise ValueError("S3 client is not available") - - file_name = url.path.lstrip("/") - logger.info("File name: %s", file_name) - object_information = s3.head_object(Bucket=url.netloc, Key=file_name) - - if object_information["ResponseMetadata"]["HTTPStatusCode"] != 200: - raise ValueError( - f"GeoPackage file {file_name} not found in S3 bucket {url.netloc}" - ) - logger.info("Downloading GerryDB view. Got response:\n%s", object_information) + @wraps(f) + def decorator(*args: Any, **kwargs: Any) -> T: + with session_scope() as session: + kwargs["session"] = session + return f(*args, **kwargs) - path = os.path.join(settings.VOLUME_PATH, file_name) + return decorator - if os.path.exists(path) and not replace: - logger.info("File already exists. Skipping download.") - else: - logger.info("Downloading file...") - s3.download_file(url.netloc, file_name, path) - return path +@click.group() +def cli(): + pass @cli.command("import-gerrydb-view") @@ -66,7 +68,10 @@ def download_file_from_s3(s3, url: ParseResult, replace=False) -> str: @click.option("--gpkg", "-g", help="Path or URL to GeoPackage file", required=True) @click.option("--replace", "-f", help="Replace the file if it exists", is_flag=True) @click.option("--rm", "-r", help="Delete file after loading to postgres", is_flag=True) -def import_gerrydb_view(layer: str, gpkg: str, replace: bool, rm: bool): +@with_session +def import_gerrydb_view( + session: Session, layer: str, gpkg: str, replace: bool, rm: bool +): logger.info("Importing GerryDB view...") url = urlparse(gpkg) @@ -110,9 +115,6 @@ def import_gerrydb_view(layer: str, gpkg: str, replace: bool, rm: bool): logger.info("GerryDB view imported successfully") - _session = get_session() - session = next(_session) - upsert_query = text( """ INSERT INTO gerrydbtable (uuid, name, updated_at) @@ -123,50 +125,39 @@ def import_gerrydb_view(layer: str, gpkg: str, replace: bool, rm: bool): """ ) - try: - session.execute( - upsert_query, - { - "name": layer, - }, - ) - session.commit() - logger.info("GerryDB view upserted successfully.") - except Exception as e: - session.rollback() - logger.error("Failed to upsert GerryDB view. Got %s", e) - raise ValueError(f"Failed to upsert GerryDB view. Got {e}") - - session.close() + session.execute( + upsert_query, + { + "name": layer, + }, + ) + logger.info("GerryDB view upserted successfully.") @cli.command("create-parent-child-edges") @click.option("--districtr-map", "-d", help="Districtr map name", required=True) -def create_parent_child_edges(districtr_map: str): +@with_session +def create_parent_child_edges(session: Session, districtr_map: str): logger.info("Creating parent-child edges...") - session = next(get_session()) stmt = text( "SELECT uuid FROM districtrmap WHERE gerrydb_table_name = :districtrmap_name" ) (districtr_map_uuid,) = session.execute( stmt, params={"districtrmap_name": districtr_map} ).one() - print(f"Found districtmap uuid: {districtr_map_uuid}") + logger.info(f"Found districtmap uuid: {districtr_map_uuid}") + _create_parent_child_edges(session=session, districtr_map_uuid=districtr_map_uuid) - session.commit() logger.info("Parent-child relationship upserted successfully.") - session.close() - @cli.command("delete-parent-child-edges") @click.option("--districtr-map", "-d", help="Districtr map name", required=True) -def delete_parent_child_edges(districtr_map: str): +@with_session +def delete_parent_child_edges(session: Session, districtr_map: str): logger.info("Deleting parent-child edges...") - session = next(get_session()) - delete_query = text( """ DELETE FROM parentchildedges @@ -179,11 +170,8 @@ def delete_parent_child_edges(districtr_map: str): "districtr_map": districtr_map, }, ) - session.commit() logger.info("Parent-child relationship upserted successfully.") - session.close() - @cli.command("create-districtr-map") @click.option("--name", help="Name of the districtr map", required=True) @@ -204,7 +192,9 @@ def delete_parent_child_edges(districtr_map: str): default=None, nargs=4, ) +@with_session def create_districtr_map( + session: Session, name: str, parent_layer_name: str, child_layer_name: str | None, @@ -215,7 +205,6 @@ def create_districtr_map( bounds: list[float] | None = None, ): logger.info("Creating districtr map...") - session = next(get_session()) (districtr_map_uuid,) = _create_districtr_map( session=session, name=name, @@ -238,28 +227,88 @@ def create_districtr_map( session=session, districtr_map_uuid=districtr_map_uuid ) - session.commit() logger.info(f"Districtr map created successfully {districtr_map_uuid}") +@cli.command("update-districtr-map") +@click.option( + "--gerrydb-table-name", + "-n", + help="Name of the GerryDB table", + type=str, + required=True, +) +@click.option("--name", help="Name of the districtr map", type=str, required=False) +@click.option( + "--parent-layer-name", help="Parent gerrydb layer name", type=str, required=False +) +@click.option( + "--child-layer-name", help="Child gerrydb layer name", type=str, required=False +) +@click.option("--num-districts", help="Number of districts", type=str, required=False) +@click.option( + "--tiles-s3-path", help="S3 path to the tileset", type=str, required=False +) +@click.option("--visibility", "-v", help="Visibility", type=bool, required=False) +@click.option( + "--bounds", + "-b", + help="Bounds of the extent as `--bounds x_min y_min x_max y_max`", + required=False, + type=float, + default=None, + nargs=4, +) +@with_session +def update_districtr_map( + session: Session, + gerrydb_table_name: str, + name: str | None, + parent_layer_name: str | None, + child_layer_name: str | None, + num_districts: int | None, + tiles_s3_path: str | None, + visibility: bool = False, + bounds: list[float] | None = None, +): + logger.info("Updating districtr map...") + + _bounds = None + if bounds and len(bounds) == 4: + _bounds = bounds + + result = _update_districtrmap( + session=session, + gerrydb_table_name=gerrydb_table_name, + name=name, + parent_layer=parent_layer_name, + child_layer=child_layer_name, + num_districts=num_districts, + tiles_s3_path=tiles_s3_path, + visible=visibility, + bounds=_bounds, + ) + logger.info(f"Districtr map updated successfully {result}") + + @cli.command("create-shatterable-districtr-view") @click.option("--parent-layer-name", help="Parent gerrydb layer name", required=True) @click.option("--child-layer-name", help="Child gerrydb layer name", required=False) @click.option("--gerrydb-table-name", help="Name of the GerryDB table", required=False) +@with_session def create_shatterable_gerrydb_view( + session: Session, parent_layer_name: str, child_layer_name: str, gerrydb_table_name: str, ): logger.info("Creating materialized shatterable gerrydb view...") - session = next(get_session()) inserted_uuid = _create_shatterable_gerrydb_view( session=session, parent_layer_name=parent_layer_name, child_layer_name=child_layer_name, gerrydb_table_name=gerrydb_table_name, ) - session.commit() logger.info( f"Materialized shatterable gerrydb view created successfully {inserted_uuid}" ) @@ -276,10 +325,12 @@ def create_shatterable_gerrydb_view( default=None, nargs=4, ) -def add_extent_to_districtr_map(districtr_map: str, bounds: list[float] | None = None): +@with_session +def add_extent_to_districtr_map( + session: Session, districtr_map: str, bounds: list[float] | None = None +): logger.info(f"User provided bounds: {bounds}") - session = next(get_session()) stmt = text( "SELECT uuid FROM districtrmap WHERE gerrydb_table_name = :districtrmap_name" ) @@ -291,16 +342,13 @@ def add_extent_to_districtr_map(districtr_map: str, bounds: list[float] | None = _add_extent_to_districtrmap( session=session, districtr_map_uuid=districtr_map_uuid, bounds=bounds ) - session.commit() logger.info("Updated extent successfully.") - session.close() - @cli.command("add-available-summary-stats-to-districtr-map") @click.option("--districtr-map", "-d", help="Districtr map name", required=True) -def add_available_summary_stats_to_districtr_map(districtr_map: str): - session = next(get_session()) +@with_session +def add_available_summary_stats_to_districtr_map(session: Session, districtr_map: str): stmt = text( "SELECT uuid FROM districtrmap WHERE gerrydb_table_name = :districtrmap_name" ) @@ -313,9 +361,7 @@ def add_available_summary_stats_to_districtr_map(districtr_map: str): session=session, districtr_map_uuid=districtr_map_uuid ) - session.commit() logger.info("Updated available summary stats successfully.") - session.close() if __name__ == "__main__": diff --git a/backend/tests/test_utils.py b/backend/tests/test_utils.py index 7e03b154a..e730ff424 100644 --- a/backend/tests/test_utils.py +++ b/backend/tests/test_utils.py @@ -6,10 +6,12 @@ create_parent_child_edges, add_extent_to_districtrmap, get_available_summary_stats, + update_districtrmap, ) from sqlmodel import Session import subprocess from app.constants import GERRY_DB_SCHEMA +from app.models import DistrictrMap from tests.constants import OGR2OGR_PG_CONNECTION_STRING, FIXTURES_PATH from sqlalchemy import text @@ -193,6 +195,36 @@ def test_create_districtr_map_some_nulls(session: Session, simple_parent_geos_ge session.commit() +@pytest.fixture(name="simple_parent_geos_districtrmap") +def simple_parent_geos_districtrmap_fixture( + session: Session, simple_parent_geos_gerrydb, simple_child_geos_gerrydb +): + gerrydb_name = "simple_geos_test" + (inserted_districtr_map,) = create_districtr_map( + session, + name="Simple shatterable layer", + gerrydb_table_name=gerrydb_name, + num_districts=10, + tiles_s3_path="tilesets/simple_shatterable_layer.pmtiles", + parent_layer_name="simple_parent_geos", + child_layer_name="simple_child_geos", + visibility=True, + ) + session.commit() + return gerrydb_name + + +def test_update_districtr_map(session: Session, simple_parent_geos_districtrmap): + result = update_districtrmap( + session=session, + gerrydb_table_name=simple_parent_geos_districtrmap, + visible=False, + ) + session.commit() + districtr_map = DistrictrMap.model_validate(result) + assert not districtr_map.visible + + def test_add_extent_to_districtrmap(session: Session, simple_parent_geos_gerrydb): (inserted_districtr_map,) = create_districtr_map( session,