from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from config import Config
from airflow.models import BaseOperator
from airflow.exceptions import AirflowSkipException
from config.fips import FipsCodes
[docs]
class CensusDownloaderOperator(BaseOperator):
ui_color = "#A6E6A6"
DATA_PROD_VERSION = "v1"
DATA_PROD_SUFFIX = "r_census"
FTP_SERVER = "ftp.census.gov"
HTTP_SERVER = "https://www2.census.gov"
[docs]
@staticmethod
def get_s3_prefix(vintage: str, geolevel: str) -> str:
s3_path = f"s3://{Config.DATA_LAKE_BUCKET}/{Config.RAW_PREFIX}/census_tiger/{CensusDownloaderOperator.get_table_name(vintage,geolevel)}/vintage={vintage}/"
print(f"s3_path: {s3_path}")
return s3_path
[docs]
@staticmethod
def get_table_name(vintage: str, geolevel: str) -> str:
table_name = f"{CensusDownloaderOperator.DATA_PROD_SUFFIX}_geo_{geolevel}_{vintage%100}_{CensusDownloaderOperator.DATA_PROD_VERSION}"
print(f"table_name: {table_name}")
return table_name
@apply_defaults
def __init__(
self,
vintage: str,
state_abbr: str,
geolevel: str,
partition_cols: list[str],
datalake_db: str = Config.DATA_LAKE_DB,
use_ftp: bool = True,
*args,
**kwargs,
):
"""
An operator to download census shape files from the Census FTP server and store them in S3 and Athena.
:param vintage: the census vintage, e.g. 2023
:type vintage: str
:param state_abbr: the state abbreviation, e.g. 'AR' for Arkansas
:type state_abbr: str
:param geolevel: the geographic level, e.g. 'county', 'state', 'city', etc.
:type geolevel: str
:param partition_cols: the partition columns for the Athena table, e.g. ['statefp']
:type partition_cols: list[str]
:param datalake_db: the datalake database, defaults to Config.DATA_LAKE_DB
:type datalake_db: str, optional
"""
self.datalake_db = datalake_db
self.geolevel = geolevel
self.state_abbr = state_abbr
self.partition_cols = partition_cols
self.vintage = vintage
self.statefip = FipsCodes.ALL.get(state_abbr, "")
self.output_athena_dataset_name = CensusDownloaderOperator.get_table_name(
vintage, geolevel
)
self.s3_root = CensusDownloaderOperator.get_s3_prefix(vintage, geolevel)
self.map_name = self.output_athena_dataset_name
self.use_ftp = use_ftp
super().__init__(*args, **kwargs)
[docs]
def process_zip_file(self, downloaded_file, current_file, update_table=False):
import awswrangler as wr
import geopandas as gpd
import pandas as pd
if self.state_abbr:
partition_cols = ["statefp"]
else:
partition_cols = []
try:
df = gpd.read_file(downloaded_file)
if df.empty:
raise AirflowSkipException(
f"Empty dataframe for {self.geolevel} and {self.statefip}"
)
if "STATEFP20" in df.columns:
df["STATEFP"] = df["STATEFP20"]
if "STATEFP" not in df.columns:
df["STATEFP"] = self.statefip
df["place_type_original"] = self.geolevel
df.to_crs("EPSG:4326", inplace=True)
df["wkt"] = df.geometry.to_wkt() # TODO better handle this in db
df = pd.DataFrame(
df.drop(columns="geometry")
) # to get rid of the geo pandas
if self.geolevel == "roads_county":
# monkey patch
# "https://www2.census.gov/geo/tiger/TIGER2023/ROADS/tl_2023_69120_roads.zip".rsplit('/')[-1].split('_')[2]
county = current_file.split("_")[2]
df["countyfp"] = county
if self.geolevel == "addrfeat":
# Extract 5-digit county FIPS from filename: tl_2025_SSCCC_addrfeat.zip
county = current_file.split("_")[2]
df["countyfp"] = county
df["STATEFP"] = county[:2]
partition_cols = ["statefp"]
if self.geolevel == "areawater":
# AREAWATER shapefiles carry no STATEFP/COUNTYFP attributes, so both
# are derived from the filename: tl_2025_SSCCC_areawater.zip
county = current_file.split("_")[2]
df["countyfp"] = county
df["STATEFP"] = county[:2]
partition_cols = ["statefp"]
if self.geolevel == "zipcode20":
df["STATEFP"] = df["GEOID20"].str.slice(stop=2)
partition_cols = ["statefp"]
if self.geolevel == "vtd20":
partition_cols = ["statefp"]
mode = "append"
# Cast all object columns to string to avoid awswrangler
# UndetectedType errors on null-only columns
for col in df.select_dtypes(include=["object"]).columns:
df[col] = df[col].astype("string")
wr.s3.to_parquet(
df=df,
path=self.s3_root,
database=self.datalake_db,
dataset=update_table,
index=False,
mode=mode,
partition_cols=partition_cols,
table=self.output_athena_dataset_name,
dtype={
"cnectafp": "string",
"nectafp": "string",
"nctadvfp": "string",
"statefp": "string",
"plus4l": "string",
"plus4r": "string",
},
)
# df_count = wr.athena.read_sql_query(
# f"select count (*) as count from {self.output_athena_dataset_name}",
# database=self.datalake_db,
# )
# print(
# f"{self.output_athena_dataset_name} processed. Inserted rows: {df.shape} of total rows {df_count['count'][0]} "
# )
except Exception as ex:
print(f"shape_type:{self.geolevel}, statefb:{self.statefip}, {ex}")
raise
[docs]
def get_ftp_client(self):
from ftplib import FTP
ftp = FTP(self.FTP_SERVER)
ftp.login()
return ftp
[docs]
def get_urls(self, shp_type: str, year: str, stat_abbr=""):
"""Current tiger shape fils"""
state_fip = FipsCodes.ALL.get(stat_abbr, None)
url_dict: dict[str, str] = {
"state": f"https://www2.census.gov/geo/tiger/TIGER2023/STATE/tl_2023_us_state.zip",
"county": f"https://www2.census.gov/geo/tiger/TIGER{year}/COUNTY/tl_{year}_us_county.zip",
"county10": f"https://www2.census.gov/geo/tiger/TIGER{year}/COUNTY/{year}/tl_{year}_us_county10.zip", # SP
"county_sub": f"https://www2.census.gov/geo/tiger/TIGER{year}/COUSUB/tl_{year}_{state_fip}_cousub.zip",
"city": f"https://www2.census.gov/geo/tiger/TIGER{year}/PLACE/tl_{year}_{state_fip}_place.zip",
"tribal_aiannh": f"https://www2.census.gov/geo/tiger/TIGER{year}/AIANNH/tl_{year}_us_aiannh.zip",
"tribal_aitsn": f"https://www2.census.gov/geo/tiger/TIGER{year}/AITSN/tl_{year}_us_aitsn.zip",
"tribal_anrc": f"https://www2.census.gov/geo/tiger/TIGER{year}/ANRC/tl_{year}_02_anrc.zip",
"consolidate_city": f"https://www2.census.gov/geo/tiger/TIGER{year}/CONCITY/tl_{year}_{state_fip}_concity.zip",
"cb20": f"https://www2.census.gov/geo/tiger/TIGER{year}/TABBLOCK20/tl_{year}_{state_fip}_tabblock20.zip",
"cb10": f"https://www2.census.gov/geo/tiger/TIGER2010/TABBLOCK/2010/tl_2010_{state_fip}_tabblock10.zip",
"cbg20": f"https://www2.census.gov/geo/tiger/TIGER{year}/BG/tl_{year}_{state_fip}_bg.zip",
"tract": f"https://www2.census.gov/geo/tiger/TIGER{year}/TRACT/tl_{year}_{state_fip}_tract.zip",
"zipcode20": f"https://www2.census.gov/geo/tiger/TIGER{year}/ZCTA520/tl_{year}_us_zcta520.zip",
"consolidate_city20": f"https://www2.census.gov/geo/tiger/TIGER_RD18/LAYER/CONCITY/tl_rd22_{state_fip}_concity.zip",
"congressional_district_cb20": f"https://www2.census.gov/geo/tiger/TIGER{year}/CD/tl_{year}_us_cd116.zip",
"congressional_district": f"https://www2.census.gov/geo/tiger/TIGER{year}/CD/tl_{year}_{state_fip}_cd119.zip", # 119 is the latest congress number for 2024 2025
"metropolitan_micropolitan_statistical_area": f"https://www2.census.gov/geo/tiger/TIGER{year}/CBSA/tl_{year}_us_cbsa.zip",
"combined_statistical_area": f"https://www2.census.gov/geo/tiger/TIGER{year}/CSA/tl_{year}_us_csa.zip",
"metropolitan_division": f"https://www2.census.gov/geo/tiger/TIGER{year}/METDIV/tl_{year}_us_metdiv.zip",
"uac20": f"https://www2.census.gov/geo/tiger/TIGER2020/UAC/tl_2020_us_uac20.zip",
"uac10": f"https://www2.census.gov/geo/tiger/TIGER2010/UAC/tl_2010_us_uac10.zip",
"place": f"https://www2.census.gov/geo/tiger/TIGER{year}/PLACE/tl_{year}_{state_fip}_place.zip",
"sldu": f"https://www2.census.gov/geo/tiger/TIGER{year}/SLDU/tl_{year}_{state_fip}_sldu.zip",
"sldl": f"https://www2.census.gov/geo/tiger/TIGER{year}/SLDL/tl_{year}_{state_fip}_sldl.zip",
}
if shp_type in url_dict.keys():
return [
url_dict[shp_type].replace("https://www2.census.gov", self.FTP_SERVER)
]
if shp_type == "roads_county":
directory = f"geo/tiger/TIGER{year}/ROADS/"
files_to_download = self.get_file_list(directory)
return [
f"{self.FTP_SERVER}/{directory}/{file}" for file in files_to_download
]
if shp_type == "vtd20":
directory = f"geo/tiger/TIGER2020PL/LAYER/VTD/2020/"
files_to_download = self.get_file_list(directory)
return [
f"{self.FTP_SERVER}/{directory}/{file}" for file in files_to_download
]
if shp_type == "addrfeat":
directory = f"geo/tiger/TIGER{year}/ADDRFEAT/"
all_files = self.get_file_list(directory)
# Filter to files matching this state's FIPS code
# Filename pattern: tl_2025_SSCCC_addrfeat.zip (SS = state FIPS)
if state_fip:
files_to_download = [
f for f in all_files
if f.split("_")[2][:2] == state_fip
]
else:
files_to_download = all_files
return [
f"{self.FTP_SERVER}/{directory}/{file}" for file in files_to_download
]
if shp_type == "areawater":
directory = f"geo/tiger/TIGER{year}/AREAWATER/"
all_files = self.get_file_list(directory)
# Filter to files matching this state's FIPS code
# Filename pattern: tl_2025_SSCCC_areawater.zip (SS = state FIPS)
if state_fip:
files_to_download = [
f for f in all_files if f.split("_")[2][:2] == state_fip
]
else:
files_to_download = all_files
return [
f"{self.FTP_SERVER}/{directory}/{file}" for file in files_to_download
]
[docs]
def get_file_details_from_url(self, url):
current_directory = "/".join(url.split("/")[1:-1])
current_file = url.split("/")[-1]
print(f"current_directory: {current_directory}")
print(f"current_file: {current_file}")
return current_directory, current_file
[docs]
def get_file_list(self, directory: str):
if self.use_ftp:
ftp = self.get_ftp_client()
try:
ftp.cwd(directory)
file_list = ftp.nlst()
ftp.quit()
return file_list
except Exception as e:
raise Exception(f"Failed to list files in {directory}: {e}")
else:
import requests
from bs4 import BeautifulSoup
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
url = f"{self.HTTP_SERVER}/{directory}/"
try:
session = requests.Session()
retry_strategy = Retry(
total=3,
backoff_factor=2,
status_forcelist=[429, 500, 502, 503, 504],
allowed_methods=["GET"],
)
session.mount("http://", HTTPAdapter(max_retries=retry_strategy))
session.mount("https://", HTTPAdapter(max_retries=retry_strategy))
response = session.get(url, timeout=(30, 120))
response.raise_for_status()
soup = BeautifulSoup(response.text, "html.parser")
file_list = []
for a in soup.find_all("a", href=True):
if a["href"].endswith(".zip"):
file_list.append(a["href"])
print(f"file_list: {file_list}")
return file_list
except Exception as e:
raise Exception(f"Failed to list files in {directory}: {e}")
[docs]
def download_using_http(self, url, retry_count=3, retry_sleep=5):
"""
Download Census TIGER shapefile with retry logic and error handling.
Args:
url: URL of the file to download
output_dir: Directory to save the file
max_retries: Maximum number of retry attempts
"""
import requests
import os
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from tempfile import NamedTemporaryFile
temp_file = NamedTemporaryFile(suffix=".zip", mode="w+", delete=False)
temp_name = temp_file.name
# geo/tiger/TIGER2025/CD/tl_2025_01_cd119.zip
# geo/tiger/TIGER2025/CD/tl_2025_60_cd119.zip
# check if url is ftp or http
if url.startswith("ftp"):
url = url.replace(self.FTP_SERVER, self.HTTP_SERVER)
# Configure session with retry strategy
session = requests.Session()
retry_strategy = Retry(
total=retry_count,
backoff_factor=retry_sleep, # Wait 2, 4, 8 seconds between retries
status_forcelist=[429, 500, 502, 503, 504],
allowed_methods=["GET"],
)
adapter = HTTPAdapter(max_retries=retry_strategy)
session.mount("http://", adapter)
session.mount("https://", adapter)
# Set headers
session.headers.update(
{
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
"Connection": "keep-alive",
}
)
print(f"Starting download: {url}")
try:
# Download with streaming to handle large files
response = session.get(
url, stream=True, timeout=(30, 120) # (connect timeout, read timeout)
)
response.raise_for_status()
# Get file size if available
total_size = int(response.headers.get("content-length", 0))
# Download in chunks
chunk_size = 8192
downloaded = 0
with open(temp_name, "wb") as f:
for chunk in response.iter_content(chunk_size=chunk_size):
if chunk:
f.write(chunk)
downloaded += len(chunk)
# Print progress
if total_size > 0:
progress = (downloaded / total_size) * 100
print(
f"\rProgress: {progress:.1f}% ({downloaded:,} / {total_size:,} bytes)",
end="",
)
print(f"\n✓ Download completed: {temp_name}")
print(f"File size: {os.path.getsize(temp_name):,} bytes")
return temp_name
except Exception as e:
print(f"✗ Unexpected error: {e}")
raise
finally:
session.close()
[docs]
def download_file(self, url, retry_count=3, retry_sleep=5):
from tempfile import NamedTemporaryFile
import time
current_directory, current_file = self.get_file_details_from_url(url)
print(f"Downloading file from {url}")
print(f"Current directory: {current_directory}")
print(f"Current file: {current_file}")
print(f"Use FTP: {self.use_ftp}")
if self.use_ftp:
ftp = self.get_ftp_client()
for attempt in range(retry_count):
print(f"Attempt: {attempt}")
try:
ftp.cwd(current_directory)
temp_file = NamedTemporaryFile(
suffix=".zip", mode="w+", delete=False
)
temp_name = temp_file.name
ftp.retrbinary(f"RETR {current_file}", open(temp_name, "wb").write)
ftp.quit()
return temp_name
except Exception as e:
print(f"Calling {url} failed: {e}, rtrying...")
time.sleep(retry_sleep)
continue
raise Exception("Failed to download file after multiple attempts.")
else:
return self.download_using_http(url, retry_count, retry_sleep)
[docs]
def remove_temp_file(self, temp_file):
import os
try:
os.remove(temp_file)
except OSError as e:
print(f"Error: {e.strerror}")
[docs]
def execute(self, context):
try:
urls = self.get_urls(self.geolevel, self.vintage, self.state_abbr)
if not urls:
raise AirflowSkipException(
f"No files found for {self.geolevel} and {self.state_abbr}"
)
print(f"urls: {urls}")
is_multi_file = len(urls) > 1
for i, url in enumerate(urls):
update_table = True
print(f"Downloading file from {url}")
current_directory, current_file = self.get_file_details_from_url(url)
downloaded_file = self.download_file(url)
self.process_zip_file(downloaded_file, current_file, update_table)
self.remove_temp_file(downloaded_file)
context["params"][
"map_name"
] = f"{self.output_athena_dataset_name}-{current_file}"
except AirflowSkipException:
raise
except Exception as ex:
context["params"]["map_name"] = f"{self.output_athena_dataset_name}-{ex}"
raise
context["map_index_template"] = "{{ params.map_name }}"