spb/company-atlas

1"""Connector SDK (spec §8, §105–106): metadata, base class, registry.23    @register4    class GreenhouseConnector(Connector):5        meta = ConnectorMeta(connector_id="greenhouse-v1", name="Greenhouse job board", version="1", category=Surface.JOBS_BOARD, …)6        def extract(self, sensor, result) -> Extraction: ...78A *connector* is a strategy; a *sensor* row binds a connector to a URL for one company. `connector_id` (which embeds the version)9is stored on every observation/snapshot as `connector_version`, so history remains reproducible when connectors evolve — bump the10version (new id) when extraction behaviour changes materially.1112Only `fetch.Fetcher` touches the network; connectors receive it and must not open sockets themselves.13"""14from __future__ import annotations1516import importlib17import logging18import pkgutil19import re20from collections.abc import Mapping21from dataclasses import dataclass, field22from typing import Any, ClassVar2324from companyatlas.fetch import Fetcher, FetchResult25from companyatlas.sdk.models import DiscoveredUrl, Extraction26from companyatlas.taxonomy import SURFACE_BASE_INTERVAL_S, FetchMode, Surface2728log = logging.getLogger(__name__)293031@dataclass(frozen=True, slots=True)32class ConnectorMeta:33    connector_id: str                       # e.g. "generic-html-v1" — stable id incl. version34    name: str35    version: str36    category: str                           # primary Surface value37    fetch_mode: str = FetchMode.HTTP38    supports_discovery: bool = False39    supports_incremental: bool = True       # honours etag / last-modified / delta semantics40    default_interval_s: int = 8640041    surfaces: tuple[str, ...] = ()          # surfaces this connector can serve (empty = category only)42    url_pattern: str | None = None          # regex on the sensor URL; when it matches the connector is preferred43    pattern_required: bool = False          # vendor / API connectors: never auto-selected unless the URL pattern matches44    priority: int = 10                      # higher wins among candidates for the same surface45    accept: str | None = None               # Accept header override46    max_bytes: int | None = None47    description: str = ""484950@dataclass(slots=True)51class ConnectorContext:52    """What a connector may know about the company it is observing (read-only)."""53    company: Mapping[str, Any]54    extra: dict[str, Any] = field(default_factory=dict)5556    @property57    def canonical_domain(self) -> str:58        return str(self.company.get("canonical_domain") or "")596061class Connector:62    meta: ClassVar[ConnectorMeta]6364    # ---------------------------------------------------------------- fetch65    async def fetch(self, ctx: ConnectorContext, sensor: Mapping[str, Any], fetcher: Fetcher) -> FetchResult:66        """Default: conditional GET using the sensor's stored validators. Raises NotModified / FetchError."""67        return await fetcher.get(str(sensor["url"]), etag=sensor.get("etag"), last_modified=sensor.get("last_modified"),68                                 accept=self.meta.accept, max_bytes=self.meta.max_bytes)6970    # ---------------------------------------------------------------- extract71    def extract(self, sensor: Mapping[str, Any], result: FetchResult) -> Extraction:  # pragma: no cover - abstract72        raise NotImplementedError7374    # ---------------------------------------------------------------- discover (optional)75    async def discover(self, ctx: ConnectorContext, sensor: Mapping[str, Any], extraction: Extraction) -> list[DiscoveredUrl]:76        return list(extraction.discovered)7778    # ---------------------------------------------------------------- selection79    @classmethod80    def score(cls, surface: str, url: str) -> int:81        """0 = cannot serve; higher = better fit."""82        m = cls.meta83        s = 084        if m.url_pattern and re.search(m.url_pattern, url, re.IGNORECASE):85            s += 10086        elif m.pattern_required:87            return 088        if surface == m.category:89            s += 2090        elif surface in m.surfaces:91            s += 1092        elif "*" in m.surfaces:93            s += 594        return s + m.priority if s > 0 else 09596    @property97    def connector_id(self) -> str:98        return self.meta.connector_id99100    def __repr__(self) -> str:101        return f"<Connector {self.meta.connector_id}>"102103104# ------------------------------------------------------------------------------------------------------------ registry105106_REGISTRY: dict[str, type[Connector]] = {}107_INSTANCES: dict[str, Connector] = {}108_LOADED = False109110111def register(cls: type[Connector]) -> type[Connector]:112    meta = getattr(cls, "meta", None)113    if not isinstance(meta, ConnectorMeta):114        raise TypeError(f"{cls.__name__} must define `meta = ConnectorMeta(...)`")115    if meta.connector_id in _REGISTRY and _REGISTRY[meta.connector_id] is not cls:116        log.warning("connector id re-registered", extra={"connector_id": meta.connector_id})117    _REGISTRY[meta.connector_id] = cls118    _INSTANCES.pop(meta.connector_id, None)119    return cls120121122def load_all() -> None:123    """Import every `companyatlas.connectors.*` module once so decorators run."""124    global _LOADED125    if _LOADED:126        return127    _LOADED = True128    try:129        import companyatlas.connectors as pkg130    except ImportError:  # pragma: no cover131        return132    for mod in pkgutil.iter_modules(pkg.__path__):133        if mod.name.startswith("_"):134            continue135        try:136            importlib.import_module(f"companyatlas.connectors.{mod.name}")137        except Exception:138            log.exception("connector module failed to import", extra={"connector_module": mod.name})139140141def get(connector_id: str) -> Connector:142    load_all()143    inst = _INSTANCES.get(connector_id)144    if inst is None:145        cls = _REGISTRY.get(connector_id)146        if cls is None:147            # tolerate an older version id (e.g. "greenhouse-v0") by falling back to the newest of the same family148            family = connector_id.rsplit("-v", 1)[0]149            candidates = sorted((k for k in _REGISTRY if k.rsplit("-v", 1)[0] == family), reverse=True)150            if not candidates:151                raise KeyError(f"unknown connector {connector_id!r}")152            cls = _REGISTRY[candidates[0]]153        inst = _INSTANCES[connector_id] = cls()154    return inst155156157def all_connectors() -> list[Connector]:158    load_all()159    return [get(k) for k in sorted(_REGISTRY)]160161162def for_surface(surface: str, url: str) -> Connector:163    """Best connector for a (surface, url) pair — ATS/JSON/feed/sitemap by URL pattern, else the generic HTML connector."""164    load_all()165    best: tuple[int, str] | None = None166    for cid, cls in _REGISTRY.items():167        s = cls.score(surface, url)168        if s > 0 and (best is None or s > best[0] or (s == best[0] and cid < best[1])):169            best = (s, cid)170    if best is None:171        return get("generic-html-v1")172    return get(best[1])173174175def default_interval(surface: str, connector: Connector | None = None) -> int:176    if connector is not None and connector.meta.default_interval_s:177        return min(int(SURFACE_BASE_INTERVAL_S.get(surface, 86400)), connector.meta.default_interval_s)178    return int(SURFACE_BASE_INTERVAL_S.get(surface, 86400))179180181async def sync_connectors_table(conn: Any) -> int:182    """Upsert one `connectors` row per registered connector (idempotent). Returns the number of connectors."""183    from companyatlas.db import execute184185    n = 0186    for c in all_connectors():187        m = c.meta188        await execute(conn, """189            insert into connectors (id, name, version, category, fetch_mode, supports_discovery, supports_incremental, default_interval_s)190            values (:id, :name, :version, :category, :fetch_mode, :sd, :si, :interval)191            on conflict (id) do update set name = excluded.name, version = excluded.version, category = excluded.category,192                fetch_mode = excluded.fetch_mode, supports_discovery = excluded.supports_discovery,193                supports_incremental = excluded.supports_incremental, default_interval_s = excluded.default_interval_s, updated_at = now()194        """, id=m.connector_id, name=m.name, version=m.version, category=str(m.category), fetch_mode=str(m.fetch_mode),195                      sd=m.supports_discovery, si=m.supports_incremental, interval=int(m.default_interval_s))196        n += 1197    return n198199200__all__ = ["Connector", "ConnectorContext", "ConnectorMeta", "Surface", "all_connectors", "default_interval", "for_surface", "get", "load_all",201           "register", "sync_connectors_table"]202