[DEV-15416] CSV header names interpolated into CREATE TABLE

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
 ]

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,32 +60,40 @@ 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) == len(columns):
+            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(
             [model(**{col: row[csv] for csv, col in self._columns_mapper.items()}) for row in data]
         )
         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

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,31 @@ 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 = 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} (test 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")
@@ -79,18 +89,44 @@ 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 = 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} (test 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")
@@ -106,7 +142,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."
+        )
     )