politics-tracking/backend/politrack/core/loader.py

"""The Loader: the only code that writes canonical rows.

Invariants enforced here (not in adapters):
  * no claim without evidence; claims whose person cannot be resolved are skipped (and counted),
  * full text is kept only for sources with store_fulltext and responses without TDM reservation,
  * claim identity is computed from *resolved* entity ids + type-specific fields, so the same fact
    from different sources becomes one claim with several evidence rows ("corroborated"),
  * claims marked `review=True` start as review_status='pending' (hidden until a human approves),
  * persons are only created by adapters with creates_persons=True.
"""

from __future__ import annotations

from datetime import date, datetime
from decimal import Decimal, InvalidOperation
from typing import Any

from psycopg import Connection
from psycopg.types.json import Jsonb

from politrack.core.adapter import RunContext
from politrack.core.records import (
    Claim,
    Document,
    Event,
    Evidence,
    Legal,
    Money,
    OrgRecord,
    OrgRef,
    Participation,
    PersonRecord,
    PersonRef,
    Role,
    SocialAccount,
    Statement,
    VoteEvent,
    VotePosition,
)
from politrack.core.resolve import Resolver
from politrack.text import excerpt, fingerprint, make_slug, normalize, normalize_org_name, strip_titles

# Source policy is read from the DB once per process (synced from SourceSpecs).
_SOURCES: dict[str, dict[str, Any]] = {}


def source_row(conn: Connection, slug: str) -> dict[str, Any]:
    if slug not in _SOURCES:
        row = conn.execute("SELECT id, store_fulltext, kind FROM source WHERE slug = %s", (slug,)).fetchone()
        if row is None:
            raise KeyError(f"unknown source {slug!r}: run `politrack sources sync`")
        _SOURCES[slug] = row
    return _SOURCES[slug]


class Loader:
    def __init__(self, conn: Connection, ctx: RunContext) -> None:
        self.conn = conn
        self.ctx = ctx
        self.resolver = Resolver(conn)

    # -- dispatch ---------------------------------------------------------------------

    def load(self, record: Any) -> int | None:
        if isinstance(record, Claim):
            return self.claim(record)
        if isinstance(record, Document):
            return self.document(record.evidence)
        if isinstance(record, PersonRecord):
            return self.person(record)
        if isinstance(record, OrgRecord):
            return self.organization(record)
        if isinstance(record, SocialAccount):
            return self.social_account(record)
        if isinstance(record, VoteEvent):
            return self.vote_event(record)
        if isinstance(record, Event):
            return self.event(record)
        raise TypeError(f"unknown record type {type(record).__name__}")

    # -- documents --------------------------------------------------------------------

    def document(self, ev: Evidence) -> int:
        slug = ev.source or self.ctx.spec.slug
        src = source_row(self.conn, slug)
        response = self.ctx.response
        tdm = bool(response and response.tdm_reserved)
        full_text = ev.full_text if (src["store_fulltext"] and not tdm) else None
        short = None if tdm else (excerpt(ev.excerpt, 400) if ev.excerpt else None)
        # Link the document to the raw response only if it *is* that response (not e.g. a protocol it cites).
        fetch_id = response.fetch_id if response and ev.url in (response.url, response.request.url) else None
        row = self.conn.execute(
            """
            INSERT INTO document (source_id, url, title, doc_type, published_at, excerpt, full_text, fetch_id, meta)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
            ON CONFLICT (url) DO UPDATE SET
                title = COALESCE(EXCLUDED.title, document.title),
                published_at = COALESCE(EXCLUDED.published_at, document.published_at),
                excerpt = COALESCE(EXCLUDED.excerpt, document.excerpt),
                full_text = COALESCE(EXCLUDED.full_text, document.full_text),
                fetch_id = COALESCE(EXCLUDED.fetch_id, document.fetch_id),
                meta = document.meta || EXCLUDED.meta,
                retrieved_at = now()
            RETURNING id
            """,
            (src["id"], ev.url, ev.title, ev.doc_type, ev.published_at, short, full_text, fetch_id,
             Jsonb({**ev.meta, **({"tdm_reserved": True} if tdm else {})})),
        ).fetchone()
        self.ctx.stats["documents"] += 1
        return row["id"]

    # -- entities ---------------------------------------------------------------------

    def person(self, rec: PersonRecord) -> int | None:
        person_id = self.resolver.person_by_ids(rec.ref.ids)
        if person_id is None and rec.match_by_name:
            person_id = self.resolver.person_by_name(rec.full_name)
        if person_id is None:
            if not self.ctx.creates_persons:
                self.ctx.stats["persons_not_created"] += 1
                return None
            base = make_slug(rec.given_name or "", rec.family_name or strip_titles(rec.full_name)) or "person"
            slug, n = base, 1
            while self.conn.execute("SELECT 1 FROM person WHERE slug = %s", (slug,)).fetchone():
                n += 1
                slug = f"{base}-{n}"
            person_id = self.conn.execute(
                """INSERT INTO person (slug, full_name, given_name, family_name, honorific_prefix, honorific_suffix,
                                       birth_date, death_date, gender, photo_url)
                   VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s) RETURNING id""",
                (slug, rec.full_name, rec.given_name, rec.family_name, rec.honorific_prefix, rec.honorific_suffix,
                 rec.birth_date, rec.death_date, rec.gender, rec.photo_url),
            ).fetchone()["id"]
            self.ctx.stats["persons_created"] += 1
        else:
            self.conn.execute(
                """UPDATE person SET full_name = %s,
                       given_name = COALESCE(%s, given_name), family_name = COALESCE(%s, family_name),
                       honorific_prefix = COALESCE(%s, honorific_prefix), honorific_suffix = COALESCE(%s, honorific_suffix),
                       birth_date = COALESCE(%s, birth_date), death_date = COALESCE(%s, death_date),
                       gender = COALESCE(%s, gender), photo_url = COALESCE(%s, photo_url), updated_at = now()
                   WHERE id = %s""",
                (rec.full_name, rec.given_name, rec.family_name, rec.honorific_prefix, rec.honorific_suffix,
                 rec.birth_date, rec.death_date, rec.gender, rec.photo_url, person_id),
            )
        for scheme, value in rec.ref.ids.items():
            if value:
                self.conn.execute(
                    """INSERT INTO identifier (entity_type, entity_id, scheme, value, source_slug)
                       VALUES ('person', %s, %s, %s, %s) ON CONFLICT DO NOTHING""",
                    (person_id, scheme, str(value), self.ctx.spec.slug),
                )
        names = {rec.full_name, strip_titles(rec.full_name), *rec.aliases}
        if rec.given_name and rec.family_name:
            names.add(f"{rec.given_name} {rec.family_name}")
        for name in names:
            if name:
                self.conn.execute(
                    """INSERT INTO alias (entity_type, entity_id, alias, alias_norm) VALUES ('person', %s, %s, %s)
                       ON CONFLICT DO NOTHING""",
                    (person_id, name, normalize(strip_titles(name))),
                )
        self.ctx.stats["persons"] += 1
        return person_id

    def organization(self, rec: OrgRecord) -> int | None:
        org_id = self.resolver.org(rec.ref)
        if org_id is None:
            return None
        if rec.color:
            self.conn.execute("UPDATE organization SET color = COALESCE(color, %s) WHERE id = %s", (rec.color, org_id))
        if rec.parent:
            parent = self.resolver.org(rec.parent)
            if parent and parent != org_id:
                self.conn.execute("UPDATE organization SET parent_id = %s WHERE id = %s", (parent, org_id))
        self.resolver.add_org_aliases(org_id, list(rec.aliases))
        return org_id

    def social_account(self, rec: SocialAccount) -> int | None:
        person_id = self.resolver.person(rec.person)
        org_id = self.resolver.org(rec.org) if rec.org else None  # OrgRef.create decides
        if person_id is None and org_id is None:
            self.ctx.stats["social_unresolved"] += 1
            return None
        row = self.conn.execute(
            """INSERT INTO social_account (person_id, organization_id, platform, handle, url, verified, source_slug)
               VALUES (%s, %s, %s, %s, %s, %s, %s)
               ON CONFLICT (platform, url) DO UPDATE SET
                   handle = COALESCE(EXCLUDED.handle, social_account.handle),
                   verified = social_account.verified OR EXCLUDED.verified,
                   person_id = COALESCE(social_account.person_id, EXCLUDED.person_id),
                   organization_id = COALESCE(social_account.organization_id, EXCLUDED.organization_id)
               RETURNING id""",
            (person_id, org_id, rec.platform.lower(), rec.handle, rec.url, rec.verified, self.ctx.spec.slug),
        ).fetchone()
        self.ctx.stats["social_accounts"] += 1
        return row["id"]

    # -- standalone facets ------------------------------------------------------------

    def vote_event(self, rec: VoteEvent) -> int:
        chamber = self.resolver.org(rec.chamber, create=False) if rec.chamber else None
        row = self.conn.execute(
            """INSERT INTO vote_event (event_key, chamber_id, item_ref, item_title, stage, voted_on, result, roll_call,
                                       url, protocol_url, raw_text)
               VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
               ON CONFLICT (event_key) DO UPDATE SET result = EXCLUDED.result, raw_text = EXCLUDED.raw_text,
                   roll_call = EXCLUDED.roll_call, item_title = EXCLUDED.item_title,
                   protocol_url = COALESCE(EXCLUDED.protocol_url, vote_event.protocol_url)
               RETURNING id""",
            (rec.key, chamber, rec.item_ref, rec.title, rec.stage, rec.voted_on, rec.result, rec.roll_call, rec.url,
             rec.protocol_url, rec.raw_text),
        ).fetchone()
        self.ctx.stats["vote_events"] += 1
        return row["id"]

    def event(self, rec: Event) -> int:
        body = self.resolver.org(rec.body, create=False) if rec.body else None
        source_id = source_row(self.conn, self.ctx.spec.slug)["id"]
        row = self.conn.execute(
            """INSERT INTO event (event_key, kind, title, starts_at, ends_at, location_name, country, url, description,
                                  organization_id, source_id)
               VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
               ON CONFLICT (event_key) DO UPDATE SET title = EXCLUDED.title,
                   starts_at = COALESCE(EXCLUDED.starts_at, event.starts_at), ends_at = EXCLUDED.ends_at,
                   location_name = COALESCE(EXCLUDED.location_name, event.location_name),
                   description = COALESCE(EXCLUDED.description, event.description)
               RETURNING id""",
            (rec.key, rec.kind, rec.title, rec.starts_at, rec.ends_at, rec.location, rec.country, rec.url,
             rec.description, body, source_id),
        ).fetchone()
        self.ctx.stats["events"] += 1
        return row["id"]

    # -- claims -----------------------------------------------------------------------

    def claim(self, c: Claim) -> int | None:
        if not c.evidence:
            raise ValueError("a claim needs at least one piece of evidence")
        person_id = self.resolver.person(c.person)
        if c.person is not None and person_id is None:
            self.ctx.stats["unresolved_person"] += 1
            self.ctx.log.debug("unresolved person %s for %s", c.person, c.summary_de[:80])
            return None
        org_id = self.resolver.org(c.org) if c.org else None
        cp_person = self.resolver.person(c.counterpart_person)
        cp_org = self.resolver.org(c.counterpart_org) if c.counterpart_org else None
        if person_id is None and org_id is None and cp_org is None:
            self.ctx.stats["claims_without_subject"] += 1
            return None

        fp = fingerprint(*self.identity(c, person_id, org_id, cp_org))
        row = self.conn.execute(
            """
            INSERT INTO claim (claim_type, fingerprint, person_id, organization_id, counterpart_person_id,
                               counterpart_org_id, occurred_at, valid_from, valid_to, summary_de, summary_en,
                               data, sensitive, review_status, created_by)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            ON CONFLICT (fingerprint) DO UPDATE SET
                summary_de = EXCLUDED.summary_de,
                summary_en = COALESCE(EXCLUDED.summary_en, claim.summary_en),
                valid_to = COALESCE(EXCLUDED.valid_to, claim.valid_to),
                occurred_at = COALESCE(claim.occurred_at, EXCLUDED.occurred_at),
                data = claim.data || EXCLUDED.data,
                sensitive = claim.sensitive OR EXCLUDED.sensitive,
                updated_at = now()
            RETURNING id, (xmax = 0) AS inserted
            """,
            (c.claim_type, fp, person_id, org_id, cp_person, cp_org, c.occurred_at, c.valid_from, c.valid_to,
             c.summary_de, c.summary_en, Jsonb(c.data), c.sensitive or c.claim_type == "legal_proceeding",
             "pending" if c.review else "published", self.ctx.adapter),
        ).fetchone()
        claim_id = row["id"]
        self.ctx.stats["claims_new" if row["inserted"] else "claims_updated"] += 1

        for ev in c.evidence:
            doc_id = self.document(ev)
            self.conn.execute(
                """INSERT INTO claim_evidence (claim_id, document_id, stance, quote, locator, extracted_by, confidence)
                   VALUES (%s, %s, %s, %s, %s, %s, %s)
                   ON CONFLICT (claim_id, document_id, locator) DO UPDATE SET
                       quote = COALESCE(EXCLUDED.quote, claim_evidence.quote),
                       confidence = GREATEST(claim_evidence.confidence, EXCLUDED.confidence)""",
                (claim_id, doc_id, ev.stance, ev.quote and excerpt(ev.quote, 600), ev.locator, self.ctx.adapter,
                 ev.confidence),
            )
        for topic in c.topics:
            self.assign_topic(claim_id, topic)
        if c.facet is not None:
            self.facet(claim_id, c, person_id, org_id)
        self.refresh_status(claim_id)
        return claim_id

    def identity(self, c: Claim, person_id: int | None, org_id: int | None, cp_org: int | None) -> tuple:
        """Semantic identity of a claim (see docs/architecture/01-design-from-scratch.md §3.3)."""
        if c.key is not None:
            return (c.claim_type, "key", *c.key)
        f = c.facet
        if isinstance(f, Role):
            body = self.resolver.org(f.body) if f.body else cp_org
            return ("role", person_id, body, normalize(f.role), f.start)
        if isinstance(f, VotePosition):
            return ("vote", f.event_key, person_id, org_id)
        if isinstance(f, Money):
            def side(ref: Any, name: str | None) -> Any:
                if isinstance(ref, PersonRef):
                    return ("p", self.resolver.person(ref))
                if isinstance(ref, OrgRef):
                    return ("o", self.resolver.org(ref))
                return ("n", normalize_org_name(name or ""))
            return ("money", f.kind, side(f.payer, f.payer_name), side(f.payee, f.payee_name), _amount(f.amount),
                    f.paid_on, f.period, f.reference)
        if isinstance(f, Statement):
            return ("statement", person_id, org_id, f.url)
        if isinstance(f, Participation):
            return ("participation", f.event_key, person_id, org_id)
        if c.claim_type == "media_mention":
            return ("mention", c.evidence[0].url, person_id, org_id)
        raise ValueError(f"claim of type {c.claim_type} needs an explicit key")

    def facet(self, claim_id: int, c: Claim, person_id: int | None, org_id: int | None) -> None:
        f = c.facet
        if isinstance(f, Role):
            body = self.resolver.org(f.body) if f.body else None
            party = self.resolver.org(f.party) if f.party else None
            self.conn.execute(
                """INSERT INTO role (claim_id, person_id, organization_id, party_id, role, level, region, start_date, end_date)
                   VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
                   ON CONFLICT (claim_id) DO UPDATE SET end_date = COALESCE(EXCLUDED.end_date, role.end_date),
                       party_id = COALESCE(EXCLUDED.party_id, role.party_id), region = COALESCE(EXCLUDED.region, role.region),
                       organization_id = COALESCE(EXCLUDED.organization_id, role.organization_id)""",
                (claim_id, person_id, body, party, f.role, f.level, f.region, f.start, f.end),
            )
        elif isinstance(f, VotePosition):
            event = self.conn.execute(
                """SELECT ve.id, ve.voted_on, o.slug AS chamber FROM vote_event ve
                   LEFT JOIN organization o ON o.id = ve.chamber_id WHERE ve.event_key = %s""",
                (f.event_key,),
            ).fetchone()
            if event is None:
                raise ValueError(f"vote event {f.event_key} must be loaded before its positions")
            party = self.resolver.org(f.party) if f.party else None
            if person_id is not None and party is None and event["chamber"]:
                # individual (roll-call) vote: remember the voter's faction on that day (dissent detection)
                party = self.resolver.faction_of(person_id, event["chamber"], event["voted_on"])
            self.conn.execute(
                """INSERT INTO vote_position (claim_id, vote_event_id, organization_id, person_id, party_id, position)
                   VALUES (%s, %s, %s, %s, %s, %s)
                   ON CONFLICT (claim_id) DO UPDATE SET position = EXCLUDED.position,
                       party_id = COALESCE(EXCLUDED.party_id, vote_position.party_id)""",
                (claim_id, event["id"], org_id if person_id is None else None, person_id, party, f.position),
            )
        elif isinstance(f, Money):
            payer_org = self.resolver.org(f.payer) if isinstance(f.payer, OrgRef) else None
            payee_org = self.resolver.org(f.payee) if isinstance(f.payee, OrgRef) else None
            payer_person = self.resolver.person(f.payer) if isinstance(f.payer, PersonRef) else None
            payee_person = self.resolver.person(f.payee) if isinstance(f.payee, PersonRef) else None
            self.conn.execute(
                """INSERT INTO money_flow (claim_id, kind, payer_name, payer_org_id, payer_person_id, payee_name,
                                           payee_org_id, payee_person_id, amount, amount_category, currency, paid_on,
                                           period, reference, category)
                   VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                   ON CONFLICT (claim_id) DO UPDATE SET amount = EXCLUDED.amount,
                       payer_org_id = COALESCE(EXCLUDED.payer_org_id, money_flow.payer_org_id),
                       payee_org_id = COALESCE(EXCLUDED.payee_org_id, money_flow.payee_org_id)""",
                (claim_id, f.kind, f.payer_name, payer_org, payer_person, f.payee_name, payee_org, payee_person,
                 _amount(f.amount), f.amount_category, f.currency, f.paid_on, f.period, f.reference, f.category),
            )
        elif isinstance(f, Participation):
            event = self.conn.execute("SELECT id FROM event WHERE event_key = %s", (f.event_key,)).fetchone()
            if event is None:
                raise ValueError(f"event {f.event_key} must be loaded before participations")
            self.conn.execute(
                """INSERT INTO event_participant (claim_id, event_id, person_id, organization_id, participation)
                   VALUES (%s, %s, %s, %s, %s) ON CONFLICT (claim_id) DO NOTHING""",
                (claim_id, event["id"], person_id, org_id, f.participation),
            )
        elif isinstance(f, Statement):
            self.conn.execute(
                """INSERT INTO statement (claim_id, person_id, context, platform, medium, title, said_at, excerpt, url)
                   VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) ON CONFLICT (claim_id) DO UPDATE SET
                       excerpt = COALESCE(EXCLUDED.excerpt, statement.excerpt)""",
                (claim_id, person_id, f.context, f.platform, f.medium, f.title, f.said_at,
                 f.excerpt and excerpt(f.excerpt, 600), f.url),
            )
        elif isinstance(f, Legal):
            self.conn.execute(
                """INSERT INTO legal_proceeding (claim_id, person_id, organization_id, stage, offense, authority, case_ref,
                                                 stage_date)
                   VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
                   ON CONFLICT (claim_id) DO UPDATE SET stage = EXCLUDED.stage,
                       offense = COALESCE(EXCLUDED.offense, legal_proceeding.offense),
                       authority = COALESCE(EXCLUDED.authority, legal_proceeding.authority),
                       case_ref = COALESCE(EXCLUDED.case_ref, legal_proceeding.case_ref)""",
                (claim_id, person_id, org_id, f.stage, f.offense, f.authority, f.case_ref, f.stage_date),
            )

    # -- helpers ----------------------------------------------------------------------

    def assign_topic(self, claim_id: int, topic_name: str, method: str | None = None, confidence: float = 1.0) -> None:
        norm = normalize(topic_name)
        row = self.conn.execute(
            "SELECT id FROM topic WHERE lower(unaccent(name_de)) = %s OR slug = %s", (norm, topic_name)
        ).fetchone()
        if row is None:
            row = self.conn.execute(
                "INSERT INTO topic (slug, name_de, name_en) VALUES (%s, %s, %s) ON CONFLICT (slug) DO UPDATE "
                "SET name_de = topic.name_de RETURNING id",
                (make_slug(topic_name), topic_name, topic_name),
            ).fetchone()
        self.conn.execute(
            "INSERT INTO topic_assignment (topic_id, claim_id, confidence, method) VALUES (%s, %s, %s, %s) "
            "ON CONFLICT DO NOTHING",
            (row["id"], claim_id, confidence, method or self.ctx.adapter),
        )

    def refresh_status(self, claim_id: int) -> None:
        """corroborated = backed by >= 2 different sources; disputed = any contradicting evidence."""
        self.conn.execute(
            """
            UPDATE claim SET status = CASE
                WHEN status = 'retracted' THEN status
                WHEN EXISTS (SELECT 1 FROM claim_evidence e WHERE e.claim_id = claim.id AND e.stance = 'contradicts')
                    THEN 'disputed'
                WHEN (SELECT count(DISTINCT d.source_id) FROM claim_evidence e JOIN document d ON d.id = e.document_id
                      WHERE e.claim_id = claim.id AND e.stance IN ('supports', 'mentions')) >= 2 THEN 'corroborated'
                ELSE 'reported' END
            WHERE id = %s
            """,
            (claim_id,),
        )


def _amount(value: Any) -> Decimal | None:
    if value is None or value == "":
        return None
    try:
        return Decimal(str(value)).quantize(Decimal("0.01"))
    except InvalidOperation:
        return None


def as_date(value: Any) -> date | None:
    if isinstance(value, datetime):
        return value.date()
    if isinstance(value, date):
        return value
    return None