initial code
This commit is contained in:
@@ -0,0 +1,234 @@
|
||||
"""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())
|
||||
Reference in New Issue
Block a user