235 lines
8.5 KiB
Python
235 lines
8.5 KiB
Python
"""Export complete histories for a private development-only vehicle sample.
|
|
|
|
The source sample is page-based and therefore not population-representative.
|
|
Use it to develop and test the episode/model pipeline, never for final rates.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import gzip
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import sys
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import psycopg2
|
|
|
|
try:
|
|
from . import export_inspection_batch as secure_export
|
|
except ImportError: # Direct execution: python scripts/export_history_sample.py
|
|
import export_inspection_batch as secure_export
|
|
|
|
|
|
QUERY_VERSION = "inspection_history_development_sample_v3"
|
|
|
|
|
|
def parse_timestamp(value: str) -> datetime:
|
|
try:
|
|
return datetime.fromisoformat(value)
|
|
except ValueError as exc:
|
|
raise argparse.ArgumentTypeError(f"Invalid ISO timestamp: {value}") from exc
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--vehicles", type=int, default=10_000)
|
|
parser.add_argument("--sample-percent", type=float, default=0.5)
|
|
parser.add_argument("--seed", type=int, default=20260715)
|
|
parser.add_argument("--start", type=parse_timestamp, default=datetime(2010, 1, 1))
|
|
parser.add_argument("--end", type=parse_timestamp, default=datetime(2026, 7, 1))
|
|
parser.add_argument("--fetch-size", type=int, default=5_000)
|
|
parser.add_argument("--output", type=Path)
|
|
parser.add_argument("--overwrite", action="store_true")
|
|
return parser.parse_args()
|
|
|
|
|
|
def validate_args(args: argparse.Namespace) -> Path:
|
|
if not 100 <= args.vehicles <= 100_000:
|
|
raise ValueError("--vehicles must be between 100 and 100000")
|
|
if not 0 < args.sample_percent <= 5:
|
|
raise ValueError("--sample-percent must be greater than 0 and at most 5")
|
|
if not 100 <= args.fetch_size <= 50_000:
|
|
raise ValueError("--fetch-size must be between 100 and 50000")
|
|
if args.end <= args.start:
|
|
raise ValueError("--end must be later than --start")
|
|
|
|
output = args.output
|
|
if output is None:
|
|
output = (
|
|
secure_export.PRIVATE_DATA_ROOT
|
|
/ "development"
|
|
/ f"history_sample_{args.vehicles}.csv.gz"
|
|
)
|
|
output = output.resolve()
|
|
if not output.is_relative_to(secure_export.PRIVATE_DATA_ROOT):
|
|
raise ValueError(
|
|
f"Development histories must stay under {secure_export.PRIVATE_DATA_ROOT}"
|
|
)
|
|
manifest = output.with_name(output.name + ".manifest.json")
|
|
if (output.exists() or manifest.exists()) and not args.overwrite:
|
|
raise FileExistsError(
|
|
f"Refusing to overwrite {output} or its manifest; use --overwrite"
|
|
)
|
|
return output
|
|
|
|
|
|
def make_sql(sample_percent: float, seed: int) -> str:
|
|
safe_percent = f"{sample_percent:.6f}"
|
|
safe_seed = int(seed)
|
|
return f"""
|
|
WITH candidate_vins AS MATERIALIZED (
|
|
SELECT DISTINCT vin
|
|
FROM data.inspection_search
|
|
TABLESAMPLE SYSTEM ({safe_percent}) REPEATABLE ({safe_seed})
|
|
WHERE vin IS NOT NULL
|
|
AND length(btrim(vin)) = 17
|
|
),
|
|
sampled_vins AS MATERIALIZED (
|
|
SELECT vin
|
|
FROM candidate_vins
|
|
ORDER BY md5(vin)
|
|
LIMIT %(vehicle_limit)s
|
|
)
|
|
SELECT
|
|
s.id AS internal_event_id,
|
|
s.vin,
|
|
s.test_start AS event_ts,
|
|
lower(btrim(s.county)) AS source_era,
|
|
CASE
|
|
WHEN lower(btrim(s.county)) IN ('slc', 'slco') THEN 'salt_lake'
|
|
ELSE lower(btrim(s.county))
|
|
END AS public_county,
|
|
{secure_export.OUTCOME_SELECT_SQL},
|
|
lower(nullif(btrim(s.program_type), '')) AS program_type,
|
|
upper(nullif(btrim(s.test_type), '')) AS test_type,
|
|
upper(nullif(btrim(v.make), '')) AS observed_make,
|
|
upper(nullif(btrim(v.model), '')) AS observed_model,
|
|
v.year AS observed_model_year
|
|
FROM data.inspection_search AS s
|
|
JOIN data.inspection_vehicle AS v USING (id)
|
|
JOIN sampled_vins ON sampled_vins.vin = s.vin
|
|
WHERE s.test_start >= %(start_ts)s
|
|
AND s.test_start < %(end_ts)s
|
|
ORDER BY s.test_start, s.id
|
|
"""
|
|
|
|
|
|
def file_sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
os.umask(0o077)
|
|
try:
|
|
output = validate_args(args)
|
|
key, key_version, key_fingerprint = secure_export.get_hmac_configuration()
|
|
except (FileExistsError, ValueError) as exc:
|
|
print(f"Configuration error: {exc}", file=sys.stderr)
|
|
return 2
|
|
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
manifest = output.with_name(output.name + ".manifest.json")
|
|
partial = output.with_name(output.name + ".partial")
|
|
partial_manifest = manifest.with_name(manifest.name + ".partial")
|
|
sql = make_sql(args.sample_percent, args.seed)
|
|
|
|
try:
|
|
connection = secure_export.connect_read_only()
|
|
except (ValueError, RuntimeError, psycopg2.Error) as exc:
|
|
print(f"Secure read-only connection failed: {type(exc).__name__}", file=sys.stderr)
|
|
return 1
|
|
|
|
source_rows_read = 0
|
|
rows_written = 0
|
|
invalid_vins_skipped = 0
|
|
last_order_key: Optional[tuple[object, object]] = None
|
|
try:
|
|
with gzip.open(partial, "wt", encoding="utf-8", newline="") as handle:
|
|
writer = csv.writer(handle)
|
|
writer.writerow(secure_export.OUTPUT_COLUMNS)
|
|
with connection.cursor(name="uvh_history_sample") as cursor:
|
|
cursor.itersize = args.fetch_size
|
|
cursor.execute(
|
|
sql,
|
|
{
|
|
"vehicle_limit": args.vehicles,
|
|
"start_ts": args.start,
|
|
"end_ts": args.end,
|
|
},
|
|
)
|
|
while True:
|
|
rows = cursor.fetchmany(args.fetch_size)
|
|
if not rows:
|
|
break
|
|
for source_row in rows:
|
|
source_rows_read += 1
|
|
order_key = (source_row[2], source_row[0])
|
|
if last_order_key is not None and order_key <= last_order_key:
|
|
raise RuntimeError(
|
|
"Source join did not return a unique increasing "
|
|
"event_ts/internal_event_id order"
|
|
)
|
|
last_order_key = order_key
|
|
transformed = secure_export.private_row(
|
|
source_row, key, key_version
|
|
)
|
|
if transformed is None:
|
|
invalid_vins_skipped += 1
|
|
continue
|
|
writer.writerow(transformed)
|
|
rows_written += 1
|
|
connection.rollback()
|
|
except Exception:
|
|
partial.unlink(missing_ok=True)
|
|
partial_manifest.unlink(missing_ok=True)
|
|
raise
|
|
finally:
|
|
connection.close()
|
|
|
|
manifest_data = {
|
|
"export_kind": "vehicle_history_development_sample",
|
|
"query_version": QUERY_VERSION,
|
|
"query_sha256": hashlib.sha256(sql.encode("utf-8")).hexdigest(),
|
|
"generated_at_utc": datetime.now(timezone.utc).isoformat(),
|
|
"source_start_inclusive": args.start.isoformat(),
|
|
"source_end_exclusive": args.end.isoformat(),
|
|
"vehicle_sample_limit": args.vehicles,
|
|
"sample_method": "page_sample_candidates_then_md5_ordered_distinct_vins",
|
|
"sample_percent": args.sample_percent,
|
|
"sample_seed": args.seed,
|
|
"population_estimate_allowed": False,
|
|
"vehicle_token_key_version": key_version,
|
|
"vehicle_token_key_fingerprint": key_fingerprint,
|
|
"source_rows_read": source_rows_read,
|
|
"rows": rows_written,
|
|
"invalid_vins_skipped": invalid_vins_skipped,
|
|
"columns": secure_export.OUTPUT_COLUMNS,
|
|
"compressed_file_sha256": file_sha256(partial),
|
|
"compressed_file_bytes": partial.stat().st_size,
|
|
"classification": "private_pseudonymized_development_only",
|
|
"source_transaction": "read_only_repeatable_read",
|
|
}
|
|
partial_manifest.write_text(
|
|
json.dumps(manifest_data, indent=2) + "\n", encoding="utf-8"
|
|
)
|
|
os.replace(partial, output)
|
|
os.replace(partial_manifest, manifest)
|
|
print(f"Exported {rows_written:,} rows for a development vehicle sample")
|
|
print(f"Private output: {output}")
|
|
print("This sample is not valid for population-rate estimates.")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|