spb/company-atlas
Public
Python 66.3%
TypeScript 22.7%
JavaScript 8.6%
HTML 1.4%
CSS 0.7%
1"""Idempotent seed loader: registry files → `industries`, `countries`, `companies` (+ aliases, domains, relationships, discover queue).23Contract with the crawl core (docs/ARCHITECTURE.md, "Seeds → Crawl"): new companies get `onboarding_status='pending'` and a4`queue_jobs(kind='discover', key='discover:<company_id>')` row. Re-running never overwrites a non-null value except5`importance`, `tier` and the provenance in `source_meta` (spec: factual provenance, historical-first).6"""7from __future__ import annotations89import csv10import json11import logging12from collections.abc import Iterable13from datetime import UTC, datetime14from pathlib import Path15from typing import Any1617from sqlalchemy.ext.asyncio import AsyncConnection1819from companyatlas import ids20from companyatlas.db import execute, execute_many, fetch_all, jsonb21from companyatlas.registry.industries import REGISTRY_DIR, is_valid_slug, load_industries22from companyatlas.urls import registrable_domain2324log = logging.getLogger(__name__)2526COMPANIES_DIR = REGISTRY_DIR / "companies"27COUNTRIES_FILE = REGISTRY_DIR / "countries.csv"28DEFAULT_IMPORTANCE = 0.229DEFAULT_TIER = 430GENERIC_TLD_NAMES = {"www", "web", "site", "home", "online", "official"}313233# ------------------------------------------------------------------------------------------------------------ registry files34def registry_files() -> list[Path]:35 return sorted(COMPANIES_DIR.glob("*.ndjson"))363738def load_registry_rows(files: Iterable[Path] | None = None) -> list[dict[str, Any]]:39 rows: list[dict[str, Any]] = []40 for path in files or registry_files():41 with Path(path).open(encoding="utf-8") as f:42 for n, line in enumerate(f, 1):43 line = line.strip()44 if not line:45 continue46 try:47 rows.append(json.loads(line))48 except json.JSONDecodeError as e:49 raise ValueError(f"{path}:{n}: invalid JSON ({e})") from e50 return rows515253def load_countries_file(path: Path = COUNTRIES_FILE) -> list[dict[str, Any]]:54 with path.open(encoding="utf-8") as f:55 out = []56 for r in csv.DictReader(f):57 out.append({"code": r["code"].strip().upper(), "name": r["name"], "region": r.get("region") or None,58 "subregion": r.get("subregion") or None, "lat": float(r["lat"]) if r.get("lat") else None,59 "lon": float(r["lon"]) if r.get("lon") else None})60 return out616263def read_rows_file(path: Path) -> list[dict[str, Any]]:64 """Manual import: `.ndjson`/`.jsonl`/`.json` (list) or `.csv` with a header (`website` required)."""65 path = Path(path)66 if path.suffix.lower() == ".csv":67 with path.open(encoding="utf-8-sig") as f:68 rows = []69 for r in csv.DictReader(f):70 row = {k.strip(): (v.strip() if isinstance(v, str) else v) for k, v in r.items() if k}71 for key in ("industries", "aliases"):72 if isinstance(row.get(key), str):73 row[key] = [x.strip() for x in row[key].replace(";", ",").split(",") if x.strip()]74 rows.append(row)75 return rows76 if path.suffix.lower() == ".json":77 data = json.loads(path.read_text(encoding="utf-8"))78 return list(data if isinstance(data, list) else data.get("companies", []))79 return load_registry_rows([path])808182# ------------------------------------------------------------------------------------------------------------ normalisation83def name_from_domain(domain: str) -> str:84 label = domain.split(".")[0]85 if label in GENERIC_TLD_NAMES and domain.count(".") >= 2:86 label = domain.split(".")[1]87 return label.replace("-", " ").replace("_", " ").title()888990def normalise_row(row: dict[str, Any], *, source: str) -> dict[str, Any] | None:91 website = (row.get("website") or "").strip()92 if not website:93 return None94 if "://" not in website:95 website = "https://" + website96 if not website.startswith(("http://", "https://")):97 return None98 domain = (row.get("canonical_domain") or "").strip().lower() or registrable_domain(website)99 if not domain or "." not in domain:100 return None101 display = (row.get("display_name") or "").strip() or name_from_domain(domain)102 industries = [s for s in (row.get("industries") or []) if isinstance(s, str) and is_valid_slug(s)]103 if row.get("industry") and is_valid_slug(row["industry"]) and row["industry"] not in industries:104 industries.insert(0, row["industry"])105 country = (row.get("country") or "").strip().upper() or None106 parent = row.get("parent") if isinstance(row.get("parent"), dict) else None107 try:108 importance = float(row.get("importance", DEFAULT_IMPORTANCE))109 except (TypeError, ValueError):110 importance = DEFAULT_IMPORTANCE111 try:112 tier = int(row.get("tier", DEFAULT_TIER))113 except (TypeError, ValueError):114 tier = DEFAULT_TIER115 meta = {116 "source": row.get("source") or source,117 "harvested_at": row.get("harvested_at"),118 "sitelinks": row.get("sitelinks"),119 "industry_labels": row.get("industry_labels") or [],120 "parent_wikidata_id": parent.get("wikidata_id") if parent else None,121 "parent_name": parent.get("name") if parent else None,122 "notes": row.get("notes") or [],123 "seeded_at": datetime.now(UTC).replace(microsecond=0).isoformat(),124 }125 meta = {k: v for k, v in meta.items() if v not in (None, [], "")}126 return {127 "wikidata_id": (row.get("wikidata_id") or None), "display_name": display[:200], "legal_name": (row.get("legal_name") or None),128 "aliases": [a for a in (row.get("aliases") or []) if isinstance(a, str) and a.strip()][:12], "website": website,129 "canonical_domain": domain, "country": country, "hq_city": row.get("hq_city") or None, "hq_region": row.get("hq_region") or None,130 "industries": industries, "industry_primary": industries[0] if industries else None, "founded_year": row.get("founded_year"),131 "employees": row.get("employees"), "public_company": bool(row.get("public_company")), "ticker": row.get("ticker") or None,132 "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,133 "logo_url": row.get("logo_url") or None, "description": (row.get("description") or None), "importance": max(0.0, min(1.0, importance)),134 "tier": min(4, max(1, tier)), "company_type": "public" if row.get("public_company") else None, "source_meta": meta,135 "parent_wikidata_id": parent.get("wikidata_id") if parent else None,136 }137138139def unique_slug(base: str, country: str | None, taken: set[str]) -> str:140 if base not in taken:141 return base142 if country and f"{base}-{country.lower()}" not in taken:143 return f"{base}-{country.lower()}"144 n = 2145 while f"{base}-{n}" in taken:146 n += 1147 return f"{base}-{n}"148149150# ------------------------------------------------------------------------------------------------------------ reference tables151async def upsert_industries(conn: AsyncConnection) -> int:152 rows = [{"slug": i.slug, "name": i.name, "parent": i.parent, "description": i.description, "keywords": list(i.keywords),153 "sort_order": i.sort_order} for i in load_industries()]154 # Parents first (FK on parent_slug); insert parent_slug in a second pass so ordering inside the file does not matter.155 await execute_many(conn, """156 insert into industries (slug, name, description, keywords, sort_order) values (:slug, :name, :description, cast(:keywords as text[]), :sort_order)157 on conflict (slug) do update set name = excluded.name, description = excluded.description, keywords = excluded.keywords,158 sort_order = excluded.sort_order""", rows)159 await execute_many(conn, "update industries set parent_slug = :parent where slug = :slug", [r for r in rows if r["parent"]])160 return len(rows)161162163async def upsert_countries(conn: AsyncConnection) -> int:164 rows = load_countries_file()165 await execute_many(conn, """166 insert into countries (code, name, region, subregion, lat, lon) values (:code, :name, :region, :subregion, :lat, :lon)167 on conflict (code) do update set name = excluded.name, region = excluded.region, subregion = excluded.subregion,168 lat = excluded.lat, lon = excluded.lon""", rows)169 return len(rows)170171172# ------------------------------------------------------------------------------------------------------------ companies173class _Index:174 """In-memory view of existing companies for one loader run (avoids a lookup per row)."""175176 def __init__(self, by_wikidata: dict[str, str], by_domain: dict[str, str], slugs: set[str], countries: set[str]) -> None:177 self.by_wikidata = by_wikidata178 self.by_domain = by_domain179 self.slugs = slugs180 self.countries = countries181182 @classmethod183 async def load(cls, conn: AsyncConnection) -> _Index:184 rows = await fetch_all(conn, "select id, slug, canonical_domain, wikidata_id from companies")185 countries = {r["code"] for r in await fetch_all(conn, "select code from countries")}186 return cls({r["wikidata_id"]: r["id"] for r in rows if r["wikidata_id"]}, {r["canonical_domain"]: r["id"] for r in rows},187 {r["slug"] for r in rows}, countries)188189190def _alias_rows(company_id: str, row: dict[str, Any], source: str) -> list[dict[str, Any]]:191 seen: set[str] = set()192 out: list[dict[str, Any]] = []193 for alias, kind in ([(row["display_name"], "brand"), (row.get("legal_name"), "legal"), (row.get("ticker"), "ticker")]194 + [(a, "alias") for a in row.get("aliases", [])]):195 if not alias:196 continue197 norm = ids.normalize_alias(alias)198 if not norm or norm in seen:199 continue200 seen.add(norm)201 out.append({"company_id": company_id, "alias": alias[:200], "alias_norm": norm[:200], "kind": kind, "source": source})202 return out203204205async def _insert_company(conn: AsyncConnection, row: dict[str, Any], idx: _Index) -> str:206 company_id = ids.new_id("company")207 slug = unique_slug(ids.slugify(row["display_name"]), row["country"], idx.slugs)208 idx.slugs.add(slug)209 await execute(conn, """210 insert into companies (id, slug, legal_name, display_name, canonical_domain, website, description, industries, industry_primary, country,211 hq_city, hq_region, founded_year, company_type, public_company, ticker, exchange, employees, wikidata_id, lei,212 sec_cik, logo_url, onboarding_status, importance, tier, source_meta)213 values (:id, :slug, :legal_name, :display_name, :canonical_domain, :website, :description, cast(:industries as text[]), :industry_primary,214 :country, :hq_city, :hq_region, :founded_year, :company_type, :public_company, :ticker, :exchange, :employees, :wikidata_id, :lei,215 :sec_cik, :logo_url, 'pending', :importance, :tier, cast(:source_meta as jsonb))""",216 id=company_id, slug=slug, legal_name=row["legal_name"], display_name=row["display_name"], canonical_domain=row["canonical_domain"],217 website=row["website"], description=row["description"], industries=row["industries"], industry_primary=row["industry_primary"],218 country=row["country"], hq_city=row["hq_city"], hq_region=row["hq_region"], founded_year=row["founded_year"],219 company_type=row["company_type"], public_company=row["public_company"], ticker=row["ticker"], exchange=row["exchange"],220 employees=row["employees"], wikidata_id=row["wikidata_id"], lei=row["lei"], sec_cik=row["sec_cik"], logo_url=row["logo_url"],221 importance=row["importance"], tier=row["tier"], source_meta=jsonb(row["source_meta"]))222 idx.by_domain[row["canonical_domain"]] = company_id223 if row["wikidata_id"]:224 idx.by_wikidata[row["wikidata_id"]] = company_id225 return company_id226227228async def _update_company(conn: AsyncConnection, company_id: str, row: dict[str, Any]) -> None:229 """Fill only null columns; refresh importance/tier and merge provenance. Never touches names, website or history."""230 await execute(conn, """231 update companies set232 legal_name = coalesce(legal_name, :legal_name), description = coalesce(description, :description),233 industries = case when cardinality(industries) = 0 then cast(:industries as text[]) else industries end,234 industry_primary = coalesce(industry_primary, :industry_primary), country = coalesce(country, :country),235 hq_city = coalesce(hq_city, :hq_city), hq_region = coalesce(hq_region, :hq_region), founded_year = coalesce(founded_year, :founded_year),236 company_type = coalesce(company_type, :company_type), public_company = public_company or :public_company,237 ticker = coalesce(ticker, :ticker), exchange = coalesce(exchange, :exchange), employees = coalesce(employees, :employees),238 wikidata_id = coalesce(wikidata_id, :wikidata_id), lei = coalesce(lei, :lei), sec_cik = coalesce(sec_cik, :sec_cik),239 logo_url = coalesce(logo_url, :logo_url), importance = :importance, tier = :tier,240 source_meta = source_meta || cast(:source_meta as jsonb), updated_at = now()241 where id = :id""",242 id=company_id, legal_name=row["legal_name"], description=row["description"], industries=row["industries"],243 industry_primary=row["industry_primary"], country=row["country"], hq_city=row["hq_city"], hq_region=row["hq_region"],244 founded_year=row["founded_year"], company_type=row["company_type"], public_company=row["public_company"], ticker=row["ticker"],245 exchange=row["exchange"], employees=row["employees"], wikidata_id=row["wikidata_id"], lei=row["lei"], sec_cik=row["sec_cik"],246 logo_url=row["logo_url"], importance=row["importance"], tier=row["tier"], source_meta=jsonb(row["source_meta"]))247248249async def _link_parents(conn: AsyncConnection, pairs: list[tuple[str, str]], source: str) -> int:250 """(child_id, parent_id) → PARENT_OF + SUBSIDIARY_OF when missing. Confidence 0.8 (Wikidata P749, not verified on the web)."""251 if not pairs:252 return 0253 existing = await fetch_all(conn, """254 select from_company_id, to_company_id, kind from company_relationships255 where kind in ('PARENT_OF', 'SUBSIDIARY_OF') and from_company_id = any(cast(:ids as text[]))""",256 ids=sorted({c for c, _ in pairs} | {p for _, p in pairs}))257 have = {(r["from_company_id"], r["to_company_id"], r["kind"]) for r in existing}258 rows = []259 for child, parent in pairs:260 for frm, to, kind in ((parent, child, "PARENT_OF"), (child, parent, "SUBSIDIARY_OF")):261 if (frm, to, kind) not in have:262 have.add((frm, to, kind))263 rows.append({"id": ids.new_id("relationship"), "frm": frm, "to": to, "kind": kind,264 "prov": jsonb({"source": source, "property": "P749"})})265 await execute_many(conn, """266 insert into company_relationships (id, from_company_id, to_company_id, kind, confidence, provenance)267 values (:id, :frm, :to, :kind, 0.8, cast(:prov as jsonb))""", rows)268 return len(rows)269270271async def upsert_companies(conn: AsyncConnection, rows: Iterable[dict[str, Any]], *, source: str = "registry", limit: int | None = None,272 enqueue: bool = True) -> dict[str, int]:273 idx = await _Index.load(conn)274 counters: dict[str, int] = {"companies_seen": 0, "companies_new": 0, "companies_updated": 0, "companies_skipped": 0, "aliases": 0,275 "domains": 0, "relationships": 0, "queue_jobs": 0}276 parents: list[tuple[str, str]] = [] # (child_id, parent_wikidata_id)277 new_ids: list[tuple[str, float]] = []278 seen_domains: set[str] = set()279 for raw in rows:280 if limit is not None and counters["companies_seen"] >= limit:281 break282 counters["companies_seen"] += 1283 row = normalise_row(raw, source=source)284 if row is None or row["canonical_domain"] in seen_domains:285 counters["companies_skipped"] += 1286 continue287 seen_domains.add(row["canonical_domain"])288 if row["country"] and row["country"] not in idx.countries:289 row["source_meta"]["unknown_country"] = row["country"]290 row["country"] = None291 company_id = idx.by_wikidata.get(row["wikidata_id"] or "") or idx.by_domain.get(row["canonical_domain"])292 if company_id:293 await _update_company(conn, company_id, row)294 counters["companies_updated"] += 1295 if row["wikidata_id"]:296 idx.by_wikidata.setdefault(row["wikidata_id"], company_id) # matched by domain: parent links must resolve this run297 else:298 company_id = await _insert_company(conn, row, idx)299 counters["companies_new"] += 1300 new_ids.append((company_id, row["importance"]))301 aliases = _alias_rows(company_id, row, source)302 await execute_many(conn, """303 insert into company_aliases (company_id, alias, alias_norm, kind, source) values (:company_id, :alias, :alias_norm, :kind, :source)304 on conflict (company_id, alias_norm) do nothing""", aliases)305 counters["aliases"] += len(aliases)306 await execute(conn, """307 insert into domains (id, company_id, domain, kind) values (:id, :company_id, :domain, 'primary')308 on conflict (domain, company_id) do nothing""", id=ids.new_id("domain"), company_id=company_id, domain=row["canonical_domain"])309 counters["domains"] += 1310 if row["parent_wikidata_id"] and row["parent_wikidata_id"] != row["wikidata_id"]:311 parents.append((company_id, row["parent_wikidata_id"]))312 pairs = [(child, idx.by_wikidata[pq]) for child, pq in parents if pq in idx.by_wikidata and idx.by_wikidata[pq] != child]313 counters["relationships"] = await _link_parents(conn, pairs, source)314 if enqueue:315 counters["queue_jobs"] = await enqueue_discovery(conn, [cid for cid, _ in new_ids])316 return counters317318319async def enqueue_discovery(conn: AsyncConnection, company_ids: list[str] | None = None) -> int:320 """`discover` jobs for pending companies (all pending when `company_ids` is None). Idempotent on the job key."""321 if company_ids is None:322 pending = await fetch_all(conn, "select id, importance from companies where onboarding_status = 'pending'")323 elif company_ids:324 pending = await fetch_all(conn, "select id, importance from companies where id = any(cast(:ids as text[])) and onboarding_status = 'pending'",325 ids=company_ids)326 else:327 return 0328 rows = [{"id": ids.new_id("queue_job"), "key": f"discover:{r['id']}", "priority": float(r["importance"]),329 "payload": jsonb({"company_id": r["id"]})} for r in pending]330 await execute_many(conn, """331 insert into queue_jobs (id, kind, key, payload, priority) values (:id, 'discover', :key, cast(:payload as jsonb), :priority)332 on conflict (key) do nothing""", rows)333 return len(rows)334335336# ------------------------------------------------------------------------------------------------------------ public API337async def seed(conn: AsyncConnection, *, companies: bool = True, limit: int | None = None, files: Iterable[Path] | None = None) -> dict[str, int]:338 """Upsert industries + countries, then (optionally) every company of the registry NDJSON files. Idempotent."""339 started = datetime.now(UTC)340 counters: dict[str, int] = {"industries": await upsert_industries(conn), "countries": await upsert_countries(conn)}341 if companies:342 rows = load_registry_rows(files)343 counters.update(await upsert_companies(conn, rows, source="wikidata", limit=limit))344 await execute(conn, """345 insert into settings_kv (key, value, updated_at) values ('seed:last_run', cast(:v as jsonb), now())346 on conflict (key) do update set value = excluded.value, updated_at = now()""",347 v=jsonb({"started_at": started.isoformat(), "finished_at": datetime.now(UTC).isoformat(), "counters": counters}))348 log.info("seed done", extra=counters)349 return counters350351352async def import_companies(conn: AsyncConnection, rows: Iterable[dict[str, Any]], *, source: str = "manual") -> dict[str, int]:353 """Manual CSV/NDJSON rows (`website` required; `display_name` derived from the domain when missing)."""354 await upsert_industries(conn)355 await upsert_countries(conn)356 return await upsert_companies(conn, rows, source=source)357358359async def add_company(conn: AsyncConnection, website: str, **fields: Any) -> dict[str, Any]:360 """Add (or top up) a single company by website. Returns the stored row (id, slug, canonical_domain, onboarding_status)."""361 counters = await import_companies(conn, [{"website": website, **fields}], source=fields.pop("source", "manual"))362 domain = registrable_domain(website if "://" in website else "https://" + website)363 row = await fetch_all(conn, "select id, slug, display_name, canonical_domain, onboarding_status from companies where canonical_domain = :d",364 d=domain)365 return {**(row[0] if row else {}), "counters": counters}366367368__all__ = ["COMPANIES_DIR", "add_company", "enqueue_discovery", "import_companies", "load_registry_rows", "normalise_row", "read_rows_file",369 "registry_files", "seed", "unique_slug", "upsert_companies", "upsert_countries", "upsert_industries"]370