spb/company-atlas

1"""Reusable, parameterised SQL for the public/admin API.23Rules: user input is always bound (`:param`), never interpolated; sort keys go through whitelists; every list is bounded; the4company card is built in two phases (page of ids with an indexed ORDER BY, then one card query for those ids) so the per-company5lateral aggregates only run for the rows that are returned.6"""7from __future__ import annotations89from datetime import UTC, date, datetime, timedelta10from typing import Any1112from fastapi import HTTPException13from sqlalchemy.ext.asyncio import AsyncConnection1415from companyatlas.db import fetch_all, fetch_one, fetch_val16from companyatlas.taxonomy import EventType1718# ------------------------------------------------------------------------------------------------ windows / params1920WINDOWS: dict[str, timedelta] = {"24h": timedelta(hours=24), "7d": timedelta(days=7), "30d": timedelta(days=30), "90d": timedelta(days=90),21                                 "1y": timedelta(days=365)}22COUNT_CAP = 10_000          # `total` is exact up to this bound (keeps deep pagination cheap on the big tables)23MAX_LIMIT = 500242526def now_utc() -> datetime:27    return datetime.now(UTC)282930def window_start(window: str, default: str = "7d") -> datetime:31    return now_utc() - WINDOWS.get(window, WINDOWS[default])323334def days_ago(days: int) -> datetime:35    return now_utc() - timedelta(days=days)363738def parse_iso(value: str | None, name: str = "since") -> datetime | None:39    if not value:40        return None41    v = value.strip()42    if v.endswith(("Z", "z")):43        v = v[:-1] + "+00:00"44    try:45        dt = datetime.fromisoformat(v)46    except ValueError:47        try:48            dt = datetime.combine(date.fromisoformat(v), datetime.min.time())49        except ValueError as exc:50            raise HTTPException(status_code=422, detail=f"{name}: invalid ISO-8601 timestamp") from exc51    if dt.tzinfo is None:52        dt = dt.replace(tzinfo=UTC)53    return dt.astimezone(UTC)545556def parse_bool(value: str | bool | None) -> bool | None:57    if value is None or value == "":58        return None59    if isinstance(value, bool):60        return value61    return value.strip().lower() in ("1", "true", "yes", "on")626364def clamp(value: int, lo: int, hi: int) -> int:65    return max(lo, min(hi, value))666768def csv_list(value: str | None, *, maxlen: int = 20) -> list[str]:69    if not value:70        return []71    return [x.strip() for x in value.split(",") if x.strip()][:maxlen]727374# ------------------------------------------------------------------------------------------------ companies757677async def get_company(conn: AsyncConnection, key: str) -> dict[str, Any] | None:78    key = (key or "").strip()79    if not key or len(key) > 200:80        return None81    return await fetch_one(conn, "select * from companies where slug = :k or id = :k limit 1", k=key)828384async def require_company(conn: AsyncConnection, key: str) -> dict[str, Any]:85    row = await get_company(conn, key)86    if row is None:87        raise HTTPException(status_code=404, detail="company not found")88    return row899091COMPANY_CARD_SQL = """92select c.id, c.slug, c.display_name, c.legal_name, c.canonical_domain, c.website, c.description, c.industries, c.industry_primary,93       c.country, c.hq_city, c.hq_region, c.public_company, c.ticker, c.exchange, c.founded_year, c.employees_band, c.logo_url,94       c.status, c.onboarding_status, c.onboarding_error, c.importance, c.tier, c.indexed, c.stats, c.last_event_at, c.last_observed_at,95       c.first_observed_at, c.discovered_at, c.created_at, c.updated_at, c.source_meta->'profile' as profile,96       mx.metrics, sc.sensors, sc.observations, sc.changes, sc.events, jc.jobs_open{sparkline_col}97from companies c98left join lateral (select jsonb_object_agg(m.metric, m.value) as metrics from metrics_current m where m.company_id = c.id) mx on true99left join lateral (select count(*) filter (where s.status <> 'retired') as sensors, coalesce(sum(s.observation_count), 0) as observations,100                          coalesce(sum(s.change_count), 0) as changes, coalesce(sum(s.event_count), 0) as events101                   from sensors s where s.company_id = c.id) sc on true102left join lateral (select count(*) as jobs_open from jobs j where j.company_id = c.id and j.status = 'open') jc on true103{sparkline_join}104where c.id = any(cast(:ids as text[]))105"""106SPARKLINE_COL = ", sp.sparkline"107SPARKLINE_JOIN = ("left join lateral (select array_agg(x.value order by x.day) as sparkline from (select day, value from metric_series ms "108                  "where ms.company_id = c.id and ms.metric = 'activity_score' and ms.day >= :spark_from order by day desc limit 30) x) sp on true")109110COMPANY_SORTS: dict[str, str] = {111    "activity": "ma.value desc nulls last, c.importance desc, c.id",112    "events": "c.last_event_at desc nulls last, c.importance desc, c.id",113    "hiring": "mh.value desc nulls last, c.importance desc, c.id",114    "name": "c.display_name asc, c.id",115    "importance": "c.importance desc, c.display_name asc, c.id",116    "recent": "c.discovered_at desc, c.id",117    # only valid together with a `q` filter (binds :q); list endpoints switch to it automatically when q is present118    "relevance": ("greatest(coalesce(ts_rank(c.search, websearch_to_tsquery('simple', :q)), 0), similarity(c.display_name, :q), "119                  "similarity(c.canonical_domain, :q)) desc, c.importance desc, c.id"),120}121COMPANY_SORT_JOINS: dict[str, str] = {122    "activity": "left join metrics_current ma on ma.company_id = c.id and ma.metric = 'activity_score'",123    "hiring": "left join metrics_current mh on mh.company_id = c.id and mh.metric = 'hiring_momentum_30d'",124}125126127async def fetch_cards_by_ids(conn: AsyncConnection, ids: list[str], *, sparkline: bool = False) -> list[dict[str, Any]]:128    """Card rows for the given ids, in the given order."""129    if not ids:130        return []131    sql = COMPANY_CARD_SQL.format(sparkline_col=SPARKLINE_COL if sparkline else "", sparkline_join=SPARKLINE_JOIN if sparkline else "")132    params: dict[str, Any] = {"ids": ids}133    if sparkline:134        params["spark_from"] = (now_utc() - timedelta(days=30)).date()135    rows = await fetch_all(conn, sql, **params)136    by_id = {r["id"]: r for r in rows}137    return [by_id[i] for i in ids if i in by_id]138139140async def company_page_ids(conn: AsyncConnection, where: list[str], params: dict[str, Any], *, sort: str = "activity", limit: int = 25,141                           offset: int = 0, extra_join: str = "") -> tuple[list[str], int]:142    """Phase 1 of a company list: ids for the page + bounded total."""143    sort = sort if sort in COMPANY_SORTS else "activity"144    if sort == "relevance" and "q" not in params:145        sort = "activity"146    elif sort == "activity" and params.get("q"):147        sort = "relevance"148    join = " ".join(x for x in (COMPANY_SORT_JOINS.get(sort, ""), extra_join) if x)149    where_sql = (" where " + " and ".join(where)) if where else ""150    rows = await fetch_all(conn, f"select c.id from companies c {join}{where_sql} order by {COMPANY_SORTS[sort]} limit :limit offset :offset",151                           **params, limit=limit, offset=offset)152    total = await bounded_count(conn, f"from companies c {extra_join}{where_sql}", params)153    return [r["id"] for r in rows], total154155156async def bounded_count(conn: AsyncConnection, from_where_sql: str, params: dict[str, Any], cap: int = COUNT_CAP) -> int:157    val = await fetch_val(conn, f"select count(*) from (select 1 {from_where_sql} limit :cap) t", **params, cap=cap)158    return int(val or 0)159160161def company_filters(*, q: str | None = None, country: str | None = None, industry: str | None = None, tier: int | None = None,162                    public: bool | None = None, status: str | None = None, has_events: bool | None = None,163                    onboarding_status: str | None = None) -> tuple[list[str], dict[str, Any]]:164    where: list[str] = []165    params: dict[str, Any] = {}166    if q:167        q = q.strip()[:200]168        params["q"] = q169        if len(q) >= 3:170            where.append("(c.search @@ websearch_to_tsquery('simple', :q) or c.display_name % :q or c.canonical_domain % :q)")171        else:172            where.append("(c.display_name ilike :q_prefix or c.canonical_domain ilike :q_prefix)")173            params["q_prefix"] = q.replace("%", "").replace("_", "") + "%"174    if country:175        where.append("c.country = cast(:country as char(2))")176        params["country"] = country.strip().upper()[:2]177    if industry:178        where.append("cast(:industry as text) = any(c.industries)")179        params["industry"] = industry.strip()[:80]180    if tier is not None:181        where.append("c.tier = :tier")182        params["tier"] = tier183    if public is not None:184        where.append("c.public_company = :public")185        params["public"] = public186    if status:187        where.append("c.status = cast(:status as text)")188        params["status"] = status.strip().upper()[:40]189    if has_events is True:190        where.append("c.last_event_at is not null")191    elif has_events is False:192        where.append("c.last_event_at is null")193    if onboarding_status:194        where.append("c.onboarding_status = cast(:onboarding_status as text)")195        params["onboarding_status"] = onboarding_status.strip().lower()[:40]196    return where, params197198199# ------------------------------------------------------------------------------------------------ events200201EVENT_SELECT = """202select e.id, e.company_id, e.sensor_id, e.change_id, e.cluster_id, e.surface, e.event_type, e.event_subtype, e.importance, e.confidence,203       e.confidence_label, e.title, e.summary, e.old_value, e.new_value, e.payload, e.entities, e.tags, e.detected_at, e.effective_at,204       e.published_at, e.source_url, e.snapshot_before, e.snapshot_after, e.language, e.origin, e.model_name, e.prompt_version, e.status,205       e.retracted_reason,206       c.slug as company_slug, c.display_name as company_display_name, c.canonical_domain as company_domain, c.country as company_country,207       c.logo_url as company_logo_url208from events e join companies c on c.id = e.company_id209"""210EVENT_SORTS: dict[str, str] = {"recent": "e.detected_at desc, e.id desc", "importance": "e.importance desc, e.detected_at desc, e.id desc"}211_EVENT_TYPES = {t.value for t in EventType}212213214def event_filters(*, company_id: str | None = None, event_type: str | None = None, event_subtype: str | None = None, country: str | None = None,215                  industry: str | None = None, since: datetime | None = None, until: datetime | None = None, min_importance: float | None = None,216                  min_confidence: float | None = None, q: str | None = None, surface: str | None = None, origin: str | None = None,217                  status: str | None = "active", event_types: list[str] | None = None,218                  event_subtypes: list[str] | None = None) -> tuple[list[str], dict[str, Any]]:219    where: list[str] = []220    params: dict[str, Any] = {}221    if status:222        where.append("e.status = cast(:e_status as text)")223        params["e_status"] = status224    if company_id:225        where.append("e.company_id = :e_company_id")226        params["e_company_id"] = company_id227    if event_type:228        where.append("e.event_type = cast(:e_type as text)")229        params["e_type"] = event_type.strip().upper()[:40]230    if event_types:231        where.append("e.event_type = any(cast(:e_types as text[]))")232        params["e_types"] = [t.upper()[:40] for t in event_types][:30]233    if event_subtype:234        where.append("e.event_subtype = cast(:e_subtype as text)")235        params["e_subtype"] = event_subtype.strip().upper()[:60]236    if event_subtypes:237        where.append("e.event_subtype = any(cast(:e_subtypes as text[]))")238        params["e_subtypes"] = [t.upper()[:60] for t in event_subtypes][:60]239    if country:240        where.append("c.country = cast(:e_country as char(2))")241        params["e_country"] = country.strip().upper()[:2]242    if industry:243        where.append("cast(:e_industry as text) = any(c.industries)")244        params["e_industry"] = industry.strip()[:80]245    if since is not None:246        where.append("e.detected_at >= :e_since")247        params["e_since"] = since248    if until is not None:249        where.append("e.detected_at < :e_until")250        params["e_until"] = until251    if min_importance is not None:252        where.append("e.importance >= :e_min_imp")253        params["e_min_imp"] = float(min_importance)254    if min_confidence is not None:255        where.append("e.confidence >= :e_min_conf")256        params["e_min_conf"] = float(min_confidence)257    if q:258        where.append("e.search @@ websearch_to_tsquery('english', :e_q)")259        params["e_q"] = q.strip()[:200]260    if surface:261        where.append("e.surface = cast(:e_surface as text)")262        params["e_surface"] = surface.strip().lower()[:40]263    if origin:264        where.append("e.origin = cast(:e_origin as text)")265        params["e_origin"] = origin.strip().lower()[:20]266    return where, params267268269async def fetch_events(conn: AsyncConnection, where: list[str], params: dict[str, Any], *, sort: str = "recent", limit: int = 50,270                       offset: int = 0) -> list[dict[str, Any]]:271    order = EVENT_SORTS.get(sort, EVENT_SORTS["recent"])272    where_sql = (" where " + " and ".join(where)) if where else ""273    return await fetch_all(conn, f"{EVENT_SELECT}{where_sql} order by {order} limit :limit offset :offset", **params,274                           limit=clamp(limit, 1, MAX_LIMIT), offset=max(0, offset))275276277async def count_events(conn: AsyncConnection, where: list[str], params: dict[str, Any]) -> int:278    where_sql = (" where " + " and ".join(where)) if where else ""279    return await bounded_count(conn, f"from events e join companies c on c.id = e.company_id{where_sql}", params)280281282async def fetch_event(conn: AsyncConnection, event_id: str) -> dict[str, Any] | None:283    return await fetch_one(conn, f"{EVENT_SELECT} where e.id = :id", id=event_id)284285286# ------------------------------------------------------------------------------------------------ metrics helpers287288289async def metric_series(conn: AsyncConnection, company_id: str, metric: str, days: int) -> list[dict[str, Any]]:290    return await fetch_all(conn, "select day, value, confidence from metric_series where company_id = :cid and metric = :m and day >= :d "291                                 "order by day asc limit :lim", cid=company_id, m=metric, d=days_ago(days).date(), lim=clamp(days, 1, 2000))292293294async def metric_values_at(conn: AsyncConnection, ids: list[str], metric: str, on_or_before: date) -> dict[str, float]:295    """Latest series value on/before a day per company (used for window deltas)."""296    if not ids:297        return {}298    rows = await fetch_all(conn, "select distinct on (company_id) company_id, value from metric_series where metric = :m and "299                                 "company_id = any(cast(:ids as text[])) and day <= :d order by company_id, day desc", m=metric, ids=ids, d=on_or_before)300    return {r["company_id"]: float(r["value"]) for r in rows}301302303async def settings_value(conn: AsyncConnection, key: str) -> Any:304    return await fetch_val(conn, "select value from settings_kv where key = :k", k=key)305306307__all__ = ["COMPANY_SORTS", "COUNT_CAP", "EVENT_SORTS", "MAX_LIMIT", "WINDOWS", "bounded_count", "clamp", "company_filters", "company_page_ids",308           "count_events", "csv_list", "days_ago", "event_filters", "fetch_cards_by_ids", "fetch_event", "fetch_events", "get_company",309           "metric_series", "metric_values_at", "now_utc", "parse_bool", "parse_iso", "require_company", "settings_value", "window_start"]310