Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
197 changes: 127 additions & 70 deletions data_engineering/data_sources/census_acs/census_extractor.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,25 +36,33 @@
RATE_LIMITS = {
"requests_per_day": 500,
"requests_per_minute": 10,
"delay_between_requests": 0.2 # seconds
"delay_between_requests": 0.2, # seconds
}


class CensusExtractor:
"""Extract data from US Census API"""

def __init__(self, api_key: Optional[str] = None):
"""Initialize Census extractor with API key"""
self.api_key = api_key or os.getenv('CENSUS_API_KEY')
self.api_key = api_key or os.getenv("CENSUS_API_KEY")
if not self.api_key:
# We don't raise error here to allow the script to be imported/tested
# but methods will fail if they need the key
logger.warning("CENSUS_API_KEY not found in environment. API requests will likely fail.")
logger.warning(
"CENSUS_API_KEY not found in environment. API requests will likely fail."
)

self.base_url = CENSUS_API_BASE_URL
self.session = requests.Session()
self.session.headers.update({
'User-Agent': 'Fentanyl-Awareness-Pipeline/1.0'
})
self.session.headers.update({"User-Agent": "Fentanyl-Awareness-Pipeline/1.0"})

def _redact(self, error: Exception) -> str:
"""Redact API key from error messages"""
error_msg = str(error)
if self.api_key and self.api_key in error_msg:
return error_msg.replace(self.api_key, "***REDACTED***")
return error_msg

def get_state_population_estimates(self, years: List[int] = None) -> pd.DataFrame:
"""
Expand All @@ -80,48 +88,55 @@ def get_state_population_estimates(self, years: List[int] = None) -> pd.DataFram
# ACS 5-Year Estimates endpoint
url = f"{self.base_url}/{year}/acs/acs5"

params = {
'get': 'NAME,B01001_001E', # B01001_001E is total population
'for': 'state:*',
'key': self.api_key
} if self.api_key else {
'get': 'NAME,B01001_001E',
'for': 'state:*'
}
params = (
{
"get": "NAME,B01001_001E", # B01001_001E is total population
"for": "state:*",
"key": self.api_key,
}
if self.api_key
else {"get": "NAME,B01001_001E", "for": "state:*"}
)

logger.info(f"Fetching data for year {year}...")
response = self.session.get(url, params=params, timeout=30)

if response.status_code == 404:
logger.warning(f"Data for year {year} not available yet (404). Skipping.")
logger.warning(
f"Data for year {year} not available yet (404). Skipping."
)
continue

response.raise_for_status()

# Check if the response is actually JSON
if 'application/json' not in response.headers.get('Content-Type', ''):
logger.error(f"Error fetching data for year {year}: Expected JSON but received {response.headers.get('Content-Type')}. "
f"This often indicates a missing API key or an invalid endpoint.")
if "application/json" not in response.headers.get("Content-Type", ""):
logger.error(
f"Error fetching data for year {year}: Expected JSON but received {response.headers.get('Content-Type')}. "
f"This often indicates a missing API key or an invalid endpoint."
)
continue

try:
data = response.json()
except Exception as e:
logger.error(f"Error parsing JSON for year {year}: {e}")
logger.error(
f"Error parsing JSON for year {year}: {self._redact(e)}"
)
continue

# Convert to DataFrame
df = pd.DataFrame(data[1:], columns=data[0])
df['year'] = year
df['extracted_at'] = datetime.now()
df["year"] = year
df["extracted_at"] = datetime.now()

all_data.append(df)

# Rate limiting
time.sleep(RATE_LIMITS["delay_between_requests"])

except Exception as e:
logger.error(f"Error fetching data for year {year}: {e}")
logger.error(f"Error fetching data for year {year}: {self._redact(e)}")
continue

if not all_data:
Expand Down Expand Up @@ -159,48 +174,60 @@ def get_state_economic_data(self, years: List[int] = None) -> pd.DataFrame:
# ACS 5-Year Estimates endpoint
url = f"{self.base_url}/{year}/acs/acs5"

params = {
'get': 'B19013_001E,B19301_001E,B23025_002E,B23025_003E,B23025_004E,B23025_005E,NAME',
'for': 'state:*',
'key': self.api_key
} if self.api_key else {
'get': 'B19013_001E,B19301_001E,B23025_002E,B23025_003E,B23025_004E,B23025_005E,NAME',
'for': 'state:*'
}
params = (
{
"get": "B19013_001E,B19301_001E,B23025_002E,B23025_003E,B23025_004E,B23025_005E,NAME",
"for": "state:*",
"key": self.api_key,
}
if self.api_key
else {
"get": "B19013_001E,B19301_001E,B23025_002E,B23025_003E,B23025_004E,B23025_005E,NAME",
"for": "state:*",
}
)

logger.info(f"Fetching economic data for year {year}...")
response = self.session.get(url, params=params, timeout=30)

if response.status_code == 404:
logger.warning(f"Economic data for year {year} not available yet (404). Skipping.")
logger.warning(
f"Economic data for year {year} not available yet (404). Skipping."
)
continue

response.raise_for_status()

# Check if the response is actually JSON
if 'application/json' not in response.headers.get('Content-Type', ''):
logger.error(f"Error fetching economic data for year {year}: Expected JSON but received {response.headers.get('Content-Type')}. "
f"This often indicates a missing API key or an invalid endpoint.")
if "application/json" not in response.headers.get("Content-Type", ""):
logger.error(
f"Error fetching economic data for year {year}: Expected JSON but received {response.headers.get('Content-Type')}. "
f"This often indicates a missing API key or an invalid endpoint."
)
continue

try:
data = response.json()
except Exception as e:
logger.error(f"Error parsing economic JSON for year {year}: {e}")
logger.error(
f"Error parsing economic JSON for year {year}: {self._redact(e)}"
)
continue

# Convert to DataFrame
df = pd.DataFrame(data[1:], columns=data[0])
df['year'] = year
df['extracted_at'] = datetime.now()
df["year"] = year
df["extracted_at"] = datetime.now()

all_data.append(df)

# Rate limiting
time.sleep(RATE_LIMITS["delay_between_requests"])

except Exception as e:
logger.error(f"Error fetching economic data for year {year}: {e}")
logger.error(
f"Error fetching economic data for year {year}: {self._redact(e)}"
)
continue

if not all_data:
Expand All @@ -219,26 +246,32 @@ def _clean_population_data(self, df: pd.DataFrame) -> pd.DataFrame:
"""Clean and standardize population data from ACS"""

# Convert numeric columns
df['B01001_001E'] = pd.to_numeric(df['B01001_001E'], errors='coerce')
df['state'] = pd.to_numeric(df['state'], errors='coerce')
df["B01001_001E"] = pd.to_numeric(df["B01001_001E"], errors="coerce")
df["state"] = pd.to_numeric(df["state"], errors="coerce")

# Rename columns for consistency
df = df.rename(columns={
'B01001_001E': 'population',
'state': 'state_code',
'NAME': 'state_name'
})
df = df.rename(
columns={
"B01001_001E": "population",
"state": "state_code",
"NAME": "state_name",
}
)

# Clean state names (remove extra text)
df['state_name'] = df['state_name'].str.replace(',', '').str.strip()
df["state_name"] = df["state_name"].str.replace(",", "").str.strip()

# Add data description
df['date_description'] = 'ACS 5-Year Estimate'
df["date_description"] = "ACS 5-Year Estimate"

# Select final columns
final_columns = [
'year', 'state_code', 'state_name', 'population',
'date_description', 'extracted_at'
"year",
"state_code",
"state_name",
"population",
"date_description",
"extracted_at",
]

return df[final_columns].dropna()
Expand All @@ -247,42 +280,60 @@ def _clean_economic_data(self, df: pd.DataFrame) -> pd.DataFrame:
"""Clean and standardize economic data"""

# Convert numeric columns
numeric_columns = ['B19013_001E', 'B19301_001E', 'B23025_002E',
'B23025_003E', 'B23025_004E', 'B23025_005E']
numeric_columns = [
"B19013_001E",
"B19301_001E",
"B23025_002E",
"B23025_003E",
"B23025_004E",
"B23025_005E",
]

for col in numeric_columns:
df[col] = pd.to_numeric(df[col], errors='coerce')
df[col] = pd.to_numeric(df[col], errors="coerce")

# Convert state code to numeric
df['state'] = pd.to_numeric(df['state'], errors='coerce')
df["state"] = pd.to_numeric(df["state"], errors="coerce")

# Rename columns to meaningful names
df = df.rename(columns={
'B19013_001E': 'median_household_income',
'B19301_001E': 'per_capita_income',
'B23025_002E': 'labor_force_total',
'B23025_003E': 'labor_force_civilian',
'B23025_004E': 'employed',
'B23025_005E': 'unemployed',
'NAME': 'state_name'
})
df = df.rename(
columns={
"B19013_001E": "median_household_income",
"B19301_001E": "per_capita_income",
"B23025_002E": "labor_force_total",
"B23025_003E": "labor_force_civilian",
"B23025_004E": "employed",
"B23025_005E": "unemployed",
"NAME": "state_name",
}
)

# Calculate unemployment rate (handle division by zero and NaN)
df['unemployment_rate'] = (df['unemployed'] / df['labor_force_civilian'] * 100).round(2)
df['unemployment_rate'] = df['unemployment_rate'].fillna(0)
df["unemployment_rate"] = (
df["unemployed"] / df["labor_force_civilian"] * 100
).round(2)
df["unemployment_rate"] = df["unemployment_rate"].fillna(0)

# Use the state field directly (already converted to numeric above)
df['state_code'] = df['state'].astype(int)
df["state_code"] = df["state"].astype(int)

# Select final columns
final_columns = [
'year', 'state_code', 'state_name', 'median_household_income',
'per_capita_income', 'unemployment_rate', 'labor_force_total',
'employed', 'unemployed', 'extracted_at'
"year",
"state_code",
"state_name",
"median_household_income",
"per_capita_income",
"unemployment_rate",
"labor_force_total",
"employed",
"unemployed",
"extracted_at",
]

return df[final_columns].dropna()


def test_census_api():
"""Test Census API connection and data extraction"""
try:
Expand All @@ -299,6 +350,7 @@ def test_census_api():
print(f"❌ Error during test: {e}")
return False


def main():
"""Main function to extract Census data"""

Expand All @@ -316,7 +368,9 @@ def main():
# Save to CSV files (overwrite if they exist)
# Using Path(__file__) for robust resolution
script_path = Path(__file__).resolve()
output_dir = script_path.parent.parent.parent / "data_build_tool" / "dbt" / "seeds"
output_dir = (
script_path.parent.parent.parent / "data_build_tool" / "dbt" / "seeds"
)
output_dir.mkdir(parents=True, exist_ok=True)

population_file = output_dir / "census_state_population.csv"
Expand All @@ -338,7 +392,9 @@ def main():
print(f"\n📊 Census Data Extraction Summary:")
print(f"Population records: {len(population_df)}")
print(f"Economic records: {len(economic_df)}")
print(f"Years covered: {population_df['year'].min()}-{population_df['year'].max()}")
print(
f"Years covered: {population_df['year'].min()}-{population_df['year'].max()}"
)
print(f"States covered: {population_df['state_name'].nunique()}")

except Exception as e:
Expand All @@ -347,5 +403,6 @@ def main():

return 0


if __name__ == "__main__":
exit(main())
Loading