spb/company-atlas

1"""Streamed exports (json / ndjson / csv) with bounded limits and keyset iteration — derived data only, never raw page content."""2from __future__ import annotations34import csv5import io6from collections.abc import AsyncIterator, Callable7from typing import Any89import orjson10from fastapi import APIRouter, HTTPException, Query11from fastapi.responses import StreamingResponse1213from companyatlas.api import queries as q14from companyatlas.api import serializers as ser15from companyatlas.api.common import _default16from companyatlas.db import connection, fetch_all1718ORDER = 2019router = APIRouter(prefix="/api/v1", tags=["exports"])20BATCH = 100021MAX_EVENTS = 10_00022MAX_COMPANIES = 20_00023MAX_JOBS = 10_00024MEDIA = {"json": "application/json", "ndjson": "application/x-ndjson", "csv": "text/csv; charset=utf-8"}25EVENT_COLUMNS = ["id", "detected_at", "company_slug", "company_name", "company_domain", "country", "event_type", "event_subtype", "importance", "confidence",26                 "confidence_label", "title", "summary", "old_value", "new_value", "source_url", "surface", "origin", "status"]27COMPANY_COLUMNS = ["id", "slug", "display_name", "canonical_domain", "website", "country", "hq_city", "industry_primary", "industries", "public_company", "ticker",28                   "tier", "importance", "activity_score", "hiring_momentum_30d", "open_jobs", "ai_adoption", "sensors", "events", "last_event_at"]29JOB_COLUMNS = ["id", "company_slug", "title", "department", "location_text", "city", "country", "remote", "employment_type", "seniority", "url", "posted_at",30               "first_seen_at", "last_seen_at", "removed_at", "status", "is_ai"]313233def _dumps(obj: Any) -> bytes:34    return orjson.dumps(obj, option=orjson.OPT_UTC_Z | orjson.OPT_NON_STR_KEYS, default=_default)353637def _csv_row(values: list[Any]) -> bytes:38    buf = io.StringIO()39    csv.writer(buf, lineterminator="\n").writerow(["" if v is None else (v.isoformat() if hasattr(v, "isoformat") else v) for v in values])40    return buf.getvalue().encode("utf-8")414243def _event_flat(ev: dict[str, Any]) -> list[Any]:44    c = ev["company"]45    return [ev["id"], ev["detected_at"], c["slug"], c["display_name"], c["canonical_domain"], c["country"], ev["event_type"], ev["event_subtype"], ev["importance"],46            ev["confidence"], ev["confidence_label"], ev["title"], ev["summary"], ev["old_value"], ev["new_value"], ev["source_url"], ev["surface"], ev["origin"],47            ev["status"]]484950def _company_flat(c: dict[str, Any]) -> list[Any]:51    m, n = c["metrics"], c["counts"]52    return [c["id"], c["slug"], c["display_name"], c["canonical_domain"], c["website"], c["country"], c["hq_city"], c["industry_primary"], "|".join(c["industries"]),53            c["public_company"], c["ticker"], c["tier"], c["importance"], m.get("activity_score"), m.get("hiring_momentum_30d"), m.get("open_jobs"), m.get("ai_adoption"),54            n["sensors"], n["events"], c["last_event_at"]]555657def _job_flat(j: dict[str, Any]) -> list[Any]:58    return [j["id"], j.get("company_slug"), j["title"], j["department"], j["location_text"], j["city"], j["country"], j["remote"], j["employment_type"], j["seniority"],59            j["url"], j["posted_at"], j["first_seen_at"], j["last_seen_at"], j["removed_at"], j["status"], j["is_ai"]]606162async def _stream(fmt: str, rows: AsyncIterator[dict[str, Any]], columns: list[str], flat: Callable[[dict[str, Any]], list[Any]]) -> AsyncIterator[bytes]:63    if fmt == "csv":64        yield _csv_row(columns)65        async for r in rows:66            yield _csv_row(flat(r))67        return68    if fmt == "ndjson":69        async for r in rows:70            yield _dumps(r) + b"\n"71        return72    yield b"["73    first = True74    async for r in rows:75        yield (b"" if first else b",") + _dumps(r)76        first = False77    yield b"]"787980def _response(fmt: str, name: str, body: AsyncIterator[bytes]) -> StreamingResponse:81    if fmt not in MEDIA:82        raise HTTPException(status_code=404, detail="unsupported format (json, ndjson, csv)")83    return StreamingResponse(body, media_type=MEDIA[fmt], headers={"content-disposition": f'attachment; filename="company-atlas-{name}.{fmt}"',84                                                                  "cache-control": "public, max-age=300", "x-accel-buffering": "no"})858687async def _iter_events(where: list[str], params: dict[str, Any], limit: int) -> AsyncIterator[dict[str, Any]]:88    sent = 089    cursor: tuple[Any, str] | None = None90    async with connection() as conn:91        while sent < limit:92            w = list(where)93            p = dict(params)94            if cursor:95                w.append("(e.detected_at, e.id) < (:cur_at, :cur_id)")96                p.update(cur_at=cursor[0], cur_id=cursor[1])97            rows = await q.fetch_events(conn, w, p, sort="recent", limit=min(BATCH, limit - sent))98            if not rows:99                return100            for r in rows:101                yield ser.event(r)102            sent += len(rows)103            cursor = (rows[-1]["detected_at"], rows[-1]["id"])104105106@router.get("/export/events.{fmt}", summary="Export events (≤ 10 000 rows)")107async def export_events(fmt: str, since: str | None = None, until: str | None = None, event_type: str | None = None, country: str | None = None,108                        industry: str | None = None, company: str | None = None, min_importance: float | None = Query(None, ge=0, le=1),109                        limit: int = Query(MAX_EVENTS, ge=1, le=MAX_EVENTS)) -> StreamingResponse:110    company_id = None111    if company:112        async with connection() as conn:113            company_id = (await q.require_company(conn, company))["id"]114    where, params = q.event_filters(company_id=company_id, event_type=event_type, country=country, industry=industry, since=q.parse_iso(since),115                                    until=q.parse_iso(until, "until"), min_importance=min_importance)116    return _response(fmt, "events", _stream(fmt, _iter_events(where, params, limit), EVENT_COLUMNS, _event_flat))117118119async def _iter_companies(where: list[str], params: dict[str, Any], limit: int) -> AsyncIterator[dict[str, Any]]:120    sent = 0121    last_id = ""122    async with connection() as conn:123        while sent < limit:124            w = list(where) + ["c.id > :cur_id"]125            rows = await fetch_all(conn, f"select c.id from companies c where {' and '.join(w)} order by c.id limit :lim", **params, cur_id=last_id,126                                   lim=min(BATCH, limit - sent))127            if not rows:128                return129            ids = [r["id"] for r in rows]130            for card in await q.fetch_cards_by_ids(conn, ids):131                yield ser.company_card(card)132            sent += len(ids)133            last_id = ids[-1]134135136@router.get("/export/companies.{fmt}", summary="Export companies (≤ 20 000 rows)")137async def export_companies(fmt: str, country: str | None = None, industry: str | None = None, status: str | None = None, tier: int | None = Query(None, ge=1, le=4),138                           limit: int = Query(MAX_COMPANIES, ge=1, le=MAX_COMPANIES)) -> StreamingResponse:139    where, params = q.company_filters(country=country, industry=industry, status=status, tier=tier)140    return _response(fmt, "companies", _stream(fmt, _iter_companies(where or ["true"], params, limit), COMPANY_COLUMNS, _company_flat))141142143async def _iter_jobs(where: list[str], params: dict[str, Any], limit: int) -> AsyncIterator[dict[str, Any]]:144    sent = 0145    last_id = ""146    async with connection() as conn:147        while sent < limit:148            w = list(where) + ["j.id > :cur_id"]149            rows = await fetch_all(conn, f"select j.*, c.slug as company_slug from jobs j join companies c on c.id = j.company_id where {' and '.join(w)} "150                                         "order by j.id limit :lim", **params, cur_id=last_id, lim=min(BATCH, limit - sent))151            if not rows:152                return153            for r in rows:154                item = ser.job(r)155                item["company_slug"] = r["company_slug"]156                yield item157            sent += len(rows)158            last_id = rows[-1]["id"]159160161@router.get("/export/jobs.{fmt}", summary="Export jobs (≤ 10 000 rows)")162async def export_jobs(fmt: str, company: str | None = None, since: str | None = None, status: str = Query("open", pattern="^(open|removed|all)$"),163                      country: str | None = None, ai: str | None = None, limit: int = Query(MAX_JOBS, ge=1, le=MAX_JOBS)) -> StreamingResponse:164    where: list[str] = ["true"]165    params: dict[str, Any] = {}166    if company:167        async with connection() as conn:168            params["cid"] = (await q.require_company(conn, company))["id"]169        where.append("j.company_id = :cid")170    since_dt = q.parse_iso(since)171    if since_dt:172        where.append("j.first_seen_at >= :since")173        params["since"] = since_dt174    if status == "open":175        where.append("j.status = 'open'")176    elif status == "removed":177        where.append("j.status <> 'open'")178    if country:179        where.append("j.country = cast(:country as char(2))")180        params["country"] = country.upper()[:2]181    if q.parse_bool(ai):182        where.append("j.is_ai")183    return _response(fmt, "jobs", _stream(fmt, _iter_jobs(where, params, limit), JOB_COLUMNS, _job_flat))184