"""Idempotent seed loader: registry files → `industries`, `countries`, `companies` (+ aliases, domains, relationships, discover queue). Contract with the crawl core (docs/ARCHITECTURE.md, "Seeds → Crawl"): new companies get `onboarding_status='pending'` and a `queue_jobs(kind='discover', key='discover:')` row. Re-running never overwrites a non-null value except `importance`, `tier` and the provenance in `source_meta` (spec: factual provenance, historical-first). """ from __future__ import annotations import csv import json import logging from collections.abc import Iterable from datetime import UTC, datetime from pathlib import Path from typing import Any from sqlalchemy.ext.asyncio import AsyncConnection from companyatlas import ids from companyatlas.db import execute, execute_many, fetch_all, jsonb from companyatlas.registry.industries import REGISTRY_DIR, is_valid_slug, load_industries from companyatlas.urls import registrable_domain log = logging.getLogger(__name__) COMPANIES_DIR = REGISTRY_DIR / "companies" COUNTRIES_FILE = REGISTRY_DIR / "countries.csv" DEFAULT_IMPORTANCE = 0.2 DEFAULT_TIER = 4 GENERIC_TLD_NAMES = {"www", "web", "site", "home", "online", "official"} # ------------------------------------------------------------------------------------------------------------ registry files def registry_files() -> list[Path]: return sorted(COMPANIES_DIR.glob("*.ndjson")) def load_registry_rows(files: Iterable[Path] | None = None) -> list[dict[str, Any]]: rows: list[dict[str, Any]] = [] for path in files or registry_files(): with Path(path).open(encoding="utf-8") as f: for n, line in enumerate(f, 1): line = line.strip() if not line: continue try: rows.append(json.loads(line)) except json.JSONDecodeError as e: raise ValueError(f"{path}:{n}: invalid JSON ({e})") from e return rows def load_countries_file(path: Path = COUNTRIES_FILE) -> list[dict[str, Any]]: with path.open(encoding="utf-8") as f: out = [] for r in csv.DictReader(f): out.append({"code": r["code"].strip().upper(), "name": r["name"], "region": r.get("region") or None, "subregion": r.get("subregion") or None, "lat": float(r["lat"]) if r.get("lat") else None, "lon": float(r["lon"]) if r.get("lon") else None}) return out def read_rows_file(path: Path) -> list[dict[str, Any]]: """Manual import: `.ndjson`/`.jsonl`/`.json` (list) or `.csv` with a header (`website` required).""" path = Path(path) if path.suffix.lower() == ".csv": with path.open(encoding="utf-8-sig") as f: rows = [] for r in csv.DictReader(f): row = {k.strip(): (v.strip() if isinstance(v, str) else v) for k, v in r.items() if k} for key in ("industries", "aliases"): if isinstance(row.get(key), str): row[key] = [x.strip() for x in row[key].replace(";", ",").split(",") if x.strip()] rows.append(row) return rows if path.suffix.lower() == ".json": data = json.loads(path.read_text(encoding="utf-8")) return list(data if isinstance(data, list) else data.get("companies", [])) return load_registry_rows([path]) # ------------------------------------------------------------------------------------------------------------ normalisation def name_from_domain(domain: str) -> str: label = domain.split(".")[0] if label in GENERIC_TLD_NAMES and domain.count(".") >= 2: label = domain.split(".")[1] return label.replace("-", " ").replace("_", " ").title() def normalise_row(row: dict[str, Any], *, source: str) -> dict[str, Any] | None: website = (row.get("website") or "").strip() if not website: return None if "://" not in website: website = "https://" + website if not website.startswith(("http://", "https://")): return None domain = (row.get("canonical_domain") or "").strip().lower() or registrable_domain(website) if not domain or "." not in domain: return None display = (row.get("display_name") or "").strip() or name_from_domain(domain) industries = [s for s in (row.get("industries") or []) if isinstance(s, str) and is_valid_slug(s)] if row.get("industry") and is_valid_slug(row["industry"]) and row["industry"] not in industries: industries.insert(0, row["industry"]) country = (row.get("country") or "").strip().upper() or None parent = row.get("parent") if isinstance(row.get("parent"), dict) else None try: importance = float(row.get("importance", DEFAULT_IMPORTANCE)) except (TypeError, ValueError): importance = DEFAULT_IMPORTANCE try: tier = int(row.get("tier", DEFAULT_TIER)) except (TypeError, ValueError): tier = DEFAULT_TIER meta = { "source": row.get("source") or source, "harvested_at": row.get("harvested_at"), "sitelinks": row.get("sitelinks"), "industry_labels": row.get("industry_labels") or [], "parent_wikidata_id": parent.get("wikidata_id") if parent else None, "parent_name": parent.get("name") if parent else None, "notes": row.get("notes") or [], "seeded_at": datetime.now(UTC).replace(microsecond=0).isoformat(), } meta = {k: v for k, v in meta.items() if v not in (None, [], "")} return { "wikidata_id": (row.get("wikidata_id") or None), "display_name": display[:200], "legal_name": (row.get("legal_name") or None), "aliases": [a for a in (row.get("aliases") or []) if isinstance(a, str) and a.strip()][:12], "website": website, "canonical_domain": domain, "country": country, "hq_city": row.get("hq_city") or None, "hq_region": row.get("hq_region") or None, "industries": industries, "industry_primary": industries[0] if industries else None, "founded_year": row.get("founded_year"), "employees": row.get("employees"), "public_company": bool(row.get("public_company")), "ticker": row.get("ticker") or None, "exchange": row.get("exchange") or None, "lei": row.get("lei") or None, "sec_cik": str(row["sec_cik"]) if row.get("sec_cik") else None, "logo_url": row.get("logo_url") or None, "description": (row.get("description") or None), "importance": max(0.0, min(1.0, importance)), "tier": min(4, max(1, tier)), "company_type": "public" if row.get("public_company") else None, "source_meta": meta, "parent_wikidata_id": parent.get("wikidata_id") if parent else None, } def unique_slug(base: str, country: str | None, taken: set[str]) -> str: if base not in taken: return base if country and f"{base}-{country.lower()}" not in taken: return f"{base}-{country.lower()}" n = 2 while f"{base}-{n}" in taken: n += 1 return f"{base}-{n}" # ------------------------------------------------------------------------------------------------------------ reference tables async def upsert_industries(conn: AsyncConnection) -> int: rows = [{"slug": i.slug, "name": i.name, "parent": i.parent, "description": i.description, "keywords": list(i.keywords), "sort_order": i.sort_order} for i in load_industries()] # Parents first (FK on parent_slug); insert parent_slug in a second pass so ordering inside the file does not matter. await execute_many(conn, """ insert into industries (slug, name, description, keywords, sort_order) values (:slug, :name, :description, cast(:keywords as text[]), :sort_order) on conflict (slug) do update set name = excluded.name, description = excluded.description, keywords = excluded.keywords, sort_order = excluded.sort_order""", rows) await execute_many(conn, "update industries set parent_slug = :parent where slug = :slug", [r for r in rows if r["parent"]]) return len(rows) async def upsert_countries(conn: AsyncConnection) -> int: rows = load_countries_file() await execute_many(conn, """ insert into countries (code, name, region, subregion, lat, lon) values (:code, :name, :region, :subregion, :lat, :lon) on conflict (code) do update set name = excluded.name, region = excluded.region, subregion = excluded.subregion, lat = excluded.lat, lon = excluded.lon""", rows) return len(rows) # ------------------------------------------------------------------------------------------------------------ companies class _Index: """In-memory view of existing companies for one loader run (avoids a lookup per row).""" def __init__(self, by_wikidata: dict[str, str], by_domain: dict[str, str], slugs: set[str], countries: set[str]) -> None: self.by_wikidata = by_wikidata self.by_domain = by_domain self.slugs = slugs self.countries = countries @classmethod async def load(cls, conn: AsyncConnection) -> _Index: rows = await fetch_all(conn, "select id, slug, canonical_domain, wikidata_id from companies") countries = {r["code"] for r in await fetch_all(conn, "select code from countries")} return cls({r["wikidata_id"]: r["id"] for r in rows if r["wikidata_id"]}, {r["canonical_domain"]: r["id"] for r in rows}, {r["slug"] for r in rows}, countries) def _alias_rows(company_id: str, row: dict[str, Any], source: str) -> list[dict[str, Any]]: seen: set[str] = set() out: list[dict[str, Any]] = [] for alias, kind in ([(row["display_name"], "brand"), (row.get("legal_name"), "legal"), (row.get("ticker"), "ticker")] + [(a, "alias") for a in row.get("aliases", [])]): if not alias: continue norm = ids.normalize_alias(alias) if not norm or norm in seen: continue seen.add(norm) out.append({"company_id": company_id, "alias": alias[:200], "alias_norm": norm[:200], "kind": kind, "source": source}) return out async def _insert_company(conn: AsyncConnection, row: dict[str, Any], idx: _Index) -> str: company_id = ids.new_id("company") slug = unique_slug(ids.slugify(row["display_name"]), row["country"], idx.slugs) idx.slugs.add(slug) await execute(conn, """ insert into companies (id, slug, legal_name, display_name, canonical_domain, website, description, industries, industry_primary, country, hq_city, hq_region, founded_year, company_type, public_company, ticker, exchange, employees, wikidata_id, lei, sec_cik, logo_url, onboarding_status, importance, tier, source_meta) values (:id, :slug, :legal_name, :display_name, :canonical_domain, :website, :description, cast(:industries as text[]), :industry_primary, :country, :hq_city, :hq_region, :founded_year, :company_type, :public_company, :ticker, :exchange, :employees, :wikidata_id, :lei, :sec_cik, :logo_url, 'pending', :importance, :tier, cast(:source_meta as jsonb))""", id=company_id, slug=slug, legal_name=row["legal_name"], display_name=row["display_name"], canonical_domain=row["canonical_domain"], website=row["website"], description=row["description"], industries=row["industries"], industry_primary=row["industry_primary"], country=row["country"], hq_city=row["hq_city"], hq_region=row["hq_region"], founded_year=row["founded_year"], company_type=row["company_type"], public_company=row["public_company"], ticker=row["ticker"], exchange=row["exchange"], employees=row["employees"], wikidata_id=row["wikidata_id"], lei=row["lei"], sec_cik=row["sec_cik"], logo_url=row["logo_url"], importance=row["importance"], tier=row["tier"], source_meta=jsonb(row["source_meta"])) idx.by_domain[row["canonical_domain"]] = company_id if row["wikidata_id"]: idx.by_wikidata[row["wikidata_id"]] = company_id return company_id async def _update_company(conn: AsyncConnection, company_id: str, row: dict[str, Any]) -> None: """Fill only null columns; refresh importance/tier and merge provenance. Never touches names, website or history.""" await execute(conn, """ update companies set legal_name = coalesce(legal_name, :legal_name), description = coalesce(description, :description), industries = case when cardinality(industries) = 0 then cast(:industries as text[]) else industries end, industry_primary = coalesce(industry_primary, :industry_primary), country = coalesce(country, :country), hq_city = coalesce(hq_city, :hq_city), hq_region = coalesce(hq_region, :hq_region), founded_year = coalesce(founded_year, :founded_year), company_type = coalesce(company_type, :company_type), public_company = public_company or :public_company, ticker = coalesce(ticker, :ticker), exchange = coalesce(exchange, :exchange), employees = coalesce(employees, :employees), wikidata_id = coalesce(wikidata_id, :wikidata_id), lei = coalesce(lei, :lei), sec_cik = coalesce(sec_cik, :sec_cik), logo_url = coalesce(logo_url, :logo_url), importance = :importance, tier = :tier, source_meta = source_meta || cast(:source_meta as jsonb), updated_at = now() where id = :id""", id=company_id, legal_name=row["legal_name"], description=row["description"], industries=row["industries"], industry_primary=row["industry_primary"], country=row["country"], hq_city=row["hq_city"], hq_region=row["hq_region"], founded_year=row["founded_year"], company_type=row["company_type"], public_company=row["public_company"], ticker=row["ticker"], exchange=row["exchange"], employees=row["employees"], wikidata_id=row["wikidata_id"], lei=row["lei"], sec_cik=row["sec_cik"], logo_url=row["logo_url"], importance=row["importance"], tier=row["tier"], source_meta=jsonb(row["source_meta"])) async def _link_parents(conn: AsyncConnection, pairs: list[tuple[str, str]], source: str) -> int: """(child_id, parent_id) → PARENT_OF + SUBSIDIARY_OF when missing. Confidence 0.8 (Wikidata P749, not verified on the web).""" if not pairs: return 0 existing = await fetch_all(conn, """ select from_company_id, to_company_id, kind from company_relationships where kind in ('PARENT_OF', 'SUBSIDIARY_OF') and from_company_id = any(cast(:ids as text[]))""", ids=sorted({c for c, _ in pairs} | {p for _, p in pairs})) have = {(r["from_company_id"], r["to_company_id"], r["kind"]) for r in existing} rows = [] for child, parent in pairs: for frm, to, kind in ((parent, child, "PARENT_OF"), (child, parent, "SUBSIDIARY_OF")): if (frm, to, kind) not in have: have.add((frm, to, kind)) rows.append({"id": ids.new_id("relationship"), "frm": frm, "to": to, "kind": kind, "prov": jsonb({"source": source, "property": "P749"})}) await execute_many(conn, """ insert into company_relationships (id, from_company_id, to_company_id, kind, confidence, provenance) values (:id, :frm, :to, :kind, 0.8, cast(:prov as jsonb))""", rows) return len(rows) async def upsert_companies(conn: AsyncConnection, rows: Iterable[dict[str, Any]], *, source: str = "registry", limit: int | None = None, enqueue: bool = True) -> dict[str, int]: idx = await _Index.load(conn) counters: dict[str, int] = {"companies_seen": 0, "companies_new": 0, "companies_updated": 0, "companies_skipped": 0, "aliases": 0, "domains": 0, "relationships": 0, "queue_jobs": 0} parents: list[tuple[str, str]] = [] # (child_id, parent_wikidata_id) new_ids: list[tuple[str, float]] = [] seen_domains: set[str] = set() for raw in rows: if limit is not None and counters["companies_seen"] >= limit: break counters["companies_seen"] += 1 row = normalise_row(raw, source=source) if row is None or row["canonical_domain"] in seen_domains: counters["companies_skipped"] += 1 continue seen_domains.add(row["canonical_domain"]) if row["country"] and row["country"] not in idx.countries: row["source_meta"]["unknown_country"] = row["country"] row["country"] = None company_id = idx.by_wikidata.get(row["wikidata_id"] or "") or idx.by_domain.get(row["canonical_domain"]) if company_id: await _update_company(conn, company_id, row) counters["companies_updated"] += 1 if row["wikidata_id"]: idx.by_wikidata.setdefault(row["wikidata_id"], company_id) # matched by domain: parent links must resolve this run else: company_id = await _insert_company(conn, row, idx) counters["companies_new"] += 1 new_ids.append((company_id, row["importance"])) aliases = _alias_rows(company_id, row, source) await execute_many(conn, """ insert into company_aliases (company_id, alias, alias_norm, kind, source) values (:company_id, :alias, :alias_norm, :kind, :source) on conflict (company_id, alias_norm) do nothing""", aliases) counters["aliases"] += len(aliases) await execute(conn, """ insert into domains (id, company_id, domain, kind) values (:id, :company_id, :domain, 'primary') on conflict (domain, company_id) do nothing""", id=ids.new_id("domain"), company_id=company_id, domain=row["canonical_domain"]) counters["domains"] += 1 if row["parent_wikidata_id"] and row["parent_wikidata_id"] != row["wikidata_id"]: parents.append((company_id, row["parent_wikidata_id"])) pairs = [(child, idx.by_wikidata[pq]) for child, pq in parents if pq in idx.by_wikidata and idx.by_wikidata[pq] != child] counters["relationships"] = await _link_parents(conn, pairs, source) if enqueue: counters["queue_jobs"] = await enqueue_discovery(conn, [cid for cid, _ in new_ids]) return counters async def enqueue_discovery(conn: AsyncConnection, company_ids: list[str] | None = None) -> int: """`discover` jobs for pending companies (all pending when `company_ids` is None). Idempotent on the job key.""" if company_ids is None: pending = await fetch_all(conn, "select id, importance from companies where onboarding_status = 'pending'") elif company_ids: pending = await fetch_all(conn, "select id, importance from companies where id = any(cast(:ids as text[])) and onboarding_status = 'pending'", ids=company_ids) else: return 0 rows = [{"id": ids.new_id("queue_job"), "key": f"discover:{r['id']}", "priority": float(r["importance"]), "payload": jsonb({"company_id": r["id"]})} for r in pending] await execute_many(conn, """ insert into queue_jobs (id, kind, key, payload, priority) values (:id, 'discover', :key, cast(:payload as jsonb), :priority) on conflict (key) do nothing""", rows) return len(rows) # ------------------------------------------------------------------------------------------------------------ public API async def seed(conn: AsyncConnection, *, companies: bool = True, limit: int | None = None, files: Iterable[Path] | None = None) -> dict[str, int]: """Upsert industries + countries, then (optionally) every company of the registry NDJSON files. Idempotent.""" started = datetime.now(UTC) counters: dict[str, int] = {"industries": await upsert_industries(conn), "countries": await upsert_countries(conn)} if companies: rows = load_registry_rows(files) counters.update(await upsert_companies(conn, rows, source="wikidata", limit=limit)) await execute(conn, """ insert into settings_kv (key, value, updated_at) values ('seed:last_run', cast(:v as jsonb), now()) on conflict (key) do update set value = excluded.value, updated_at = now()""", v=jsonb({"started_at": started.isoformat(), "finished_at": datetime.now(UTC).isoformat(), "counters": counters})) log.info("seed done", extra=counters) return counters async def import_companies(conn: AsyncConnection, rows: Iterable[dict[str, Any]], *, source: str = "manual") -> dict[str, int]: """Manual CSV/NDJSON rows (`website` required; `display_name` derived from the domain when missing).""" await upsert_industries(conn) await upsert_countries(conn) return await upsert_companies(conn, rows, source=source) async def add_company(conn: AsyncConnection, website: str, **fields: Any) -> dict[str, Any]: """Add (or top up) a single company by website. Returns the stored row (id, slug, canonical_domain, onboarding_status).""" counters = await import_companies(conn, [{"website": website, **fields}], source=fields.pop("source", "manual")) domain = registrable_domain(website if "://" in website else "https://" + website) row = await fetch_all(conn, "select id, slug, display_name, canonical_domain, onboarding_status from companies where canonical_domain = :d", d=domain) return {**(row[0] if row else {}), "counters": counters} __all__ = ["COMPANIES_DIR", "add_company", "enqueue_discovery", "import_companies", "load_registry_rows", "normalise_row", "read_rows_file", "registry_files", "seed", "unique_slug", "upsert_companies", "upsert_countries", "upsert_industries"]