diff --git a/scripts/us_census/pep/annual_population/manifest.json b/scripts/us_census/pep/annual_population/manifest.json index 7655d3c516..4346cb61a2 100644 --- a/scripts/us_census/pep/annual_population/manifest.json +++ b/scripts/us_census/pep/annual_population/manifest.json @@ -12,7 +12,8 @@ "preprocess.py --mode=process" ], "source_files": [ - "input_files/*" + "input_files/*", + "golden_data/*.csv" ], "import_inputs": [ { diff --git a/scripts/us_census/pep/annual_population/preprocess.py b/scripts/us_census/pep/annual_population/preprocess.py index a13d5612d6..f711b9c911 100644 --- a/scripts/us_census/pep/annual_population/preprocess.py +++ b/scripts/us_census/pep/annual_population/preprocess.py @@ -43,6 +43,7 @@ import requests import time import re +import tempfile from datetime import datetime as dt warnings.filterwarnings('ignore') @@ -51,6 +52,8 @@ from absl import logging from absl import flags from retry import retry +from requests.adapters import HTTPAdapter +from urllib3.util.retry import Retry _FLAGS = flags.FLAGS @@ -66,6 +69,7 @@ def _define_flags(): _INPUT_FILE_PATH = os.path.join(_MODULE_DIR, 'input_files') os.makedirs(_INPUT_FILE_PATH, exist_ok=True) excel_file_name_pattern = r"NST-EST(\d{4})-POP\.xlsx" +popchg_file_name_pattern = r"NST-EST\d{4}-POPCHG2020_\d{4}\.csv" sys.path.insert(1, _MODULE_DIR) # pylint: disable=wrong-import-position # pylint: disable=import-error @@ -486,19 +490,40 @@ def _process_nationals_2021(file_path: str) -> pd.DataFrame: return data_df +def _extract_future_years(values, min_year: int = 2023) -> list: + """Extracts year strings > min_year from header values (handling int, float, and str).""" + years = [] + for x in values: + if pd.isna(x): + continue + if isinstance(x, (int, float)): + yr = int(x) + elif isinstance(x, str): + s = x.strip() + if re.fullmatch(r"\d{4}(?:\.0)?", s): + yr = int(float(s)) + else: + continue + else: + continue + if yr > min_year: + years.append(str(yr)) + return years + + def _process_nationals_2029(file_path: str) -> pd.DataFrame: data_df = pd.DataFrame() data_df = _load_data_df(path=file_path, file_format="xlsx", header=3) df_cols = ["Location", "042020", "072020", "072021", "2022", "2023"] unpivot_cols = ["2022", "2023"] # extend df_cols & unpivot_cols with all years > 2023 - newly_added_years = [ - str(x) - for x in data_df.columns.to_list() - if isinstance(x, int) and x > 2023 - ] + newly_added_years = _extract_future_years(data_df.columns, min_year=2023) df_cols.extend(newly_added_years) unpivot_cols.extend(newly_added_years) + if len(df_cols) != len(data_df.columns): + raise ValueError( + f"Column count mismatch in {file_path}: expected {len(df_cols)} columns ({df_cols}), " + f"found {len(data_df.columns)} columns ({list(data_df.columns)})") data_df.columns = df_cols data_df = data_df[data_df["Location"] == USA] data_df = _unpivot_data_df(data_df, ["Location"], unpivot_cols) @@ -541,23 +566,21 @@ def _process_states_2029(file_path: str) -> pd.DataFrame: "Region", "ignore_col1", "ignore_col2", "ignore_col3", "2022", "2023" ] unpivot_cols = ["2022", "2023"] - try: - # extend df_cols & unpivot_cols with all years > 2023 - year_column_row = _load_data_df( - path=file_path, file_format="xlsx", - header=0).iloc[2:3].values.flatten().tolist() - if year_column_row: - newly_added_years = [ - str(int(x)) - for x in year_column_row - if isinstance(x, float) and x > 2023 - ] - df_cols.extend(newly_added_years) - unpivot_cols.extend(newly_added_years) - except: - # this error can be ignored - pass + # extend df_cols & unpivot_cols with all years > 2023 + header_df = _load_data_df(path=file_path, file_format="xlsx", header=0) + if len(header_df) <= 2: + raise ValueError( + f"Unable to extract dynamic year columns from {file_path}: header row 2 missing" + ) + year_column_row = header_df.iloc[2:3].values.flatten().tolist() + newly_added_years = _extract_future_years(year_column_row, min_year=2023) + df_cols.extend(newly_added_years) + unpivot_cols.extend(newly_added_years) data_df = _load_data_df(path=file_path, file_format="xlsx", header=8) + if len(df_cols) != len(data_df.columns): + raise ValueError( + f"Column count mismatch in {file_path}: expected {len(df_cols)} columns ({df_cols}), " + f"found {len(data_df.columns)} columns") data_df.columns = df_cols data_df["Region"] = data_df["Region"].str.replace(".", "", regex=False) @@ -999,7 +1022,7 @@ def process(input_path, cleaned_csv_file_path: str, mcf_file_path: str, _process_nationals_2010_2020(file, op_file) data_df = _load_data_df(op_file, "csv", 0) os.remove(op_file) - elif "NST-EST2023-POPCHG2020_2023.csv" in file: + elif re.search(popchg_file_name_pattern, file_name): _process_nationals_2020_2029(file, op_file) data_df = _load_data_df(op_file, "csv", 0) os.remove(op_file) @@ -1046,10 +1069,9 @@ def process(input_path, cleaned_csv_file_path: str, mcf_file_path: str, if not data_df.empty: processed_count += 1 - # data_df['len'] = data_df['Location'].astype(str).map(len) - # geo_code_len = data_df['len'].unique() - # if 13 in geo_code_len: - # pass + vintage_matches = re.findall(r"(?:19|20)\d{2}", file_name) + data_df["_vintage"] = max([int(y) for y in vintage_matches], + default=0) final_df = pd.concat([final_df, data_df], axis=0) else: logging.error(f"Failed to process {file}") @@ -1064,8 +1086,9 @@ def process(input_path, cleaned_csv_file_path: str, mcf_file_path: str, f"Mismatch between files expected ({total_files_to_process}) and processed ({processed_count})" ) final_df["Year"] = final_df["Year"].astype("int") - final_df = final_df.sort_values(by=["Location", "Year"]) - final_df = final_df.drop_duplicates(["Year", "Location"]) + final_df = final_df.sort_values(by=["Location", "Year", "_vintage"], + kind="mergesort") + final_df = final_df.drop_duplicates(["Year", "Location"], keep="last") final_df["Count_Person"] = final_df["Count_Person"].astype("int") final_df[["Year", "Location", "Count_Person"]].to_csv(cleaned_csv_file_path, index=False) @@ -1075,6 +1098,35 @@ def process(input_path, cleaned_csv_file_path: str, mcf_file_path: str, logging.fatal("No files were successfully processed.") +def _create_retry_session() -> requests.Session: + """Creates a requests Session configured with retry and backoff for transient errors.""" + session = requests.Session() + session.headers.update({'User-Agent': 'Mozilla/5.0'}) + retries = Retry( + total=5, + backoff_factor=2, + status_forcelist=[408, 429, 500, 502, 503, 504, 520], + allowed_methods=["HEAD", "GET"], + respect_retry_after_header=True, + ) + adapter = HTTPAdapter(max_retries=retries) + session.mount("https://", adapter) + session.mount("http://", adapter) + return session + + +@retry(tries=3, delay=2, backoff=2, exceptions=(ValueError,)) +def _head_with_html_retry(session: requests.Session, + url: str, + timeout: int = 30) -> requests.Response: + """Performs a HEAD request and retries if HTTP 200 returns an HTML error page.""" + response = session.head(url, allow_redirects=True, timeout=timeout) + content_type = response.headers.get('Content-Type', '') + if response.status_code == 200 and 'html' in content_type.lower(): + raise ValueError(f"Server returned HTTP 200 HTML response for {url}") + return response + + def add_future_year_urls(): """ This method will generate future URLs specified in the urls_to_scan object. If valid URL, it will append to the _FILES_TO_DOWNLOAD global variable. @@ -1083,33 +1135,74 @@ def add_future_year_urls(): Returns: None """ global _FILES_TO_DOWNLOAD - with open(os.path.join(_MODULE_DIR, 'input_url.json'), 'r') as inpit_file: - _FILES_TO_DOWNLOAD = json.load(inpit_file) + with open(os.path.join(_MODULE_DIR, 'input_url.json'), 'r') as input_file: + _FILES_TO_DOWNLOAD = json.load(input_file) urls_to_scan = [ - "https://www2.census.gov/programs-surveys/popest/datasets/2020-{YEAR}/cities/totals/sub-est{YEAR}.csv", - "https://www2.census.gov/programs-surveys/popest/datasets/2020-{YEAR}/counties/totals/co-est{YEAR}-alldata.csv", - "https://www2.census.gov/programs-surveys/popest/tables/2020-{YEAR}/state/totals/NST-EST{YEAR}-POP.xlsx", - "https://www2.census.gov/programs-surveys/popest/datasets/2020-{YEAR}/state/totals/NST-EST{YEAR}-POPCHG2020_{YEAR}.csv" + ("https://www2.census.gov/programs-surveys/popest/datasets/2020-{YEAR}/cities/totals/sub-est{YEAR}.csv", + 2025), + ("https://www2.census.gov/programs-surveys/popest/datasets/2020-{YEAR}/counties/totals/co-est{YEAR}-alldata.csv", + 2025), + ("https://www2.census.gov/programs-surveys/popest/tables/2020-{YEAR}/state/totals/NST-EST{YEAR}-POP.xlsx", + 2025), + ("https://www2.census.gov/programs-surveys/popest/datasets/2020-{YEAR}/state/totals/NST-EST{YEAR}-POPCHG2020_{YEAR}.csv", + 2023), ] - # This method will get the latest year URLs available for the years 2023 to 2030 - for url in urls_to_scan: - for future_year in range(2030, 2022, -1): + session = _create_retry_session() + # This method will get the latest year URLs available for the years baseline_year to 2030 + for url, baseline_year in urls_to_scan: + found = False + for future_year in range(2030, baseline_year - 1, -1): url_to_check = url.format(YEAR=future_year) try: - logging.info(f"checking furute url if available {url_to_check}") - check_url = requests.head(url_to_check, allow_redirects=True) - if check_url.status_code == 200: - logging.info(f"Furute url found {url_to_check}") + logging.info(f"checking future url if available {url_to_check}") + check_url = _head_with_html_retry(session, + url_to_check, + timeout=30) + time.sleep(0.2) + content_type = check_url.headers.get('Content-Type', '') + if check_url.status_code == 200 and 'html' not in content_type.lower( + ): + logging.info(f"Future url found {url_to_check}") _FILES_TO_DOWNLOAD.append({"download_path": url_to_check}) + found = True break - except: - logging.error(f"URL is not accessable {url_to_check}") - - -@retry(tries=5, delay=5, backoff=5) -def download_with_retry(url, file_name_to_save): - logging.info(f"Downloaded file : {file_name_to_save} - URL {url}") - return requests.get(url=url, stream=True) + except (requests.RequestException, ValueError) as e: + logging.error(f"URL is not accessible {url_to_check}: {e}") + if not found: + fallback_url = url.format(YEAR=baseline_year) + logging.warning( + f"No future URL confirmed via HEAD probe for {url}; falling back to known baseline URL: {fallback_url}" + ) + _FILES_TO_DOWNLOAD.append({"download_path": fallback_url}) + + +@retry(tries=3, delay=5, backoff=2, exceptions=(ValueError,)) +def download_with_retry(session: requests.Session, url: str, dest_path: str): + logging.info(f"Downloading file from URL: {url} -> {dest_path}") + with session.get(url, stream=True, timeout=120) as response: + response.raise_for_status() + content_type = response.headers.get('Content-Type', '') + if 'html' in content_type.lower(): + raise ValueError(f"Server returned HTML error page for URL: {url}") + tmp_file_path = None + try: + dest_dir = os.path.dirname(dest_path) or '.' + os.makedirs(dest_dir, exist_ok=True) + with tempfile.NamedTemporaryFile(dir=dest_dir, + delete=False) as tmp_file: + tmp_file_path = tmp_file.name + for chunk in response.iter_content(chunk_size=65536): + if chunk: + tmp_file.write(chunk) + os.replace(tmp_file_path, dest_path) + tmp_file_path = None + finally: + if tmp_file_path and os.path.exists(tmp_file_path): + try: + os.remove(tmp_file_path) + except OSError: + pass + logging.info(f"Downloaded file: {dest_path}") def download_files(): @@ -1122,6 +1215,7 @@ def download_files(): global _FILES_TO_DOWNLOAD if not os.path.exists(_INPUT_FILE_PATH): os.makedirs(_INPUT_FILE_PATH) + session = _create_retry_session() try: for file_to_download in _FILES_TO_DOWNLOAD: file_name_to_save = None @@ -1132,26 +1226,26 @@ def download_files(): else: file_name_to_save = url.split('/')[-1] if 'file_path' in file_to_download: - if not os.path.exists( - os.path.join(_INPUT_FILE_PATH, - file_to_download['file_path'])): - os.makedirs( - os.path.join(_INPUT_FILE_PATH, - file_to_download['file_path'])) - file_name_to_save = file_to_download[ - 'file_path'] + file_name_to_save - - response = download_with_retry(url, file_name_to_save) - if response.status_code == 200: - with open(os.path.join(_INPUT_FILE_PATH, file_name_to_save), - 'wb') as f: - f.write(response.content) + target_dir = os.path.join(_INPUT_FILE_PATH, + file_to_download['file_path']) + if not os.path.exists(target_dir): + os.makedirs(target_dir, exist_ok=True) + file_name_to_save = os.path.join(file_to_download['file_path'], + file_name_to_save) + + dest_path = os.path.join(_INPUT_FILE_PATH, file_name_to_save) + try: + download_with_retry(session, url, dest_path) file_to_download['is_downloaded'] = True - else: - logging.error( - f"Failed to download {url} with status code {response.status_code}" - ) + except Exception as e: + logging.error(f"Failed to download {url}: {e}") file_to_download['is_downloaded'] = False + if os.path.exists(dest_path): + try: + os.remove(dest_path) + except OSError: + pass + time.sleep(0.2) failed_downloads = [ file_to_download['download_path'] for file_to_download in _FILES_TO_DOWNLOAD diff --git a/scripts/us_census/pep/annual_population/preprocess_test.py b/scripts/us_census/pep/annual_population/preprocess_test.py index 07001af4ba..21adc5994c 100644 --- a/scripts/us_census/pep/annual_population/preprocess_test.py +++ b/scripts/us_census/pep/annual_population/preprocess_test.py @@ -26,9 +26,23 @@ if path not in sys.path: sys.path.append(path) +import re +from unittest.mock import MagicMock, patch +import pandas as pd + # pylint: disable=wrong-import-position # pylint: disable=import-error -from us_census.pep.annual_population.preprocess import process +from us_census.pep.annual_population.preprocess import ( + process, + _process_nationals_2029, + _process_states_2029, + download_with_retry, + _create_retry_session, + add_future_year_urls, + popchg_file_name_pattern, + USA, +) +import us_census.pep.annual_population.preprocess as prep from us_census.pep.annual_population.constants import TEST_DATA_DIR # pylint: enable=import-error # pylint: enable=wrong-import-position @@ -109,5 +123,220 @@ def test_create_csv(self): self._actual_csv_data.strip()) +class TestProcessExcel2029(unittest.TestCase): + """Tests for dynamic future year processing in Excel files, POPCHG regex, and vintage precedence.""" + + def test_process_nationals_2029_float_int_and_str_headers(self): + """Tests that _process_nationals_2029 correctly parses float, int, and str headers for future years.""" + mock_df = pd.DataFrame({ + 'Geographic Area': [USA, 'Northeast'], + 'April 1, 2020': [331449281, 57000000], + '2020': [331511512, 57100000], + '2021': [332031554, 57200000], + '2022': [333287557, 57300000], + '2023': [334914895, 57400000], + 2024.0: [336000000, + 57500000], # parsed as float by pandas read_excel + '2025': [337000000, 57600000], # parsed as str + }) + with patch('us_census.pep.annual_population.preprocess._load_data_df', + return_value=mock_df): + result = _process_nationals_2029('dummy_path.xlsx') + + self.assertIn('Year', result.columns) + self.assertIn('Count_Person', result.columns) + self.assertIn('Location', result.columns) + years = result['Year'].tolist() + self.assertIn('2024', years) + self.assertIn('2025', years) + self.assertEqual( + result[result['Year'] == '2024']['Count_Person'].iloc[0], 336000000) + self.assertEqual( + result[result['Year'] == '2025']['Count_Person'].iloc[0], 337000000) + + def test_process_states_2029_column_length_mismatch_raises(self): + """Tests that _process_states_2029 fails fast with ValueError when extracted headers do not match data columns.""" + header_df = pd.DataFrame([['a'] * 8, ['b'] * 8, [None] * 8]) + data_df = pd.DataFrame([['.Alabama', 1, 2, 3, 4, 5, 6, 7]]) + with patch('us_census.pep.annual_population.preprocess._load_data_df', + side_effect=[header_df, data_df]): + with self.assertRaises(ValueError): + _process_states_2029('dummy_path.xlsx') + + def test_popchg_regex_and_vintage_precedence_in_process(self): + """Tests popchg_file_name_pattern and verifies newer vintage NST-EST2025-POP.xlsx takes precedence over NST-EST2023-POPCHG2020_2023.csv.""" + for year in [2023, 2024, 2025, 2026, 2029]: + filename = f"NST-EST{year}-POPCHG2020_{year}.csv" + self.assertIsNotNone(re.search(popchg_file_name_pattern, filename), + f"Failed to match {filename}") + + self.assertIsNone( + re.search(popchg_file_name_pattern, "NST-EST2024-POP.csv")) + self.assertIsNone( + re.search(popchg_file_name_pattern, "NST-EST-POPCHG2020.csv")) + + with tempfile.TemporaryDirectory() as tmp_dir: + in_dir = os.path.join(tmp_dir, 'input') + out_dir = os.path.join(tmp_dir, 'output') + os.makedirs(in_dir) + os.makedirs(out_dir) + + popchg_path = os.path.join(in_dir, + 'NST-EST2023-POPCHG2020_2023.csv') + with open(popchg_path, 'w') as f: + f.write( + 'NAME,POPESTIMATE2020,POPESTIMATE2021,POPESTIMATE2022,POPESTIMATE2023\n' + 'United States,331500000,332000000,333271411,334914895\n') + + excel_path = os.path.join(in_dir, 'NST-EST2025-POP.xlsx') + open(excel_path, 'w').close() + + nat_2025_df = pd.DataFrame({ + 'Year': ['2022', '2023', '2024', '2025'], + 'Location': ['country/USA'] * 4, + 'Count_Person': [334017321, 336806231, 340003797, 341784857], + }) + state_2025_df = pd.DataFrame({ + 'Year': ['2025'], + 'Location': ['geoId/01'], + 'Count_Person': [5157699], + }) + + csv_out = os.path.join(out_dir, 'usa_annual_population.csv') + mcf_out = os.path.join(out_dir, 'usa_annual_population.mcf') + tmcf_out = os.path.join(out_dir, 'usa_annual_population.tmcf') + + with patch.object(prep, '_process_nationals_2029', return_value=nat_2025_df), \ + patch.object(prep, '_process_states_2029', return_value=state_2025_df): + process(in_dir, csv_out, mcf_out, tmcf_out, False) + + out_df = pd.read_csv(csv_out) + usa_df = out_df[out_df['Location'] == 'country/USA'].set_index( + 'Year') + self.assertEqual(usa_df.loc[2022, 'Count_Person'], 334017321) + self.assertEqual(usa_df.loc[2023, 'Count_Person'], 336806231) + self.assertEqual(usa_df.loc[2025, 'Count_Person'], 341784857) + + +class TestDownloadAndUrlProbing(unittest.TestCase): + """Tests for download session, retries, HTML error detection, and URL probing.""" + + def test_create_retry_session(self): + """Tests that retry session is properly configured with status codes and backoff.""" + session = _create_retry_session() + self.assertIn('User-Agent', session.headers) + adapter = session.adapters.get('https://') + self.assertIsNotNone(adapter) + self.assertIsNotNone(adapter.max_retries) + self.assertIn(429, adapter.max_retries.status_forcelist) + self.assertIn(500, adapter.max_retries.status_forcelist) + self.assertIn(503, adapter.max_retries.status_forcelist) + + def test_download_with_retry_raises_on_html_error(self): + """Tests that download_with_retry raises ValueError on HTML error page (soft-404).""" + session = MagicMock() + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.headers = {'Content-Type': 'text/html; charset=UTF-8'} + mock_resp.__enter__.return_value = mock_resp + session.get.return_value = mock_resp + + with tempfile.TemporaryDirectory() as tmp_dir: + dest_file = os.path.join(tmp_dir, 'test.xlsx') + with self.assertRaises(ValueError) as ctx: + download_with_retry.__wrapped__(session, + 'https://census.gov/test.xlsx', + dest_file) + self.assertIn('HTML error page', str(ctx.exception)) + self.assertFalse(os.path.exists(dest_file)) + + def test_download_with_retry_writes_binary_file(self): + """Tests that download_with_retry writes binary data chunks atomically.""" + session = MagicMock() + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.headers = { + 'Content-Type': + 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet' + } + mock_resp.iter_content.return_value = [b'chunk1', b'chunk2'] + mock_resp.__enter__.return_value = mock_resp + session.get.return_value = mock_resp + + with tempfile.TemporaryDirectory() as tmp_dir: + dest_file = os.path.join(tmp_dir, 'test.xlsx') + download_with_retry.__wrapped__(session, + 'https://census.gov/test.xlsx', + dest_file) + self.assertTrue(os.path.exists(dest_file)) + with open(dest_file, 'rb') as f: + self.assertEqual(f.read(), b'chunk1chunk2') + + def test_add_future_year_urls_gatekeeper_break(self): + """Tests that add_future_year_urls stops scanning when a valid future year is found.""" + session = MagicMock() + head_calls = [] + + def mock_head(url, *args, **kwargs): + head_calls.append(url) + resp = MagicMock() + if '2026' in url: + resp.status_code = 200 + resp.headers = {'Content-Type': 'application/vnd.ms-excel'} + else: + resp.status_code = 404 + resp.headers = {'Content-Type': 'text/html'} + return resp + + session.head = mock_head + + with patch.object(prep, '_FILES_TO_DOWNLOAD', []), \ + patch('us_census.pep.annual_population.preprocess._create_retry_session', + return_value=session), \ + patch('time.sleep'): + add_future_year_urls() + download_urls = [ + entry['download_path'] for entry in prep._FILES_TO_DOWNLOAD + ] + + years_probed = [url for url in head_calls if '2026' in url] + self.assertTrue(len(years_probed) > 0) + years_2025 = [url for url in head_calls if '2025' in url] + self.assertEqual(len(years_2025), 0) + self.assertTrue(any('2020-2026' in url for url in download_urls)) + + def test_add_future_year_urls_fallback_to_baseline(self): + """Tests that add_future_year_urls stops at baseline_year and falls back without downgrading to older years.""" + session = MagicMock() + head_calls = [] + + def mock_head(url, *args, **kwargs): + head_calls.append(url) + resp = MagicMock() + resp.status_code = 404 + resp.headers = {'Content-Type': 'text/html'} + return resp + + session.head = mock_head + + with patch.object(prep, '_FILES_TO_DOWNLOAD', []), \ + patch('us_census.pep.annual_population.preprocess._create_retry_session', + return_value=session), \ + patch('time.sleep'): + prep.add_future_year_urls() + download_urls = [ + entry['download_path'] for entry in prep._FILES_TO_DOWNLOAD + ] + + # Ensure 2025-baseline templates never probed 2024 + self.assertFalse( + any('NST-EST2024-POP.xlsx' in url for url in head_calls)) + self.assertTrue( + any('NST-EST2025-POP.xlsx' in url for url in download_urls)) + self.assertTrue( + any('NST-EST2023-POPCHG2020_2023.csv' in url + for url in download_urls)) + + if __name__ == '__main__': unittest.main() diff --git a/scripts/us_census/pep/annual_population/validation_config.json b/scripts/us_census/pep/annual_population/validation_config.json index 4068a1c3f9..a5adce7684 100644 --- a/scripts/us_census/pep/annual_population/validation_config.json +++ b/scripts/us_census/pep/annual_population/validation_config.json @@ -24,6 +24,15 @@ "golden_files": "../../../../golden_data/golden_observations.csv", "input_files": "../../../../output_files/usa_annual_population.csv" } + }, + { + "rule_id": "check_date_freshness", + "description": "Verifies that all StatVars meet freshness requirements (within 2 years of current year).", + "validator": "SQL_VALIDATOR", + "params": { + "query": "SELECT StatVar, MaxDate FROM stats", + "condition": "MaxDate IS NOT NULL AND CAST(SUBSTR(CAST(MaxDate AS VARCHAR), 1, 4) AS INT) >= YEAR(CURRENT_DATE) - 2" + } } ] } \ No newline at end of file diff --git a/scripts/us_census/pep/us_pep_sex/manifest.json b/scripts/us_census/pep/us_pep_sex/manifest.json index 57568dcd5a..ee29b153a4 100644 --- a/scripts/us_census/pep/us_pep_sex/manifest.json +++ b/scripts/us_census/pep/us_pep_sex/manifest.json @@ -11,7 +11,8 @@ "process.py --config_path=unresolved_mcf/us_census/pep/population_estimate_by_sex/skip_url.json" ], "source_files": [ - "gcs_output/us_pep_sex_source_files/*" + "input_files/*", + "golden_data/*.csv" ], "import_inputs": [ { diff --git a/scripts/us_census/pep/us_pep_sex/process.py b/scripts/us_census/pep/us_pep_sex/process.py index a31429bbbf..6f68babb07 100644 --- a/scripts/us_census/pep/us_pep_sex/process.py +++ b/scripts/us_census/pep/us_pep_sex/process.py @@ -29,6 +29,8 @@ from absl import logging from absl import flags from retry import retry +from requests.adapters import HTTPAdapter +from urllib3.util.retry import Retry from google.cloud import storage # GCS Client import tempfile @@ -351,7 +353,7 @@ def _national_latest(file_path: str) -> pd.DataFrame: 2: 'Count_Person_Female' }}) column_list = [] - for year in range(2021, 2030): + for year in range(2021, 2031): column_name = f'POPESTIMATE{year}' if column_name in df.columns: df.rename(columns={column_name: str(year)}, inplace=True) @@ -598,7 +600,7 @@ def _state_latest(file_path: str) -> pd.DataFrame: df = pd.read_excel(file_path, skiprows=5, skipfooter=7, header=None) current_year = int(max_year) + 1 for year in range(2021, current_year): - if current_year < 2030: + if current_year <= 2031: base_columns.append(f'July{year}Total') base_columns.append(f'July{year}Male') base_columns.append(f'July{year}Female') @@ -614,7 +616,7 @@ def _state_latest(file_path: str) -> pd.DataFrame: df.insert(1, 'geo_ID', 'geoId/' + geoid) processed_dfs = [] # List to hold yearly DataFrames - for year in range(2020, 2030): + for year in range(2020, 2031): # Generate column names dynamically male_col = f'July{year}Male' female_col = f'July{year}Female' @@ -915,7 +917,8 @@ def _county_latest(file_path: str) -> pd.DataFrame: 8: '2026', 9: '2027', 10: '2028', - 11: '2029' + 11: '2029', + 12: '2030' } }) @@ -1016,6 +1019,60 @@ def process(self): processed_count = 0 total_files_to_process = len(ip_files) logging.info(f"No of files to be processed {len(ip_files)}") + # Define the base mappings for fixed years (e.g., 2023) + file_to_function_mapping = { + "pe-11-1": _national_1900_1979, + "us-est90int": _national_1990_2000, + "us-est90int-": _national_1990_2000, + "us-est00int": _national_2000_2010, + "us-est00int-": _national_2000_2010, + "nc-est2020-agesex-": _national_2010_2020, + "nc-est2020-agesex-r": _national_2010_2020, + "pe": _state_1970_1980, + "stiag": _state_1980_1990, + "st-est00int-02-": _state_2000_2010, + "st-est00int-02": _state_2000_2010, + "SC-EST2020-AGESEX": _state_2010_2020, + "SC-EST2020-AGESEX-": _state_2010_2020, + "co-asr-1": _county_1970_1980, + "pe-02-1": _county_1980_1990, + "stch-icen1": _county_1990_2000, + "co-est00int-agesex-": _county_2000_2010, + "co-est00int-agesex-5": _county_2000_2010, + "CC-EST2020-AGESEX-": _county_2010_2020, + "CC-EST2020-AGESEX-A": _county_2010_2020 + } + + # Iterate from 2023 to 2030 and add mappings for years dynamically + for file_year in range(2023, 2031): + # For the 'national', 'state', and 'county' entries, use the same method for years 2023 to 2030 + file_to_function_mapping[ + f"nc-est{file_year}-agesex-"] = _national_latest + file_to_function_mapping[ + f"nc-est{file_year}-agesex-r"] = _national_latest + + file_to_function_mapping[ + f"sc-est{file_year}-syasex-2"] = _state_latest + file_to_function_mapping[ + f"sc-est{file_year}-syasex-3"] = _state_latest + file_to_function_mapping[ + f"sc-est{file_year}-syasex-4"] = _state_latest + file_to_function_mapping[ + f"sc-est{file_year}-syasex-5"] = _state_latest + file_to_function_mapping[ + f"sc-est{file_year}-syasex-0"] = _state_latest + file_to_function_mapping[ + f"sc-est{file_year}-syasex-1"] = _state_latest + file_to_function_mapping[ + f"sc-est{file_year}-syasex-"] = _state_latest + file_to_function_mapping[ + f"sc-est{file_year}-agesex-"] = _state_latest + + file_to_function_mapping[ + f"cc-est{file_year}-agesex-"] = _county_latest + file_to_function_mapping[ + f"cc-est{file_year}-agesex-a"] = _county_latest + for file_path in ip_files: logging.info(f"Processing the file:{file_path}") processed_count += 1 @@ -1025,61 +1082,6 @@ def process(self): file_name = file_path.split("/")[-1][:-7] - # Define the base mappings for fixed years (e.g., 2023) - file_to_function_mapping = { - "pe-11-1": _national_1900_1979, - "us-est90int": _national_1990_2000, - "us-est90int-": _national_1990_2000, - "us-est00int": _national_2000_2010, - "us-est00int-": _national_2000_2010, - "nc-est2020-agesex-": _national_2010_2020, - "nc-est2020-agesex-r": _national_2010_2020, - "pe": _state_1970_1980, - "stiag": _state_1980_1990, - "st-est00int-02-": _state_2000_2010, - "st-est00int-02": _state_2000_2010, - "SC-EST2020-AGESEX": _state_2010_2020, - "SC-EST2020-AGESEX-": _state_2010_2020, - "co-asr-1": _county_1970_1980, - "pe-02-1": _county_1980_1990, - "stch-icen1": _county_1990_2000, - "co-est00int-agesex-": _county_2000_2010, - "co-est00int-agesex-5": _county_2000_2010, - "CC-EST2020-AGESEX-": _county_2010_2020, - "CC-EST2020-AGESEX-A": _county_2010_2020 - } - - # Iterate from 2023 to 2029 and add mappings for years dynamically - for file_year in range(2023, 2030): - # For the 'national', 'state', and 'county' entries, use the same method for years 2023 to 2029 - if file_year < dt.now().year: - file_to_function_mapping[ - f"nc-est{file_year}-agesex-"] = _national_latest - file_to_function_mapping[ - f"nc-est{file_year}-agesex-r"] = _national_latest - - file_to_function_mapping[ - f"sc-est{file_year}-syasex-2"] = _state_latest - file_to_function_mapping[ - f"sc-est{file_year}-syasex-3"] = _state_latest - file_to_function_mapping[ - f"sc-est{file_year}-syasex-4"] = _state_latest - file_to_function_mapping[ - f"sc-est{file_year}-syasex-5"] = _state_latest - file_to_function_mapping[ - f"sc-est{file_year}-syasex-0"] = _state_latest - file_to_function_mapping[ - f"sc-est{file_year}-syasex-1"] = _state_latest - file_to_function_mapping[ - f"sc-est{file_year}-syasex-"] = _state_latest - file_to_function_mapping[ - f"sc-est{file_year}-agesex-"] = _state_latest - - file_to_function_mapping[ - f"cc-est{file_year}-agesex-"] = _county_latest - file_to_function_mapping[ - f"cc-est{file_year}-agesex-a"] = _county_latest - df = file_to_function_mapping[file_name](file_path) if not df.empty: @@ -1092,7 +1094,7 @@ def process(self): logging.info(f"No of files processed {processed_count}") # Log the resulting dictionary logging.info(f"File-to-Function mappings: {file_to_function_mapping}") - if processed_count == total_files_to_process & total_files_to_process > 0: + if processed_count == total_files_to_process and total_files_to_process > 0: final_df = _states_full_to_short_form(final_df, 'geo_ID', "geo_ID") final_df = _add_geo_id(final_df, "geo_ID", "geo_ID") final_df = pd.melt( @@ -1130,23 +1132,66 @@ def fetch_skip_urls_from_gcs(GCS_BUCKET_NAME: str, return skip_urls -def is_valid_url(url): +def _create_retry_session() -> requests.Session: + """Creates a requests Session configured with retry and backoff for transient errors.""" + session = requests.Session() + session.headers.update({'User-Agent': 'Mozilla/5.0'}) + retries = Retry( + total=5, + backoff_factor=2, + status_forcelist=[408, 429, 500, 502, 503, 504, 520], + allowed_methods=["HEAD", "GET"], + respect_retry_after_header=True, + ) + adapter = HTTPAdapter(max_retries=retries) + session.mount("https://", adapter) + session.mount("http://", adapter) + return session + + +@retry(tries=3, delay=2, backoff=2, exceptions=(ValueError,)) +def _head_with_html_retry(session: requests.Session, + url: str, + timeout: int = 15) -> requests.Response: + """Performs a HEAD request and retries if HTTP 200 returns an HTML error page.""" + response = session.head(url, allow_redirects=True, timeout=timeout) + content_type = response.headers.get('Content-Type', '') + if response.status_code == 200 and 'html' in content_type.lower(): + raise ValueError(f"Server returned HTTP 200 HTML response for {url}") + return response + + +def is_valid_url(session: requests.Session, url: str) -> bool: try: - response = requests.get(url, timeout=20) - if response.status_code != 200: - return False - content_type = response.headers.get("Content-Type", "") - if "text/html" in content_type: - # Might be an error page disguised as 200 - if b"error" in response.content.lower() or len( - response.content) < 500: + with session.get(url, stream=True, timeout=20) as response: + if response.status_code != 200: + return False + content_type = response.headers.get("Content-Type", "") + if "html" in content_type.lower(): return False - return True + return True except Exception as e: logging.warning(f"Error checking URL: {url} - {e}") return False +_BASELINE_YEAR = 2025 + + +def _retain_cached_if_exists(url_to_check: str): + """If a cached file exists in persistent storage for url_to_check, retain it in _FILES_TO_DOWNLOAD.""" + file_name = url_to_check.split('/')[-1] + cached_file_path = os.path.join(_GCS_FOLDER_PERSISTENT_PATH, file_name) + if os.path.isfile(cached_file_path): + logging.warning( + f"HEAD probe failed for {url_to_check}, but cached copy exists at {cached_file_path}; retaining in download list." + ) + _FILES_TO_DOWNLOAD.append({ + "download_path": url_to_check, + "is_historical": False + }) + + def add_future_year_urls(): """ This method adds the future year urls that has to be downloaded @@ -1158,11 +1203,10 @@ def add_future_year_urls(): # Initialize the list to store files to download _FILES_TO_DOWNLOAD = [] - # Use a requests session for connection pooling to improve efficiency - # when making hundreds of HEAD requests. - session = requests.Session() - with open(os.path.join(_MODULE_DIR, 'input_url.json'), 'r') as inpit_file: - _FILES_TO_DOWNLOAD = json.load(inpit_file) + # Use a retry-enabled requests session for connection pooling and 429/5xx backoff + session = _create_retry_session() + with open(os.path.join(_MODULE_DIR, 'input_url.json'), 'r') as input_file: + _FILES_TO_DOWNLOAD = json.load(input_file) GCS_BUCKET_NAME, GCS_SKIP_FILE_PATH = extract_gcs_info() # Read URLs from the GCS-hosted config file @@ -1171,19 +1215,21 @@ def add_future_year_urls(): skip_urls = fetch_skip_urls_from_gcs(GCS_BUCKET_NAME, GCS_SKIP_FILE_PATH) logging.info(f"Fetched {len(skip_urls)} URLs to skip from GCS.") - logging.info(f"Urls fetched from the json to skip:{skip_urls}") + logging.info(f"Urls fetched from the json to skip: {skip_urls}") #Filter based on skip list + live URL check _FILES_TO_DOWNLOAD = [ url for url in _FILES_TO_DOWNLOAD if (url["download_path"] not in skip_urls # not suspicious, keep it or (url["download_path"] in skip_urls and is_valid_url( - url["download_path"])) # suspicious, but passes check + session, url["download_path"])) # suspicious, but passes check ) ] + for item in _FILES_TO_DOWNLOAD: + item["is_historical"] = True logging.info( - f"Historical urls Fetched from input_json:{_FILES_TO_DOWNLOAD}") + f"Historical urls fetched from input_json: {_FILES_TO_DOWNLOAD}") # List of URLs with placeholders for {YEAR} and {i} urls_to_scan = [ @@ -1195,173 +1241,236 @@ def add_future_year_urls(): # A set to track downloaded URLs for unique {YEAR} and URLs without {i} downloaded_year_urls = set() + selected_year = None - # Loop through years in reverse order from 2030 to 2023 - for future_year in range(2030, 2022, -1): # From 2030 to 2023 - + # Loop through years in reverse order from 2030 down to _BASELINE_YEAR + for future_year in range(2030, _BASELINE_YEAR - 1, -1): # We check the National CSV first. If it's 404, the whole year is skipped. gatekeeper_url = urls_to_scan[0].format(YEAR=future_year) try: - # Use a short 5-second timeout for the check - response = session.head(gatekeeper_url, - allow_redirects=True, - timeout=5) - if response.status_code != 200: + response = _head_with_html_retry(session, + gatekeeper_url, + timeout=15) + time.sleep(0.2) + content_type = response.headers.get('Content-Type', '') + if response.status_code != 200 or 'html' in content_type.lower(): logging.info( - f"Skipping year {future_year}: National file not found (status code: {response.status_code})." + f"Skipping year {future_year}: National file not found or is HTML (status code: {response.status_code})." ) continue - except requests.exceptions.RequestException as e: + selected_year = future_year + break + except (requests.exceptions.RequestException, ValueError) as e: logging.warning( f"Skipping year {future_year} due to an error checking the gatekeeper URL: {e}" ) continue - YEAR = future_year - # Loop through URLs - for url in urls_to_scan: - if "{i}" in url: # This URL contains the {i} variable, so we loop through i from 01 to 56 - for i in range(1, 57): # Loop i from 01 to 56 - formatted_i = f"{i:02}" # Ensure i is always 2 digits (01, 02, ..., 56) - url_to_check = url.format(YEAR=YEAR, i=formatted_i) - logging.info(f"checking url: {url_to_check}") - - try: - # HEAD calls only fetch headers and should complete quickly. - check_url = session.head(url_to_check, - allow_redirects=True, - timeout=5) - if check_url.status_code == 200: - _FILES_TO_DOWNLOAD.append( - {"download_path": url_to_check}) - - except requests.exceptions.RequestException as e: - logging.error( - f"URL is not accessible {url_to_check} due to {e}") - - else: # This URL does not contain {i}, so we only need to process it once per year - url_to_check = url.format(YEAR=YEAR) + if selected_year is None: + logging.warning( + f"No gatekeeper URL confirmed via HEAD probe; falling back to baseline year {_BASELINE_YEAR}." + ) + selected_year = _BASELINE_YEAR + + YEAR = selected_year + # Loop through URLs for the selected vintage year + for url in urls_to_scan: + if "{i}" in url: # This URL contains the {i} variable, so we loop through i from 01 to 56 + for i in range(1, 57): # Loop i from 01 to 56 + formatted_i = f"{i:02}" # Ensure i is always 2 digits (01, 02, ..., 56) + url_to_check = url.format(YEAR=YEAR, i=formatted_i) logging.info(f"checking url: {url_to_check}") - # If the URL has already been processed for this year, skip it - if url_to_check in downloaded_year_urls: - continue # Skip this URL if it's already processed try: # HEAD calls only fetch headers and should complete quickly. - check_url = session.head(url_to_check, - allow_redirects=True, - timeout=5) - if check_url.status_code == 200: - _FILES_TO_DOWNLOAD.append( - {"download_path": url_to_check}) - downloaded_year_urls.add( - url_to_check) # Mark this URL as processed - + check_url = _head_with_html_retry(session, + url_to_check, + timeout=15) + time.sleep(0.1) + content_type = check_url.headers.get('Content-Type', '') + if check_url.status_code == 200 and 'html' not in content_type.lower( + ): + _FILES_TO_DOWNLOAD.append({ + "download_path": url_to_check, + "is_historical": False + }) else: - logging.error( - f"URL returned status code {check_url.status_code}: {url_to_check}" - ) + _retain_cached_if_exists(url_to_check) - except requests.exceptions.RequestException as e: + except (requests.exceptions.RequestException, ValueError) as e: logging.error( f"URL is not accessible {url_to_check} due to {e}") + _retain_cached_if_exists(url_to_check) + + else: # This URL does not contain {i}, so we only need to process it once per year + url_to_check = url.format(YEAR=YEAR) + logging.info(f"checking url: {url_to_check}") + # If the URL has already been processed for this year, skip it + if url_to_check in downloaded_year_urls: + continue # Skip this URL if it's already processed + + try: + # HEAD calls only fetch headers and should complete quickly. + check_url = _head_with_html_retry(session, + url_to_check, + timeout=15) + time.sleep(0.1) + content_type = check_url.headers.get('Content-Type', '') + if check_url.status_code == 200 and 'html' not in content_type.lower( + ): + _FILES_TO_DOWNLOAD.append({ + "download_path": url_to_check, + "is_historical": False + }) + downloaded_year_urls.add( + url_to_check) # Mark this URL as processed + + else: + logging.error( + f"URL returned status code {check_url.status_code}: {url_to_check}" + ) + _retain_cached_if_exists(url_to_check) + + except (requests.exceptions.RequestException, ValueError) as e: + logging.error( + f"URL is not accessible {url_to_check} due to {e}") + _retain_cached_if_exists(url_to_check) + + +def _get_target_file_name(file_to_download: dict) -> str: + """Returns the local filename to save for a download entry.""" + if 'file_name' in file_to_download and len( + file_to_download['file_name']) > 5: + return file_to_download['file_name'] + return file_to_download['download_path'].split('/')[-1] def cleanup(): - """Delete all old files in the gcs_folder to prevent cache bloat.""" - if os.path.exists(_GCS_FOLDER_PERSISTENT_PATH): - for file_name in os.listdir(_GCS_FOLDER_PERSISTENT_PATH): - file_path = os.path.join(_GCS_FOLDER_PERSISTENT_PATH, file_name) - if os.path.isfile(file_path): - file_age = (time.time() - os.path.getmtime(file_path)) / (24 * - 3600) - # Delete ANY file older than the TTL - if file_age > _TTL_DAYS: - logging.info(f"Cleaning up old file: {file_name}") - os.remove(file_path) - - -@retry(tries=3, - delay=2, - backoff=2, - exceptions=(requests.RequestException, Exception)) + """Delete unreferenced or expired files in gcs_folder after downloads complete.""" + if not os.path.exists(_GCS_FOLDER_PERSISTENT_PATH): + return + active_files = set() + if _FILES_TO_DOWNLOAD: + active_files = { + _get_target_file_name(item) for item in _FILES_TO_DOWNLOAD + } + for file_name in os.listdir(_GCS_FOLDER_PERSISTENT_PATH): + file_path = os.path.join(_GCS_FOLDER_PERSISTENT_PATH, file_name) + if os.path.isfile(file_path): + file_age = (time.time() - os.path.getmtime(file_path)) / (24 * 3600) + if file_age > _TTL_DAYS and file_name not in active_files: + logging.info(f"Cleaning up old file: {file_name}") + os.remove(file_path) + + +@retry(tries=3, delay=3, backoff=2, exceptions=(ValueError,)) +def _download_single_file(session: requests.Session, url: str) -> str: + """Downloads a single URL to a temporary file, retrying on HTML error responses.""" + with session.get(url, stream=True, timeout=120) as response: + response.raise_for_status() + content_type = response.headers.get('Content-Type', '') + if 'html' in content_type.lower(): + raise ValueError(f"Server returned HTML error page for URL: {url}") + + tmp_file_path = None + try: + with tempfile.NamedTemporaryFile(delete=False) as tmp_file: + tmp_file_path = tmp_file.name + for chunk in response.iter_content(chunk_size=8192): + if chunk: + tmp_file.write(chunk) + return tmp_file_path + except Exception: + if tmp_file_path and os.path.exists(tmp_file_path): + try: + os.remove(tmp_file_path) + except OSError: + pass + raise + + def download_files(): """ Download files from URLs listed in _FILES_TO_DOWNLOAD. Skips URLs listed in skip_url.json from GCS. """ global _FILES_TO_DOWNLOAD - session = requests.session() + session = _create_retry_session() # Ensure the directory exists (it shouldn't "expect" it to be there) os.makedirs(_GCS_FOLDER_PERSISTENT_PATH, exist_ok=True) + os.makedirs(_INPUT_FILE_PATH, exist_ok=True) # Get set of already downloaded files downloaded_files = set(os.listdir(_GCS_FOLDER_PERSISTENT_PATH)) for file_to_download in _FILES_TO_DOWNLOAD: - file_name_to_save = None url = file_to_download['download_path'] - - if 'file_name' in file_to_download and len( - file_to_download['file_name']) > 5: - file_name_to_save = file_to_download['file_name'] - else: - file_name_to_save = url.split('/')[-1] - - # Skip if file already exists in cache (cleanup() already removed stale files) - if file_name_to_save in downloaded_files: - file_path = os.path.join(_GCS_FOLDER_PERSISTENT_PATH, - file_name_to_save) + file_name_to_save = _get_target_file_name(file_to_download) + cached_file_path = os.path.join(_GCS_FOLDER_PERSISTENT_PATH, + file_name_to_save) + is_historical = file_to_download.get('is_historical', False) + + # Reuse cached file if it exists (historical files never expire; dynamic files use _TTL_DAYS) + if file_name_to_save in downloaded_files and os.path.isfile( + cached_file_path): + if is_historical: + logging.info( + f"Skipping download, using immutable historical cached file: {file_name_to_save}" + ) + shutil.copy(cached_file_path, + os.path.join(_INPUT_FILE_PATH, file_name_to_save)) + file_to_download['is_downloaded'] = True + continue + file_age = (time.time() - + os.path.getmtime(cached_file_path)) / (24 * 3600) + if file_age <= _TTL_DAYS: + logging.info( + f"Skipping download, using cached file: {file_name_to_save}" + ) + shutil.copy(cached_file_path, + os.path.join(_INPUT_FILE_PATH, file_name_to_save)) + file_to_download['is_downloaded'] = True + continue logging.info( - f"Skipping download, using cached file: {file_name_to_save}") - - # Make sure to copy the cached file to the input directory! - shutil.copy(file_path, - os.path.join(_INPUT_FILE_PATH, file_name_to_save)) - continue + f"Cached file {file_name_to_save} expired (age: {file_age:.1f} days); refreshing." + ) - headers = {'User-Agent': 'Mozilla/5.0'} + tmp_file_path = None try: - with session.get(url, stream=True, timeout=120, - headers=headers) as response: - response.raise_for_status() - - content_type = response.headers.get('Content-Type', '') - - # Minimal fix: Log error and continue to skip HTML pages - if 'html' in content_type.lower(): - logging.error( - f"Server returned HTML error page for URL: {url}. Skipping." - ) - continue - - if response.status_code == 200: - with tempfile.NamedTemporaryFile(delete=False) as tmp_file: - for chunk in response.iter_content(chunk_size=8192): - if chunk: - tmp_file.write(chunk) - tmp_file_path = tmp_file.name + tmp_file_path = _download_single_file(session, url) - # Copy to local destination - shutil.copy( - tmp_file_path, + # Copy to local destination + shutil.copy(tmp_file_path, os.path.join(_INPUT_FILE_PATH, file_name_to_save)) - # Move to gcs destination (optimized from shutil.copy + os.remove) - shutil.move( - tmp_file_path, - os.path.join(_GCS_FOLDER_PERSISTENT_PATH, - file_name_to_save)) + # Move to gcs destination (optimized from shutil.copy + os.remove) + shutil.move(tmp_file_path, cached_file_path) + tmp_file_path = None + downloaded_files.add(file_name_to_save) - file_to_download['is_downloaded'] = True - logging.info(f"Downloaded file: {url}") + file_to_download['is_downloaded'] = True + logging.info(f"Downloaded file: {url}") except Exception as e: + if tmp_file_path and os.path.exists(tmp_file_path): + try: + os.remove(tmp_file_path) + except OSError: + pass + if os.path.isfile(cached_file_path): + logging.warning( + f"Failed to refresh {url} after retries ({e}); falling back to cached file: {file_name_to_save}" + ) + os.utime(cached_file_path, None) + shutil.copy(cached_file_path, + os.path.join(_INPUT_FILE_PATH, file_name_to_save)) + file_to_download['is_downloaded'] = True + continue file_to_download['is_downloaded'] = False logging.error(f"Error downloading {url}: {e}") raise - time.sleep(1) + time.sleep(0.5) return True @@ -1391,9 +1500,6 @@ def main(_): mcf_path = data_file_path + os.sep + mcf_name tmcf_path = data_file_path + os.sep + tmcf_name - # Perform cleanup of old files - cleanup() - download_status = True if mode == "" or mode == "download": # Get the config path from the flags @@ -1406,6 +1512,7 @@ def main(_): # download & process add_future_year_urls() download_status = download_files() + cleanup() if download_status and (mode == "" or mode == "process"): try: loader = PopulationEstimateBySex(_INPUT_FILE_PATH, cleaned_csv_path, diff --git a/scripts/us_census/pep/us_pep_sex/process_test.py b/scripts/us_census/pep/us_pep_sex/process_test.py index 26013e7fb0..65460c6689 100644 --- a/scripts/us_census/pep/us_pep_sex/process_test.py +++ b/scripts/us_census/pep/us_pep_sex/process_test.py @@ -19,6 +19,9 @@ import unittest import sys import tempfile +from unittest.mock import MagicMock, patch +import numpy as np +import pandas as pd from absl import flags # module_dir is the path to where this test is running from. @@ -29,7 +32,13 @@ sys.path.append(path) # pylint: disable=wrong-import-position -from us_census.pep.us_pep_sex.process import PopulationEstimateBySex +from us_census.pep.us_pep_sex.process import ( + PopulationEstimateBySex, + _state_latest, + _download_single_file, + is_valid_url, +) +import us_census.pep.us_pep_sex.process as sex_proc # pylint: enable=wrong-import-position TEST_DATASET_DIR = os.path.join(MODULE_DIR, "test_data", "datasets") @@ -111,5 +120,236 @@ def test_create_csv(self): self.actual_csv_data.strip()) +class TestStateLatestVintageBoundary(unittest.TestCase): + """Tests vintage boundary logic in _state_latest.""" + + def test_state_latest_boundary_2029(self): + """Verifies that max_year 2029 (current_year=2030) appends columns up to 2029 without being dropped.""" + header_df = pd.DataFrame( + [['dummy'] * 34, ['dummy'] * 34, + [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan] + + [2021, 2022, 2023, 2024, 2025, 2026, 2027, 2028, 2029] + + [np.nan] * 18]) + data_row = ['Total'] + [1000] * 33 + data_df = pd.DataFrame([data_row]) + + with patch('pandas.read_excel', side_effect=[header_df, data_df]): + res = _state_latest('dummy_sc-est2029-syasex-01.xlsx') + + self.assertIn('Year', res.columns) + years = res['Year'].unique().tolist() + self.assertIn('2029', years) + self.assertIn('geo_ID', res.columns) + self.assertEqual(res['geo_ID'].iloc[0], 'geoId/01') + + +class TestDynamicFunctionMapping(unittest.TestCase): + """Tests dynamic registration of file-to-function mappings up to year 2030.""" + + def test_file_to_function_mapping_years_2023_to_2030(self): + """Tests that national, state, and county future year files (e.g. 2024, 2029, 2030) are dispatched correctly.""" + with tempfile.TemporaryDirectory() as tmp_dir: + for yr in [2024, 2029, 2030]: + open(os.path.join(tmp_dir, f"nc-est{yr}-agesex-res.csv"), + 'w').close() + open(os.path.join(tmp_dir, f"sc-est{yr}-syasex-01.xlsx"), + 'w').close() + open(os.path.join(tmp_dir, f"cc-est{yr}-agesex-all.csv"), + 'w').close() + + mock_df = pd.DataFrame({ + 'Year': ['2024'], + 'geo_ID': ['geoId/01'], + 'Count_Person_Male': [10], + 'Count_Person_Female': [10], + 'Measurement_Method': ['CensusPEPSurvey'] + }) + + with patch('us_census.pep.us_pep_sex.process._national_latest', + return_value=mock_df) as mock_nat, \ + patch('us_census.pep.us_pep_sex.process._state_latest', + return_value=mock_df) as mock_st, \ + patch('us_census.pep.us_pep_sex.process._county_latest', + return_value=mock_df) as mock_co: + loader = PopulationEstimateBySex( + tmp_dir, os.path.join(tmp_dir, 'out.csv'), + os.path.join(tmp_dir, 'out.mcf'), + os.path.join(tmp_dir, 'out.tmcf')) + loader.process() + self.assertEqual(mock_nat.call_count, 3) + self.assertEqual(mock_st.call_count, 3) + self.assertEqual(mock_co.call_count, 3) + + +class TestDownloadAndCacheLifecycle(unittest.TestCase): + """Tests for streaming downloads, temp file cleanup, and cache fallback.""" + + def test_is_valid_url_rejects_html_regardless_of_size(self): + """Tests that is_valid_url returns False immediately for any HTML content-type.""" + session = MagicMock() + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.headers = {'Content-Type': 'text/html; charset=UTF-8'} + mock_resp.iter_content.return_value = [b'' + b'a' * 1000] + mock_resp.__enter__.return_value = mock_resp + session.get.return_value = mock_resp + + self.assertFalse(is_valid_url(session, 'https://census.gov/test.csv')) + + def test_download_single_file_success(self): + """Tests that _download_single_file streams chunks to temp file.""" + session = MagicMock() + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.headers = {'Content-Type': 'application/vnd.ms-excel'} + mock_resp.iter_content.return_value = [b'data1', b'data2'] + mock_resp.__enter__.return_value = mock_resp + session.get.return_value = mock_resp + + tmp_path = _download_single_file.__wrapped__( + session, 'https://census.gov/test.xlsx') + self.assertTrue(os.path.exists(tmp_path)) + with open(tmp_path, 'rb') as f: + self.assertEqual(f.read(), b'data1data2') + os.remove(tmp_path) + + def test_download_single_file_html_error(self): + """Tests that _download_single_file raises ValueError on HTML error response.""" + session = MagicMock() + mock_resp = MagicMock() + mock_resp.status_code = 200 + mock_resp.headers = {'Content-Type': 'text/html; charset=UTF-8'} + mock_resp.__enter__.return_value = mock_resp + session.get.return_value = mock_resp + + with self.assertRaises(ValueError) as ctx: + _download_single_file.__wrapped__(session, + 'https://census.gov/test.xlsx') + self.assertIn('HTML error page', str(ctx.exception)) + + def test_download_single_file_cleanup_on_exception(self): + """Tests that temporary files are deleted if an error occurs mid-stream.""" + session = MagicMock() + mock_resp = MagicMock() + mock_resp.headers = {'Content-Type': 'application/octet-stream'} + mock_resp.__enter__.return_value = mock_resp + + created_files = [] + original_named_temp = tempfile.NamedTemporaryFile + + def track_temp(*args, **kwargs): + tf = original_named_temp(*args, **kwargs) + created_files.append(tf.name) + return tf + + def failing_iter(*args, **kwargs): + yield b'chunk1' + raise IOError('Connection aborted') + + mock_resp.iter_content = failing_iter + session.get.return_value = mock_resp + + with patch('tempfile.NamedTemporaryFile', side_effect=track_temp): + with self.assertRaises(IOError): + _download_single_file.__wrapped__(session, + 'https://census.gov/broken') + + self.assertEqual(len(created_files), 1) + self.assertFalse(os.path.exists(created_files[0])) + + def test_download_files_cache_fallback_and_historical_ttl_exemption(self): + """Tests that historical files bypass TTL and expired dynamic files update mtime on fallback.""" + with tempfile.TemporaryDirectory() as tmp_dir: + gcs_persistent = os.path.join(tmp_dir, 'gcs_persistent') + input_dir = os.path.join(tmp_dir, 'input_files') + os.makedirs(gcs_persistent) + os.makedirs(input_dir) + + hist_cached = os.path.join(gcs_persistent, 'pe-19.csv') + with open(hist_cached, 'w') as f: + f.write('historical_content') + os.utime(hist_cached, (0, 0)) + + dyn_cached = os.path.join(gcs_persistent, + 'nc-est2025-agesex-res.csv') + with open(dyn_cached, 'w') as f: + f.write('cached_content') + os.utime(dyn_cached, (0, 0)) + + files_to_download = [ + { + 'download_path': 'https://census.gov/pe-19.csv', + 'is_historical': True, + }, + { + 'download_path': + 'https://census.gov/nc-est2025-agesex-res.csv', + 'is_historical': + False, + }, + ] + + with patch.object(sex_proc, '_GCS_FOLDER_PERSISTENT_PATH', gcs_persistent), \ + patch.object(sex_proc, '_INPUT_FILE_PATH', input_dir), \ + patch.object(sex_proc, '_FILES_TO_DOWNLOAD', files_to_download), \ + patch.object(sex_proc, '_download_single_file', side_effect=IOError('Network error')) as mock_dl, \ + patch('time.sleep'): + sex_proc.download_files() + # Only the expired non-historical file should have attempted a network download + self.assertEqual(mock_dl.call_count, 1) + + dest_hist = os.path.join(input_dir, 'pe-19.csv') + self.assertTrue(os.path.exists(dest_hist)) + dest_file = os.path.join(input_dir, 'nc-est2025-agesex-res.csv') + self.assertTrue(os.path.exists(dest_file)) + with open(dest_file) as f: + self.assertEqual(f.read(), 'cached_content') + self.assertTrue(files_to_download[0]['is_downloaded']) + self.assertTrue(files_to_download[1]['is_downloaded']) + # Verify mtime of fallback cached file was refreshed + self.assertGreater(os.path.getmtime(dyn_cached), 0) + + def test_add_future_year_urls_gatekeeper_break_and_cache_retention(self): + """Tests that add_future_year_urls breaks at first valid year, bounds at baseline, and retains cached files on HEAD failure.""" + session = MagicMock() + head_calls = [] + + def mock_head(url, *args, **kwargs): + head_calls.append(url) + resp = MagicMock() + if '2026' in url and 'syasex-01' not in url: + resp.status_code = 200 + resp.headers = {'Content-Type': 'text/csv'} + else: + resp.status_code = 404 + resp.headers = {'Content-Type': 'text/html'} + return resp + + session.head = mock_head + + with tempfile.TemporaryDirectory() as tmp_dir: + # Create a cached file for sc-est2026-syasex-01.xlsx whose HEAD probe will fail + cached_state = os.path.join(tmp_dir, 'sc-est2026-syasex-01.xlsx') + with open(cached_state, 'w') as f: + f.write('cached') + + with patch.object(sex_proc, '_FILES_TO_DOWNLOAD', []), \ + patch.object(sex_proc, '_GCS_FOLDER_PERSISTENT_PATH', tmp_dir), \ + patch.object(sex_proc, 'extract_gcs_info', return_value=('bucket', 'path')), \ + patch.object(sex_proc, 'fetch_skip_urls_from_gcs', return_value=set()), \ + patch.object(sex_proc, '_create_retry_session', return_value=session), \ + patch('time.sleep'): + sex_proc.add_future_year_urls() + downloaded_paths = [ + item['download_path'] + for item in sex_proc._FILES_TO_DOWNLOAD + ] + + self.assertTrue(any('2026' in url for url in head_calls)) + self.assertFalse(any('2025' in url for url in head_calls)) + self.assertTrue( + any('sc-est2026-syasex-01.xlsx' in url for url in downloaded_paths)) + + if __name__ == "__main__": unittest.main() diff --git a/scripts/us_census/pep/us_pep_sex/validation_config.json b/scripts/us_census/pep/us_pep_sex/validation_config.json index c4bde5edcc..56fa700e5f 100644 --- a/scripts/us_census/pep/us_pep_sex/validation_config.json +++ b/scripts/us_census/pep/us_pep_sex/validation_config.json @@ -24,6 +24,15 @@ "golden_files": "../../../../golden_data/golden_observations.csv", "input_files": "../../../../output/population_estimate_sex.csv" } + }, + { + "rule_id": "check_date_freshness", + "description": "Verifies that all StatVars meet freshness requirements (within 2 years of current year).", + "validator": "SQL_VALIDATOR", + "params": { + "query": "SELECT StatVar, MaxDate FROM stats", + "condition": "MaxDate IS NOT NULL AND CAST(SUBSTR(CAST(MaxDate AS VARCHAR), 1, 4) AS INT) >= YEAR(CURRENT_DATE) - 2" + } } ] }