From 90ea21f86733cb115596c2a6b4af8d3fe4a93aee Mon Sep 17 00:00:00 2001 From: David Mikola Date: Thu, 30 Jul 2026 10:26:49 -0500 Subject: [PATCH 1/4] [DEV-15416] Added basic whitelist using loaderfactory column mappings --- pyproject.toml | 1 + .../population_data_loaders/loaders.py | 48 +++++++++++-------- .../unit/test_population_data_loaders.py | 41 +++++++++++----- 3 files changed, 59 insertions(+), 31 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 8b902a5ea4..f04239b640 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -216,6 +216,7 @@ pycodestyle.max-line-length = 120 "ANN001", # missing-type-function-argument; avoid conflict with fixtures "ANN201", # missing-return-type-undocumented-public-function; avoid conflict with test case return values "ANN202", # missing-return-type-undocumented-private-function; avoid conflict with test case return values + "PLR0911", # too-many-return-statements; avoid conflict with too many return statements "PLR0913", # too-many-arguments; avoid conflict with too many fixtures "PLR0915", # too-many-statements; avoid conflict with long fixtures ] diff --git a/usaspending_api/references/management/commands/population_data_loaders/loaders.py b/usaspending_api/references/management/commands/population_data_loaders/loaders.py index afc8b0b3dd..1259a357b5 100644 --- a/usaspending_api/references/management/commands/population_data_loaders/loaders.py +++ b/usaspending_api/references/management/commands/population_data_loaders/loaders.py @@ -1,8 +1,10 @@ from abc import ABC, abstractmethod -from django.db import connection from logging import Logger from typing import List -from django.db import models + +from django.db import connection +from django.db.models import Model + from usaspending_api.references.models import PopCongressionalDistrict, PopCounty from usaspending_api.references.models.ref_country_code import RefCountryCode @@ -31,7 +33,7 @@ def create_tables(self, columns: List[str]) -> None: pass @abstractmethod - def load_data(self, data: List[dict], model: models = None) -> None: + def load_data(self, data: List[dict], model: Model = None) -> None: """Handles the loading of population data. Args: @@ -58,24 +60,32 @@ class GenericPopulationLoader(Loader): TEMP_TABLE_NAME = "temp_population_load" TEMP_TABLE_SQL = "CREATE TABLE {table} ({columns});" - def __init__(self, column_mapper, logger: Logger): + def __init__(self, column_mapper: dict[str, str], logger: Logger): self._columns_mapper = column_mapper self._logger = logger - def drop_temp_tables(self): + def drop_temp_tables(self) -> None: self._logger.info(f"Dropping temp table {self.TEMP_TABLE_NAME}") with connection.cursor() as cursor: cursor.execute(f"DROP TABLE IF EXISTS {self.TEMP_TABLE_NAME}") - def create_tables(self, columns): + def create_tables(self, columns: list[str]) -> None: self._logger.info(f"Creating temp table {self.TEMP_TABLE_NAME}") - with connection.cursor() as cursor: - cursor.execute(f"DROP TABLE IF EXISTS {self.TEMP_TABLE_NAME}") - cursor.execute( - self.TEMP_TABLE_SQL.format(table=self.TEMP_TABLE_NAME, columns=",".join([f"{c} TEXT" for c in columns])) - ) - def load_data(self, data, model=None): + # checks to see if the column is an expected value + approved_cols = [c for c in columns if c in self._columns_mapper.keys()] + + if len(approved_cols) > 0: + with connection.cursor() as cursor: + cursor.execute(f"DROP TABLE IF EXISTS {self.TEMP_TABLE_NAME}") + cursor.execute( + self.TEMP_TABLE_SQL.format( + table=self.TEMP_TABLE_NAME, + columns=",".join([f"{c} TEXT" for c in approved_cols]) + ) + ) + + def load_data(self, data: list[dict], model: Model = None) -> None: self._logger.info(f"Attempting to load {len(data)} records into {model.__name__}") model.objects.all().delete() model.objects.bulk_create( @@ -83,7 +93,7 @@ def load_data(self, data, model=None): ) self._logger.info("Success? Please Verify") - def cleanup(self): + def cleanup(self) -> None: self.drop_temp_tables() @@ -92,7 +102,7 @@ class DistrictPopulationLoader(GenericPopulationLoader): district population data. """ - def load_data(self, data, model=None): + def load_data(self, data: list[dict], model: Model = None) -> None: model = PopCongressionalDistrict super().load_data(data, model) @@ -102,7 +112,7 @@ class CountyPopulationLoader(GenericPopulationLoader): county population data. """ - def load_data(self, data, model=None): + def load_data(self, data: list[dict], model: Model = None) -> None: model = PopCounty super().load_data(data, model) @@ -114,7 +124,7 @@ class CountryPopulationLoader(GenericPopulationLoader): ERROR_THRESHOLD = 0.5 - def load_data(self, data, model=None): + def load_data(self, data: list[dict], model: Model = None) -> None: model = RefCountryCode # Ensuring that the update will result under our accepted error threshold total_countries = model.objects.all().distinct("country_code") @@ -139,14 +149,14 @@ def load_data(self, data, model=None): record.save() self._logger.info("Success? Please Verify") - def drop_temp_tables(self): + def drop_temp_tables(self) -> None: # This implementation is purely an update, no need to drop temp tables pass - def create_tables(self, columns): + def create_tables(self, columns: list[str]) -> None: # This implementation is purely an update, no need to create temp tables pass - def cleanup(self): + def cleanup(self) -> None: # This implementation is purely an update, no need to cleanup temp tables pass diff --git a/usaspending_api/references/tests/unit/test_population_data_loaders.py b/usaspending_api/references/tests/unit/test_population_data_loaders.py index 74273ca7bf..bd4289b26b 100644 --- a/usaspending_api/references/tests/unit/test_population_data_loaders.py +++ b/usaspending_api/references/tests/unit/test_population_data_loaders.py @@ -1,6 +1,7 @@ -from unittest.mock import patch, MagicMock +from unittest.mock import MagicMock, patch import pytest + from usaspending_api.references.management.commands.population_data_loaders.loaders import ( CountryPopulationLoader, CountyPopulationLoader, @@ -26,7 +27,7 @@ def test_country_population_loader_drop_temp_tables(connection_mock): logger = MagicMock() loader = CountryPopulationLoader(column_mapper, logger) loader.drop_temp_tables() - cursor_mock.execute.call_count == 0 + assert cursor_mock.execute.call_count == 0 @patch("usaspending_api.references.management.commands.population_data_loaders.loaders.connection") @@ -53,22 +54,28 @@ def test_district_population_loader_drop_temp_tables(connection_mock): def test_generic_population_loader_create_tables(connection_mock): cursor_mock = connection_mock.cursor.return_value.__enter__.return_value column_mapper = MagicMock() - test_data_cols = ["test"] + test_data_cols = ["state_code", "test"] logger = MagicMock() loader = GenericPopulationLoader(column_mapper, logger) loader.create_tables(test_data_cols) - cursor_mock.execute.assert_called_with(f"CREATE TABLE {GenericPopulationLoader.TEMP_TABLE_NAME} (test TEXT);") + assert cursor_mock.execute.call_count == 0 @patch("usaspending_api.references.management.commands.population_data_loaders.loaders.connection") def test_county_population_loader_create_tables(connection_mock): cursor_mock = connection_mock.cursor.return_value.__enter__.return_value - column_mapper = MagicMock() - test_data_cols = ["test"] + column_mapper = { + "state_code": "state_code", + "county_code": "county_number", + "state_name": "state_name", + "county_name": "county_name", + "population": "latest_population", + } + test_data_cols = ["state_code", "test"] logger = MagicMock() loader = CountyPopulationLoader(column_mapper, logger) loader.create_tables(test_data_cols) - cursor_mock.execute.assert_called_with(f"CREATE TABLE {GenericPopulationLoader.TEMP_TABLE_NAME} (test TEXT);") + cursor_mock.execute.assert_called_with(f"CREATE TABLE {GenericPopulationLoader.TEMP_TABLE_NAME} (state_code TEXT);") @patch("usaspending_api.references.management.commands.population_data_loaders.loaders.connection") @@ -79,18 +86,24 @@ def test_country_population_loader_create_tables(connection_mock): logger = MagicMock() loader = CountryPopulationLoader(column_mapper, logger) loader.create_tables(test_data_cols) - cursor_mock.execute.call_count == 0 + assert cursor_mock.execute.call_count == 0 @patch("usaspending_api.references.management.commands.population_data_loaders.loaders.connection") def test_district_population_loader_create_tables(connection_mock): cursor_mock = connection_mock.cursor.return_value.__enter__.return_value - column_mapper = MagicMock() - test_data_cols = ["test"] + column_mapper = { + "state_code": "state_code", + "state_name": "state_name", + "state_abbreviation": "state_abbreviation", + "congressional_district": "congressional_district", + "population": "latest_population", + } + test_data_cols = ["state_code", "test"] logger = MagicMock() loader = DistrictPopulationLoader(column_mapper, logger) loader.create_tables(test_data_cols) - cursor_mock.execute.assert_called_with(f"CREATE TABLE {GenericPopulationLoader.TEMP_TABLE_NAME} (test TEXT);") + cursor_mock.execute.assert_called_with(f"CREATE TABLE {GenericPopulationLoader.TEMP_TABLE_NAME} (state_code TEXT);") @patch("usaspending_api.references.management.commands.population_data_loaders.loaders.connection") @@ -106,7 +119,11 @@ def test_country_population_loader_load_data_missing_countries_error(django_mode loader.load_data(data=test_data) assert ( str(err.value) - == "The provided data contains less than 50% of the known countries. We require at least 50% of the known countries to be present in the file to load. Data had 1 countries when there are 3 countries in RefCountryCode." + == ( + "The provided data contains less than 50% of the known countries. " + "We require at least 50% of the known countries to be present in the file to load. " + "Data had 1 countries when there are 3 countries in RefCountryCode." + ) ) From c2217a58f7ad23ae60742f6fdf17b8160cd870a5 Mon Sep 17 00:00:00 2001 From: David Mikola Date: Mon, 3 Aug 2026 08:24:58 -0500 Subject: [PATCH 2/4] [Dev-15416] Updated chek for valid columns --- .../management/commands/population_data_loaders/loaders.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/usaspending_api/references/management/commands/population_data_loaders/loaders.py b/usaspending_api/references/management/commands/population_data_loaders/loaders.py index 1259a357b5..eace49a1b6 100644 --- a/usaspending_api/references/management/commands/population_data_loaders/loaders.py +++ b/usaspending_api/references/management/commands/population_data_loaders/loaders.py @@ -75,7 +75,7 @@ def create_tables(self, columns: list[str]) -> None: # checks to see if the column is an expected value approved_cols = [c for c in columns if c in self._columns_mapper.keys()] - if len(approved_cols) > 0: + if len(approved_cols) == len(columns): with connection.cursor() as cursor: cursor.execute(f"DROP TABLE IF EXISTS {self.TEMP_TABLE_NAME}") cursor.execute( From f03b4a41e4c6c62f1b1e8983d6fc4138f97f6f59 Mon Sep 17 00:00:00 2001 From: David Mikola Date: Mon, 3 Aug 2026 10:52:43 -0500 Subject: [PATCH 3/4] Updated Loader Create Tables tests --- .../unit/test_population_data_loaders.py | 31 ++++++++++++++++--- 1 file changed, 27 insertions(+), 4 deletions(-) diff --git a/usaspending_api/references/tests/unit/test_population_data_loaders.py b/usaspending_api/references/tests/unit/test_population_data_loaders.py index bd4289b26b..ec2e0c8890 100644 --- a/usaspending_api/references/tests/unit/test_population_data_loaders.py +++ b/usaspending_api/references/tests/unit/test_population_data_loaders.py @@ -71,11 +71,14 @@ def test_county_population_loader_create_tables(connection_mock): "county_name": "county_name", "population": "latest_population", } - test_data_cols = ["state_code", "test"] + test_data_cols = column_mapper.keys() logger = MagicMock() loader = CountyPopulationLoader(column_mapper, logger) loader.create_tables(test_data_cols) - cursor_mock.execute.assert_called_with(f"CREATE TABLE {GenericPopulationLoader.TEMP_TABLE_NAME} (state_code TEXT);") + cursor_mock.execute.assert_called_with( + f"CREATE TABLE {GenericPopulationLoader.TEMP_TABLE_NAME} " + + "(state_code TEXT,county_code TEXT,state_name TEXT,county_name TEXT,population TEXT);" + ) @patch("usaspending_api.references.management.commands.population_data_loaders.loaders.connection") @@ -99,11 +102,31 @@ def test_district_population_loader_create_tables(connection_mock): "congressional_district": "congressional_district", "population": "latest_population", } - test_data_cols = ["state_code", "test"] + test_data_cols = column_mapper.keys() logger = MagicMock() loader = DistrictPopulationLoader(column_mapper, logger) loader.create_tables(test_data_cols) - cursor_mock.execute.assert_called_with(f"CREATE TABLE {GenericPopulationLoader.TEMP_TABLE_NAME} (state_code TEXT);") + cursor_mock.execute.assert_called_with( + f"CREATE TABLE {GenericPopulationLoader.TEMP_TABLE_NAME}" + + "(state_code TEXT,state_name TEXT,state_abbreviation TEXT,congressional_district TEXT,population TEXT);" + ) + + +@patch("usaspending_api.references.management.commands.population_data_loaders.loaders.connection") +def test_population_loader_create_tables_fail(connection_mock): + cursor_mock = connection_mock.cursor.return_value.__enter__.return_value + column_mapper = { + "state_code": "state_code", + "state_name": "state_name", + "state_abbreviation": "state_abbreviation", + "congressional_district": "congressional_district", + "population": "latest_population", + } + test_data_cols = ["test"] + logger = MagicMock() + loader = GenericPopulationLoader(column_mapper, logger) + loader.create_tables(test_data_cols) + assert cursor_mock.execute.call_count == 0 @patch("usaspending_api.references.management.commands.population_data_loaders.loaders.connection") From 298ae56b45a8b3ede1092a697b9e8f2e9b144ae5 Mon Sep 17 00:00:00 2001 From: DavidMikolaKC Date: Mon, 3 Aug 2026 11:20:32 -0500 Subject: [PATCH 4/4] Fix loader test syntax error --- .../references/tests/unit/test_population_data_loaders.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/usaspending_api/references/tests/unit/test_population_data_loaders.py b/usaspending_api/references/tests/unit/test_population_data_loaders.py index ec2e0c8890..69091f39ac 100644 --- a/usaspending_api/references/tests/unit/test_population_data_loaders.py +++ b/usaspending_api/references/tests/unit/test_population_data_loaders.py @@ -107,7 +107,7 @@ def test_district_population_loader_create_tables(connection_mock): loader = DistrictPopulationLoader(column_mapper, logger) loader.create_tables(test_data_cols) cursor_mock.execute.assert_called_with( - f"CREATE TABLE {GenericPopulationLoader.TEMP_TABLE_NAME}" + + f"CREATE TABLE {GenericPopulationLoader.TEMP_TABLE_NAME} " + "(state_code TEXT,state_name TEXT,state_abbreviation TEXT,congressional_district TEXT,population TEXT);" )