mirror of
https://github.com/RichieCahill/dotfiles.git
synced 2026-04-17 13:08:19 -04:00
499 lines
18 KiB
Python
499 lines
18 KiB
Python
"""Ingestion pipeline for loading congress data from unitedstates/congress JSON files.
|
|
|
|
Loads legislators, bills, votes, vote records, and bill text into the data_science_dev database.
|
|
|
|
Usage:
|
|
ingest-congress /path/to/congress/data/
|
|
ingest-congress /path/to/congress/data/ --congress 118
|
|
ingest-congress /path/to/congress/data/ --congress 118 --only bills
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from pathlib import Path # noqa: TC003 needed at runtime for typer CLI argument
|
|
from typing import TYPE_CHECKING, Annotated
|
|
|
|
import orjson
|
|
import typer
|
|
from sqlalchemy import select
|
|
from sqlalchemy.orm import Session
|
|
|
|
from python.common import configure_logger
|
|
from python.orm.common import get_postgres_engine
|
|
from python.orm.data_science_dev.congress import Bill, BillText, Legislator, Vote, VoteRecord
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Iterator
|
|
|
|
from sqlalchemy.engine import Engine
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
BATCH_SIZE = 10_000
|
|
|
|
app = typer.Typer(help="Ingest unitedstates/congress data into data_science_dev.")
|
|
|
|
|
|
@app.command()
|
|
def main(
|
|
data_dir: Annotated[Path, typer.Argument(help="Path to the congress data/ directory")],
|
|
congress: Annotated[int | None, typer.Option(help="Only ingest a specific congress number")] = None,
|
|
only: Annotated[
|
|
str | None,
|
|
typer.Option(help="Only run a specific step: legislators, bills, votes, bill-text"),
|
|
] = None,
|
|
) -> None:
|
|
"""Ingest congress data from unitedstates/congress JSON files."""
|
|
configure_logger(level="INFO")
|
|
|
|
if not data_dir.is_dir():
|
|
typer.echo(f"Data directory does not exist: {data_dir}", err=True)
|
|
raise typer.Exit(code=1)
|
|
|
|
engine = get_postgres_engine(name="DATA_SCIENCE_DEV")
|
|
|
|
congress_dirs = _resolve_congress_dirs(data_dir, congress)
|
|
if not congress_dirs:
|
|
typer.echo("No congress directories found.", err=True)
|
|
raise typer.Exit(code=1)
|
|
|
|
logger.info("Found %d congress directories to process", len(congress_dirs))
|
|
|
|
steps = {
|
|
"legislators": ingest_legislators,
|
|
"bills": ingest_bills,
|
|
"votes": ingest_votes,
|
|
"bill-text": ingest_bill_text,
|
|
}
|
|
|
|
if only:
|
|
if only not in steps:
|
|
typer.echo(f"Unknown step: {only}. Choose from: {', '.join(steps)}", err=True)
|
|
raise typer.Exit(code=1)
|
|
steps = {only: steps[only]}
|
|
|
|
for step_name, step_func in steps.items():
|
|
logger.info("=== Starting step: %s ===", step_name)
|
|
step_func(engine, congress_dirs)
|
|
logger.info("=== Finished step: %s ===", step_name)
|
|
|
|
logger.info("ingest-congress done")
|
|
|
|
|
|
def _resolve_congress_dirs(data_dir: Path, congress: int | None) -> list[Path]:
|
|
"""Find congress number directories under data_dir."""
|
|
if congress is not None:
|
|
target = data_dir / str(congress)
|
|
return [target] if target.is_dir() else []
|
|
return sorted(path for path in data_dir.iterdir() if path.is_dir() and path.name.isdigit())
|
|
|
|
|
|
def _flush_batch(session: Session, batch: list[object], label: str) -> int:
|
|
"""Add a batch of ORM objects to the session and commit. Returns count added."""
|
|
if not batch:
|
|
return 0
|
|
session.add_all(batch)
|
|
session.commit()
|
|
count = len(batch)
|
|
logger.info("Committed %d %s", count, label)
|
|
batch.clear()
|
|
return count
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Legislators — extracted from vote JSON files (voter records include bioguide_id, name, party, state)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def ingest_legislators(engine: Engine, congress_dirs: list[Path]) -> None:
|
|
"""Extract unique legislators from vote data and insert them."""
|
|
with Session(engine) as session:
|
|
seen_bioguide_ids = set(session.scalars(select(Legislator.bioguide_id)).all())
|
|
logger.info("Found %d existing legislators in DB", len(seen_bioguide_ids))
|
|
|
|
total_inserted = 0
|
|
batch: list[Legislator] = []
|
|
for congress_dir in congress_dirs:
|
|
votes_dir = congress_dir / "votes"
|
|
if not votes_dir.is_dir():
|
|
continue
|
|
logger.info("Scanning legislators from %s", congress_dir.name)
|
|
for vote_file in votes_dir.rglob("data.json"):
|
|
data = _read_json(vote_file)
|
|
if data is None:
|
|
continue
|
|
_extract_legislators(data, seen_bioguide_ids, batch)
|
|
if len(batch) >= BATCH_SIZE:
|
|
total_inserted += _flush_batch(session, batch, "legislators")
|
|
|
|
total_inserted += _flush_batch(session, batch, "legislators")
|
|
logger.info("Inserted %d new legislators total", total_inserted)
|
|
|
|
|
|
def _iter_voters(position_group: object) -> Iterator[dict]:
|
|
"""Yield voter dicts from a vote position group (handles list, single dict, or string)."""
|
|
if isinstance(position_group, dict):
|
|
yield position_group
|
|
elif isinstance(position_group, list):
|
|
for voter in position_group:
|
|
if isinstance(voter, dict):
|
|
yield voter
|
|
|
|
|
|
def _extract_legislators(data: dict, seen_bioguide_ids: set[str], batch: list[Legislator]) -> None:
|
|
"""Extract unique legislators from a single vote data.json."""
|
|
for position_group in data.get("votes", {}).values():
|
|
for voter in _iter_voters(position_group):
|
|
bioguide_id = voter.get("id")
|
|
if not bioguide_id or bioguide_id in seen_bioguide_ids:
|
|
continue
|
|
seen_bioguide_ids.add(bioguide_id)
|
|
first_name, last_name = _split_name(voter.get("display_name", ""))
|
|
batch.append(
|
|
Legislator(
|
|
bioguide_id=bioguide_id,
|
|
first_name=first_name,
|
|
last_name=last_name,
|
|
current_party=voter.get("party"),
|
|
current_state=voter.get("state"),
|
|
)
|
|
)
|
|
|
|
|
|
def _split_name(display_name: str) -> tuple[str, str]:
|
|
"""Split 'Last, First' or 'Name' into (first, last)."""
|
|
if "," in display_name:
|
|
parts = display_name.split(",", 1)
|
|
return parts[1].strip(), parts[0].strip()
|
|
parts = display_name.rsplit(" ", 1)
|
|
if len(parts) > 1:
|
|
return parts[0].strip(), parts[1].strip()
|
|
return display_name, ""
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Bills
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def ingest_bills(engine: Engine, congress_dirs: list[Path]) -> None:
|
|
"""Load bill data.json files."""
|
|
with Session(engine) as session:
|
|
existing_bills = {(bill.congress, bill.bill_type, bill.number) for bill in session.scalars(select(Bill)).all()}
|
|
logger.info("Found %d existing bills in DB", len(existing_bills))
|
|
|
|
total_inserted = 0
|
|
batch: list[Bill] = []
|
|
for congress_dir in congress_dirs:
|
|
bills_dir = congress_dir / "bills"
|
|
if not bills_dir.is_dir():
|
|
continue
|
|
logger.info("Scanning bills from %s", congress_dir.name)
|
|
for bill_file in bills_dir.rglob("data.json"):
|
|
data = _read_json(bill_file)
|
|
if data is None:
|
|
continue
|
|
bill = _parse_bill(data, existing_bills)
|
|
if bill is not None:
|
|
batch.append(bill)
|
|
if len(batch) >= BATCH_SIZE:
|
|
total_inserted += _flush_batch(session, batch, "bills")
|
|
|
|
total_inserted += _flush_batch(session, batch, "bills")
|
|
logger.info("Inserted %d new bills total", total_inserted)
|
|
|
|
|
|
def _parse_bill(data: dict, existing_bills: set[tuple[int, str, int]]) -> Bill | None:
|
|
"""Parse a bill data.json dict into a Bill ORM object, skipping existing."""
|
|
raw_congress = data.get("congress")
|
|
bill_type = data.get("bill_type")
|
|
raw_number = data.get("number")
|
|
if raw_congress is None or bill_type is None or raw_number is None:
|
|
return None
|
|
congress = int(raw_congress)
|
|
number = int(raw_number)
|
|
if (congress, bill_type, number) in existing_bills:
|
|
return None
|
|
|
|
sponsor_bioguide = None
|
|
sponsor = data.get("sponsor")
|
|
if sponsor:
|
|
sponsor_bioguide = sponsor.get("bioguide_id")
|
|
|
|
return Bill(
|
|
congress=congress,
|
|
bill_type=bill_type,
|
|
number=number,
|
|
title=data.get("short_title") or data.get("official_title"),
|
|
title_short=data.get("short_title"),
|
|
official_title=data.get("official_title"),
|
|
status=data.get("status"),
|
|
status_at=data.get("status_at"),
|
|
sponsor_bioguide_id=sponsor_bioguide,
|
|
subjects_top_term=data.get("subjects_top_term"),
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Votes (and vote records)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def ingest_votes(engine: Engine, congress_dirs: list[Path]) -> None:
|
|
"""Load vote data.json files with their vote records."""
|
|
with Session(engine) as session:
|
|
legislator_map = _build_legislator_map(session)
|
|
logger.info("Loaded %d legislators into lookup map", len(legislator_map))
|
|
bill_map = _build_bill_map(session)
|
|
logger.info("Loaded %d bills into lookup map", len(bill_map))
|
|
existing_votes = {
|
|
(vote.congress, vote.chamber, vote.session, vote.number) for vote in session.scalars(select(Vote)).all()
|
|
}
|
|
logger.info("Found %d existing votes in DB", len(existing_votes))
|
|
|
|
total_inserted = 0
|
|
batch: list[Vote] = []
|
|
for congress_dir in congress_dirs:
|
|
votes_dir = congress_dir / "votes"
|
|
if not votes_dir.is_dir():
|
|
continue
|
|
logger.info("Scanning votes from %s", congress_dir.name)
|
|
for vote_file in votes_dir.rglob("data.json"):
|
|
data = _read_json(vote_file)
|
|
if data is None:
|
|
continue
|
|
vote = _parse_vote(data, legislator_map, bill_map, existing_votes)
|
|
if vote is not None:
|
|
batch.append(vote)
|
|
if len(batch) >= BATCH_SIZE:
|
|
total_inserted += _flush_batch(session, batch, "votes")
|
|
|
|
total_inserted += _flush_batch(session, batch, "votes")
|
|
logger.info("Inserted %d new votes total", total_inserted)
|
|
|
|
|
|
def _build_legislator_map(session: Session) -> dict[str, int]:
|
|
"""Build a mapping of bioguide_id -> legislator.id."""
|
|
return {legislator.bioguide_id: legislator.id for legislator in session.scalars(select(Legislator)).all()}
|
|
|
|
|
|
def _build_bill_map(session: Session) -> dict[tuple[int, str, int], int]:
|
|
"""Build a mapping of (congress, bill_type, number) -> bill.id."""
|
|
return {(bill.congress, bill.bill_type, bill.number): bill.id for bill in session.scalars(select(Bill)).all()}
|
|
|
|
|
|
def _parse_vote(
|
|
data: dict,
|
|
legislator_map: dict[str, int],
|
|
bill_map: dict[tuple[int, str, int], int],
|
|
existing_votes: set[tuple[int, str, int, int]],
|
|
) -> Vote | None:
|
|
"""Parse a vote data.json dict into a Vote ORM object with records."""
|
|
raw_congress = data.get("congress")
|
|
chamber = data.get("chamber")
|
|
raw_number = data.get("number")
|
|
vote_date = data.get("date")
|
|
if raw_congress is None or chamber is None or raw_number is None or vote_date is None:
|
|
return None
|
|
|
|
raw_session = data.get("session")
|
|
if raw_session is None:
|
|
return None
|
|
|
|
congress = int(raw_congress)
|
|
number = int(raw_number)
|
|
session_number = int(raw_session)
|
|
|
|
# Normalize chamber from "h"/"s" to "House"/"Senate"
|
|
chamber_normalized = {"h": "House", "s": "Senate"}.get(chamber, chamber)
|
|
|
|
if (congress, chamber_normalized, session_number, number) in existing_votes:
|
|
return None
|
|
|
|
# Resolve linked bill
|
|
bill_id = None
|
|
bill_ref = data.get("bill")
|
|
if bill_ref:
|
|
bill_key = (
|
|
int(bill_ref.get("congress", congress)),
|
|
bill_ref.get("type"),
|
|
int(bill_ref.get("number", 0)),
|
|
)
|
|
bill_id = bill_map.get(bill_key)
|
|
|
|
raw_votes = data.get("votes", {})
|
|
vote_counts = _count_votes(raw_votes)
|
|
vote_records = _build_vote_records(raw_votes, legislator_map)
|
|
|
|
return Vote(
|
|
congress=congress,
|
|
chamber=chamber_normalized,
|
|
session=session_number,
|
|
number=number,
|
|
vote_type=data.get("type"),
|
|
question=data.get("question"),
|
|
result=data.get("result"),
|
|
result_text=data.get("result_text"),
|
|
vote_date=vote_date[:10] if isinstance(vote_date, str) else vote_date,
|
|
bill_id=bill_id,
|
|
vote_records=vote_records,
|
|
**vote_counts,
|
|
)
|
|
|
|
|
|
def _count_votes(raw_votes: dict) -> dict[str, int]:
|
|
"""Count voters per position category, correctly handling dict and list formats."""
|
|
yea_count = 0
|
|
nay_count = 0
|
|
not_voting_count = 0
|
|
present_count = 0
|
|
|
|
for position, position_group in raw_votes.items():
|
|
voter_count = sum(1 for _ in _iter_voters(position_group))
|
|
if position in ("Yea", "Aye"):
|
|
yea_count += voter_count
|
|
elif position in ("Nay", "No"):
|
|
nay_count += voter_count
|
|
elif position == "Not Voting":
|
|
not_voting_count += voter_count
|
|
elif position == "Present":
|
|
present_count += voter_count
|
|
|
|
return {
|
|
"yea_count": yea_count,
|
|
"nay_count": nay_count,
|
|
"not_voting_count": not_voting_count,
|
|
"present_count": present_count,
|
|
}
|
|
|
|
|
|
def _build_vote_records(raw_votes: dict, legislator_map: dict[str, int]) -> list[VoteRecord]:
|
|
"""Build VoteRecord objects from raw vote data."""
|
|
records: list[VoteRecord] = []
|
|
for position, position_group in raw_votes.items():
|
|
for voter in _iter_voters(position_group):
|
|
bioguide_id = voter.get("id")
|
|
if not bioguide_id:
|
|
continue
|
|
legislator_id = legislator_map.get(bioguide_id)
|
|
if legislator_id is None:
|
|
continue
|
|
records.append(
|
|
VoteRecord(
|
|
legislator_id=legislator_id,
|
|
position=position,
|
|
)
|
|
)
|
|
return records
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Bill Text
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def ingest_bill_text(engine: Engine, congress_dirs: list[Path]) -> None:
|
|
"""Load bill text from text-versions directories."""
|
|
with Session(engine) as session:
|
|
bill_map = _build_bill_map(session)
|
|
logger.info("Loaded %d bills into lookup map", len(bill_map))
|
|
existing_bill_texts = {
|
|
(bill_text.bill_id, bill_text.version_code) for bill_text in session.scalars(select(BillText)).all()
|
|
}
|
|
logger.info("Found %d existing bill text versions in DB", len(existing_bill_texts))
|
|
|
|
total_inserted = 0
|
|
batch: list[BillText] = []
|
|
for congress_dir in congress_dirs:
|
|
logger.info("Scanning bill texts from %s", congress_dir.name)
|
|
for bill_text in _iter_bill_texts(congress_dir, bill_map, existing_bill_texts):
|
|
batch.append(bill_text)
|
|
if len(batch) >= BATCH_SIZE:
|
|
total_inserted += _flush_batch(session, batch, "bill texts")
|
|
|
|
total_inserted += _flush_batch(session, batch, "bill texts")
|
|
logger.info("Inserted %d new bill text versions total", total_inserted)
|
|
|
|
|
|
def _iter_bill_texts(
|
|
congress_dir: Path,
|
|
bill_map: dict[tuple[int, str, int], int],
|
|
existing_bill_texts: set[tuple[int, str]],
|
|
) -> Iterator[BillText]:
|
|
"""Yield BillText objects for a single congress directory, skipping existing."""
|
|
bills_dir = congress_dir / "bills"
|
|
if not bills_dir.is_dir():
|
|
return
|
|
|
|
for bill_dir in bills_dir.rglob("text-versions"):
|
|
if not bill_dir.is_dir():
|
|
continue
|
|
bill_key = _bill_key_from_dir(bill_dir.parent, congress_dir)
|
|
if bill_key is None:
|
|
continue
|
|
bill_id = bill_map.get(bill_key)
|
|
if bill_id is None:
|
|
continue
|
|
|
|
for version_dir in sorted(bill_dir.iterdir()):
|
|
if not version_dir.is_dir():
|
|
continue
|
|
if (bill_id, version_dir.name) in existing_bill_texts:
|
|
continue
|
|
text_content = _read_bill_text(version_dir)
|
|
version_data = _read_json(version_dir / "data.json")
|
|
yield BillText(
|
|
bill_id=bill_id,
|
|
version_code=version_dir.name,
|
|
version_name=version_data.get("version_name") if version_data else None,
|
|
date=version_data.get("issued_on") if version_data else None,
|
|
text_content=text_content,
|
|
)
|
|
|
|
|
|
def _bill_key_from_dir(bill_dir: Path, congress_dir: Path) -> tuple[int, str, int] | None:
|
|
"""Extract (congress, bill_type, number) from directory structure."""
|
|
congress = int(congress_dir.name)
|
|
bill_type = bill_dir.parent.name
|
|
name = bill_dir.name
|
|
# Directory name is like "hr3590" — strip the type prefix to get the number
|
|
number_str = name[len(bill_type) :]
|
|
if not number_str.isdigit():
|
|
return None
|
|
return (congress, bill_type, int(number_str))
|
|
|
|
|
|
def _read_bill_text(version_dir: Path) -> str | None:
|
|
"""Read bill text from a version directory, preferring .txt over .xml."""
|
|
for extension in ("txt", "htm", "html", "xml"):
|
|
candidates = list(version_dir.glob(f"document.{extension}"))
|
|
if not candidates:
|
|
candidates = list(version_dir.glob(f"*.{extension}"))
|
|
if candidates:
|
|
try:
|
|
return candidates[0].read_text(encoding="utf-8")
|
|
except Exception:
|
|
logger.exception("Failed to read %s", candidates[0])
|
|
return None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _read_json(path: Path) -> dict | None:
|
|
"""Read and parse a JSON file, returning None on failure."""
|
|
try:
|
|
return orjson.loads(path.read_bytes())
|
|
except FileNotFoundError:
|
|
return None
|
|
except Exception:
|
|
logger.exception("Failed to parse %s", path)
|
|
return None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app()
|