signal-platform/app/services/data_import.py

"""Source-agnostic batch import framework (Dolt/SEC bulk data → PostgreSQL).

Every bulk importer (SEC facts, Dolt earnings, later Dolt stocks) plugs into
``run_import`` and gets, for free, the plan's non-negotiables:

- **One run per source at a time** — a Postgres *session-level* advisory lock
  keyed by source. It is held on a single pinned connection for the whole run,
  so it survives the intermediate commits (the ``running`` row, then the
  promotion) and only releases at the end. No-op on non-Postgres (tests).
- **Idempotent per revision** — the cheap ``detect_revision`` probe is compared
  against the last *promoted* run; an unchanged revision records a ``no_op``
  with **zero row changes** (no expensive fetch, no writes).
- **Staging then atomic promotion** — the importer stages into an in-memory
  object (no physical staging tables), validation reads it, and only a passing
  run calls ``promote`` whose writes commit together with the run-row flip to
  ``promoted`` in a single transaction.
- **Failure is inert** — a failed validation or a mid-run exception marks the
  run ``failed``, alerts via the system-events path, and leaves the live tables
  exactly as they were (nothing is written before ``promote``).

Every attempt — promoted, no_op, or failed — is recorded in ``data_import_runs``.
KISS: no conflicts table (summaries go in ``validation_json``), no revision
table (idempotency queries the last run), no aggregate tables.
"""

from __future__ import annotations

import asyncio
import hashlib
import json
import logging
from dataclasses import dataclass, field
from datetime import date, datetime, timezone
from typing import Any, Protocol, runtime_checkable

from sqlalchemy import select, text
from sqlalchemy.engine import Engine  # noqa: F401  (typing only)
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession

from app.database import engine as app_engine
from app.models.data_import_run import DataImportRun
from app.services import system_event_service

logger = logging.getLogger(__name__)

# data_import_runs.status values
STATUS_RUNNING = "running"
STATUS_VALIDATED = "validated"
STATUS_PROMOTED = "promoted"
STATUS_NO_OP = "no_op"
STATUS_FAILED = "failed"

_MAX_ERROR_LEN = 4000


@dataclass
class ValidationResult:
    """Outcome of an importer's validation gates.

    ``summary`` is serialized into ``validation_json`` (reconciliation /
    discrepancy details live here — no separate conflicts table). ``validate``
    MUST be read-only: it reads the staged object and, if needed, live tables
    for comparison, but writes nothing — that invariant is what makes a failed
    run leave the dataset untouched.
    """

    ok: bool
    summary: dict[str, Any] = field(default_factory=dict)
    source_max_date: date | None = None
    messages: list[str] = field(default_factory=list)


@runtime_checkable
class SourceImporter(Protocol):
    """Interface a concrete bulk importer implements. All methods receive the
    session bound to the lock-holding connection; ``stage`` and ``validate``
    never write to live tables, only ``promote`` does."""

    source: str  # sec_facts | dolt_earnings | dolt_stocks

    async def detect_revision(self, db: AsyncSession) -> str | None:
        """Cheap probe of the source revision (Dolt commit / SEC archive SHA).

        Returns the revision id, or None when it can't be determined cheaply
        (in which case idempotency is skipped and the run always stages)."""
        ...

    async def stage(self, db: AsyncSession) -> Any:
        """Download/parse into an in-memory staged representation. No writes to
        live tables."""
        ...

    async def validate(self, db: AsyncSession, staged: Any) -> ValidationResult:
        """Run the source's validation gates against ``staged``. Read-only."""
        ...

    async def promote(self, db: AsyncSession, staged: Any, run_id: int) -> dict[str, int]:
        """Apply ``staged`` to the live tables. Called inside the promotion
        transaction; the caller commits. ``run_id`` is the current
        ``data_import_runs.id`` so written rows can be stamped with their
        ``import_run_id``. Returns row-count deltas."""
        ...


def _advisory_key(source: str) -> int:
    """Deterministic signed 64-bit key for a source's advisory lock."""
    digest = hashlib.blake2b(source.encode("utf-8"), digest_size=8).digest()
    return int.from_bytes(digest, "big", signed=True)


async def _last_promoted_revision(db: AsyncSession, source: str) -> str | None:
    """Revision of the most recent *promoted* run for ``source`` (the revision
    currently loaded), or None if none has promoted yet."""
    row = await db.execute(
        select(DataImportRun.revision)
        .where(
            DataImportRun.source == source,
            DataImportRun.status == STATUS_PROMOTED,
        )
        .order_by(DataImportRun.id.desc())
        .limit(1)
    )
    return row.scalar_one_or_none()


def _now() -> datetime:
    return datetime.now(timezone.utc)


async def _alert(db: AsyncSession, source: str, code: str, messages: list[str]) -> None:
    try:
        await system_event_service.log_event(
            db,
            severity="error",
            source="data_import",
            code=f"{source}_{code}",
            message=(("; ".join(messages)) or code)[:_MAX_ERROR_LEN],
            dedup_key=f"data_import:{source}:{code}",
        )
    except Exception:  # noqa: BLE001 — alerting must never mask the real outcome
        logger.exception("Failed to emit data_import alert %s/%s", source, code)


async def run_import(
    importer: SourceImporter,
    *,
    engine: AsyncEngine | None = None,
) -> DataImportRun | None:
    """Run one import for ``importer``.

    Returns the recorded ``DataImportRun`` (promoted / no_op / failed), or None
    when the per-source advisory lock is already held (another run is active).
    """
    engine = engine or app_engine
    source = importer.source
    is_pg = engine.dialect.name == "postgresql"
    key = _advisory_key(source)

    async with engine.connect() as conn:
        # Bind the session to this one connection so the session-level advisory
        # lock persists across our commits. expire_on_commit must be set here —
        # the app factory's setting doesn't carry to a directly-built session.
        session = AsyncSession(bind=conn, expire_on_commit=False)
        try:
            if is_pg:
                got = (
                    await session.execute(
                        text("SELECT pg_try_advisory_lock(:k)"), {"k": key}
                    )
                ).scalar()
                await session.commit()
                if not got:
                    logger.info("data_import %s: lock held, skipping", source)
                    return None

            # Record the attempt FIRST — before the external revision probe, the
            # most likely failure — so anything below is recorded and alerted and
            # never escapes unrecorded. Revision is filled in once detected.
            run = DataImportRun(
                source=source,
                status=STATUS_RUNNING,
                started_at=_now(),
            )
            session.add(run)
            await session.commit()
            await session.refresh(run)

            try:
                revision = await importer.detect_revision(session)
                run.revision = revision
                last_rev = await _last_promoted_revision(session, source)
                if revision is not None and revision == last_rev:
                    run.status = STATUS_NO_OP
                    run.completed_at = _now()
                    await session.commit()
                    logger.info("data_import %s: no_op (revision %s)", source, revision)
                    return run

                staged = await importer.stage(session)
                result = await importer.validate(session, staged)
                run.source_max_date = result.source_max_date
                run.validation_json = json.dumps(result.summary, default=str)

                if not result.ok:
                    run.status = STATUS_FAILED
                    run.error_details = ("; ".join(result.messages))[:_MAX_ERROR_LEN]
                    run.completed_at = _now()
                    await session.commit()
                    await _alert(session, source, "validation_failed", result.messages)
                    logger.warning(
                        "data_import %s: validation failed: %s",
                        source,
                        result.messages,
                    )
                    return run

                # Promotion: importer writes + run-row flip in one transaction.
                row_counts = await importer.promote(session, staged, run.id)
                run.status = STATUS_PROMOTED
                run.row_counts_json = json.dumps(row_counts, default=str)
                run.completed_at = _now()
                await session.commit()
                await session.refresh(run)
                logger.info(
                    "data_import %s: promoted (revision %s, rows %s)",
                    source,
                    revision,
                    row_counts,
                )
                return run

            except asyncio.CancelledError:
                # Deploy / scheduler shutdown: best-effort mark failed so no
                # ``running`` row lingers, then let the cancellation propagate —
                # never swallow it.
                try:
                    await session.rollback()
                    run.status = STATUS_FAILED
                    run.error_details = "cancelled"
                    run.completed_at = _now()
                    await session.commit()
                except BaseException:  # noqa: BLE001 — best-effort during teardown
                    logger.warning(
                        "data_import %s: could not record cancellation", source
                    )
                raise

            except Exception as exc:  # noqa: BLE001 — record + alert, don't crash the job
                await session.rollback()
                run.status = STATUS_FAILED
                run.error_details = repr(exc)[:_MAX_ERROR_LEN]
                run.completed_at = _now()
                try:
                    await session.commit()
                except Exception:  # noqa: BLE001
                    logger.exception("data_import %s: failed to record failure", source)
                await _alert(session, source, "import_error", [repr(exc)])
                logger.exception("data_import %s: import error", source)
                return run

        finally:
            if is_pg:
                try:
                    await session.execute(
                        text("SELECT pg_advisory_unlock(:k)"), {"k": key}
                    )
                    await session.commit()
                except Exception:  # noqa: BLE001
                    logger.exception("data_import %s: failed to release lock", source)
            await session.close()