diff --git a/src/spicy_regs/ingest_sam_extract.py b/src/spicy_regs/ingest_sam_extract.py new file mode 100644 index 0000000..d6d429f --- /dev/null +++ b/src/spicy_regs/ingest_sam_extract.py @@ -0,0 +1,170 @@ +"""Ingest the SAM.gov public monthly *entity extract* into ``sam_entities``. + +The SAM Entity Management API is daily-quota-gated (≈1,000 req/day for a +non-federal keyed account) and caps a single query's pagination at ~5,000 +records, so the live-API reader (``sources/sam_entities.py``) can only ever +sample the ~885K registered entities. The **public monthly extract** is the +full set in one flat file — this module parses it and publishes the complete +``sam_entities`` table. + +Get the file (no API quota involved): SAM.gov → **Data Services** → *Entity +Registration* → *Public* → download ``SAM_PUBLIC_MONTHLY_V2_YYYYMMDD.ZIP`` and +unzip the ``.dat``. (Or the extract API: +``https://api.sam.gov/data-services/v1/extracts?api_key=...&fileType=ENTITY&sensitivity=PUBLIC``.) + +Usage:: + + uv run python -m spicy_regs.ingest_sam_extract SAM_PUBLIC_MONTHLY_V2_20260705.dat + uv run python -m spicy_regs.ingest_sam_extract .dat --output out.parquet --upload + +File format: pipe-delimited, **142 fields**, no quoting; a ``BOF PUBLIC V2 ...`` +header line and an ``EOF PUBLIC V2 ...`` trailer line wrap the data records. + +Field map (1-indexed extract position → published column), confirmed against the +2026-07 extract: + +=================================== ======== ================================ +column position notes +=================================== ======== ================================ +uei 1 12-char Unique Entity ID (PK) +cage_code 4 +legal_business_name 12 +dba_name 13 +entity_structure_desc 28 raw structure CODE (see note) +state 19 physical-address state +city 18 +zip_code 20 +congressional_district 23 +primary_naics 33 primary NAICS (list starts here) +registration_status 6 A→Active, E→Expired +registration_date 8 YYYYMMDD → YYYY-MM-DD +registration_expiration_date 9 YYYYMMDD → YYYY-MM-DD +purpose_of_registration_desc 7 Z1/Z2/Z5 mapped, else raw code +entity_url 27 +=================================== ======== ================================ + +Extract limitations vs the API: the extract does not cleanly expose +``entity_type_desc``, ``profit_structure_desc``, or ``exclusion_status_flag`` +(exclusions are a separate SAM extract), so those are published NULL; and +``entity_structure_desc`` carries SAM's raw structure code rather than the API's +prose description. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import pyarrow.parquet as pq +from loguru import logger + +OUTPUT = "sam_entities.parquet" + +# 1-indexed extract position -> 0-indexed DuckDB column name (read as c0..c141). +POS = { + "uei": "c0", + "cage_code": "c3", + "legal_business_name": "c11", + "dba_name": "c12", + "entity_structure_desc": "c27", + "state": "c18", + "city": "c17", + "zip_code": "c19", + "congressional_district": "c22", + "primary_naics": "c32", + "registration_status_code": "c5", + "purpose_code": "c6", + "registration_date_raw": "c7", + "registration_expiration_date_raw": "c8", + "entity_url": "c26", +} + +# SAM code -> description for the small, stable code sets we surface. +STATUS_MAP = {"A": "Active", "E": "Expired"} +PURPOSE_MAP = {"Z1": "Federal Assistance Awards", "Z2": "All Awards", "Z5": "IGT Only"} + +_NUM_FIELDS = 142 + + +def _read_csv_expr(src: Path) -> str: + cols = "{" + ",".join(f"'c{i}':'VARCHAR'" for i in range(_NUM_FIELDS)) + "}" + # No quoting (flat file with literal quotes in data); skip the BOF header; + # the EOF trailer is dropped by the UEI-length filter below. + return ( + f"read_csv('{src}', delim='|', header=false, auto_detect=false, quote='', " + f"escape='', null_padding=true, ignore_errors=true, parallel=false, skip=1, " + f"max_line_size=20000000, columns={cols})" + ) + + +def _case_map(col: str, mapping: dict[str, str]) -> str: + whens = " ".join(f"WHEN '{k}' THEN '{v}'" for k, v in mapping.items()) + return f"CASE trim({col}) {whens} ELSE NULLIF(trim({col}),'') END" + + +def build_sam_entities_from_extract(src: Path, out: Path) -> Path: + """Parse a SAM public entity extract ``.dat`` into ``sam_entities`` parquet.""" + import duckdb + + def nz(c: str) -> str: + return f"NULLIF(trim({c}),'')" + + def dt(c: str) -> str: + return f"CASE WHEN length(trim({c}))=8 THEN substr({c},1,4)||'-'||substr({c},5,2)||'-'||substr({c},7,2) END" + + logger.info("SAM extract: parsing {}", src) + con = duckdb.connect() + con.execute( + f""" + COPY ( + SELECT uei, cage_code, legal_business_name, dba_name, entity_structure_desc, + entity_type_desc, profit_structure_desc, state, city, zip_code, + congressional_district, primary_naics, registration_status, + registration_date, registration_expiration_date, exclusion_status_flag, + purpose_of_registration_desc, entity_url + FROM ( + SELECT {nz(POS["uei"])} uei, {nz(POS["cage_code"])} cage_code, + {nz(POS["legal_business_name"])} legal_business_name, {nz(POS["dba_name"])} dba_name, + {nz(POS["entity_structure_desc"])} entity_structure_desc, + CAST(NULL AS VARCHAR) entity_type_desc, CAST(NULL AS VARCHAR) profit_structure_desc, + {nz(POS["state"])} state, {nz(POS["city"])} city, {nz(POS["zip_code"])} zip_code, + {nz(POS["congressional_district"])} congressional_district, {nz(POS["primary_naics"])} primary_naics, + {_case_map(POS["registration_status_code"], STATUS_MAP)} registration_status, + {dt(POS["registration_date_raw"])} registration_date, + {dt(POS["registration_expiration_date_raw"])} registration_expiration_date, + CAST(NULL AS VARCHAR) exclusion_status_flag, + {_case_map(POS["purpose_code"], PURPOSE_MAP)} purpose_of_registration_desc, + {nz(POS["entity_url"])} entity_url, + ROW_NUMBER() OVER ( + PARTITION BY {POS["uei"]} + ORDER BY ({POS["registration_status_code"]}='A') DESC, {POS["registration_expiration_date_raw"]} DESC + ) rn + FROM {_read_csv_expr(src)} + WHERE length(trim({POS["uei"]}))=12 + ) WHERE rn=1 + ) TO '{out}' (FORMAT PARQUET, COMPRESSION ZSTD, ROW_GROUP_SIZE 100000); + """ + ) + con.close() + rows = pq.ParquetFile(out).metadata.num_rows + logger.info("SAM extract: {:,} entities -> {}", rows, out) + return out + + +def main() -> None: + ap = argparse.ArgumentParser(description="Ingest a SAM public entity extract into sam_entities.") + ap.add_argument("dat", type=Path, help="Path to SAM_PUBLIC_MONTHLY_V2_YYYYMMDD.dat") + ap.add_argument("--output", type=Path, default=Path("output") / OUTPUT, help="Output parquet path") + ap.add_argument("--upload", action="store_true", help="Publish the result to R2 (needs R2 creds)") + args = ap.parse_args() + args.output.parent.mkdir(parents=True, exist_ok=True) + out = build_sam_entities_from_extract(args.dat, args.output) + if args.upload: + from spicy_regs.sources import r2 + + r2.upload_file(out, remote_key=OUTPUT) + logger.info("Uploaded {} to R2", OUTPUT) + + +if __name__ == "__main__": + main() diff --git a/tests/test_ingest_sam_extract.py b/tests/test_ingest_sam_extract.py new file mode 100644 index 0000000..d7788dc --- /dev/null +++ b/tests/test_ingest_sam_extract.py @@ -0,0 +1,107 @@ +"""Hermetic test for the SAM public-extract parser (no network).""" + +from __future__ import annotations + +import duckdb + +from spicy_regs.ingest_sam_extract import build_sam_entities_from_extract + + +def _record(**pos: str) -> str: + """Build one 142-field pipe-delimited record with given 0-indexed positions set.""" + fields = [""] * 142 + for k, v in pos.items(): + fields[int(k[1:])] = v # keys like "c0", "c18" + return "|".join(fields) + + +_FIXTURE = "\n".join( + [ + "BOF PUBLIC V2 00000000 20260705 0000002 0000000", + _record( + c0="ABCDEFGHIJ12", + c3="1A2B3", + c5="A", + c6="Z2", + c7="20200115", + c8="20270115", + c11="ACME WIDGETS INC", + c12="ACME", + c17="RANCHO CORDOVA", + c18="CA", + c19="95742", + c22="06", + c26="www.acme.example", + c27="2L", + c32="444110", + ), + _record( + c0="ZYXWVUTSRQ98", + c3="", + c5="E", + c6="Z1", + c7="20150301", + c8="20240301", + c11="OLD CO LLC", + c18="TX", + c19="75001", + c32="", + ), + "EOF PUBLIC V2 00000000 20260705 0000002 0000000", + ] +) + + +def test_extract_parses_and_maps(tmp_path): + src = tmp_path / "SAM_PUBLIC_MONTHLY_V2_20260705.dat" + src.write_text(_FIXTURE + "\n") + out = tmp_path / "sam_entities.parquet" + build_sam_entities_from_extract(src, out) + + rows = duckdb.sql(f"SELECT * FROM read_parquet('{out}') ORDER BY uei").fetchall() + cols = [d[0] for d in duckdb.sql(f"DESCRIBE SELECT * FROM read_parquet('{out}')").fetchall()] + assert cols == [ + "uei", + "cage_code", + "legal_business_name", + "dba_name", + "entity_structure_desc", + "entity_type_desc", + "profit_structure_desc", + "state", + "city", + "zip_code", + "congressional_district", + "primary_naics", + "registration_status", + "registration_date", + "registration_expiration_date", + "exclusion_status_flag", + "purpose_of_registration_desc", + "entity_url", + ] + assert len(rows) == 2 # BOF/EOF sentinel lines dropped + + r = dict(zip(cols, rows[0])) # ABCDEFGHIJ12 + assert r["uei"] == "ABCDEFGHIJ12" + assert r["cage_code"] == "1A2B3" + assert r["legal_business_name"] == "ACME WIDGETS INC" + assert r["dba_name"] == "ACME" + assert r["state"] == "CA" and r["city"] == "RANCHO CORDOVA" and r["zip_code"] == "95742" + assert r["congressional_district"] == "06" + assert r["primary_naics"] == "444110" + assert r["entity_structure_desc"] == "2L" + assert r["entity_url"] == "www.acme.example" + assert r["registration_status"] == "Active" # A -> Active + assert r["purpose_of_registration_desc"] == "All Awards" # Z2 -> mapped + assert r["registration_date"] == "2020-01-15" # YYYYMMDD -> YYYY-MM-DD + assert r["registration_expiration_date"] == "2027-01-15" + # Fields the extract doesn't cleanly provide are NULL. + assert r["entity_type_desc"] is None + assert r["profit_structure_desc"] is None + assert r["exclusion_status_flag"] is None + + r2 = dict(zip(cols, rows[1])) # ZYXWVUTSRQ98 + assert r2["registration_status"] == "Expired" # E -> Expired + assert r2["purpose_of_registration_desc"] == "Federal Assistance Awards" # Z1 + assert r2["cage_code"] is None and r2["primary_naics"] is None # empty -> NULL