osint-dashboard/app/ingestor.py

"""NATS JetStream consumer — ingests OSINT events from NATS streams."""

from __future__ import annotations

import json
import logging
import uuid
from datetime import datetime, timezone

import nats
from nats.errors import TimeoutError
from sqlalchemy.dialects.postgresql import insert as pg_insert

from database import async_session
from models import events as events_table
from models import fires as fires_table
from models import event_dedup as event_dedup_table
from config import NATS_URL
from sources import event_dedup_key

logger = logging.getLogger("osint.ingestor")

# NATS connection settings
NATS_URLS = NATS_URL
NATS_STREAM = "events"
NATS_DURABLE = "osint-ingestor"


# ── Active fire / hotspot routing ─────────────────────────────────────────

def _fire_row_from_msg(msg: dict) -> dict | None:
    """Map a NATS fire message to a ``fires`` table row (pre-DB, unit-testable).

    Returns None (dropped) when the idempotency key fields — latitude,
    longitude, acq_time, satellite — are missing or unparseable.
    """
    row = {
        "latitude": msg.get("latitude"),
        "longitude": msg.get("longitude"),
        "brightness": msg.get("brightness"),
        "confidence": msg.get("confidence"),
        "acq_time": msg.get("acq_time"),
        "satellite": msg.get("satellite"),
        "instrument": msg.get("instrument"),
        "bright_ti5": msg.get("bright_ti5"),
        "frp": msg.get("frp"),
        "daynight": msg.get("daynight"),
        "scan": msg.get("scan"),
        "track": msg.get("track"),
        "version": msg.get("version"),
        "raw": msg.get("raw"),
    }
    # Idempotency key must be complete; brightness/confidence may be absent in
    # malformed feeds but lat/lon/time/satellite are required for dedupe.
    if (
        row["latitude"] is None or row["longitude"] is None
        or row["acq_time"] is None or not row["satellite"]
    ):
        logger.warning(
            "dropping malformed fire message (missing natural key): %s",
            {k: msg.get(k) for k in ("latitude", "longitude", "acq_time", "satellite")},
        )
        return None
    if isinstance(row["acq_time"], str):
        try:
            row["acq_time"] = datetime.fromisoformat(row["acq_time"])
        except ValueError:
            logger.warning("dropping fire message with bad acq_time %r", row["acq_time"])
            return None
    if not isinstance(row["acq_time"], datetime):
        row["acq_time"] = datetime.fromisoformat(str(row["acq_time"]))
    return row


async def ingest_fire_row(msg: dict) -> bool:
    """Persist one FIRMS hotspot, idempotently.

    The (latitude, longitude, acq_time, satellite) primary key doubles as the
    dedupe key: ON CONFLICT DO NOTHING means a hotspot re-delivered on a later
    15-minute poll is silently ignored. Returns True if a new row was inserted,
    False if it was a duplicate (or dropped).
    """
    row = _fire_row_from_msg(msg)
    if row is None:
        return False
    async with async_session() as session:
        stmt = (
            pg_insert(fires_table)
            .values(**row)
            .on_conflict_do_nothing(constraint="pk_fires_natural_key")
        )
        result = await session.execute(stmt)
        await session.commit()
    inserted = bool(result.rowcount)
    if inserted:
        logger.info(
            "ingested fire %.5f,%.5f %s satellite=%s",
            row["latitude"], row["longitude"], row["acq_time"], row["satellite"],
        )
        lat, lon = row["latitude"], row["longitude"]
        from geofence import record_and_notify
        await record_and_notify(
            source_kind="firms",
            entity_id=f"{lat},{lon},{row['satellite']}",
            lat=lat, lon=lon, payload={"satellite": row["satellite"]},
        )
        from live_layers import aircraft_last_known
        from fire_aircraft import correlate_and_notify
        from tracks import recent_markers
        acs = list(aircraft_last_known.values()) or await recent_markers("aircraft")
        if acs:
            fire = {
                "id": f"firms:{lat:.4f},{lon:.4f}",
                "lat": lat, "lon": lon, "label": "FIRMS",
            }
            await correlate_and_notify([fire], acs)
    return inserted


async def ingest_fire_rows(msgs: list[dict]) -> int:
    """Bulk-insert FIRMS hotspots: one INSERT, one ON CONFLICT, one commit."""
    rows = []
    for msg in msgs:
        row = _fire_row_from_msg(msg)
        if row is not None:
            rows.append(row)
    if not rows:
        return 0
    async with async_session() as session:
        stmt = (
            pg_insert(fires_table)
            .values(rows)
            .on_conflict_do_nothing(constraint="pk_fires_natural_key")
        )
        result = await session.execute(stmt)
        await session.commit()
    inserted = int(result.rowcount or 0)
    if inserted:
        logger.info("bulk ingested %d/%d FIRMS hotspots", inserted, len(rows))
        from live_layers import aircraft_last_known
        from fire_aircraft import correlate_and_notify
        from tracks import recent_markers
        acs = list(aircraft_last_known.values()) or await recent_markers("aircraft")
        if acs:
            fires = [
                {
                    "id": f"firms:{r['latitude']:.4f},{r['longitude']:.4f}",
                    "lat": r["latitude"], "lon": r["longitude"], "label": "FIRMS",
                }
                for r in rows[:500]
            ]
            await correlate_and_notify(fires, acs)
    return inserted


async def ingest_event(msg: dict):
    """Ingest a single event from NATS into PostgreSQL."""
    # Active fire/hotspot messages carry a dedicated schema and land in the
    # `fires` hypertable (idempotent natural key), not the generic events feed.
    if msg.get("source_type") == "fire":
        return await ingest_fire_row(msg)
    # source_id links to a feed_sources UUID; tolerate non-UUID / missing values
    # (e.g. legacy messages that carried a URL) by coercing to None.
    raw_source_id = msg.get("source_id")
    source_id = None
    if raw_source_id is not None:
        try:
            source_id = str(uuid.UUID(str(raw_source_id)))
        except (ValueError, AttributeError, TypeError):
            source_id = None

    event_row = {
        "source_type": msg.get("source_type", "rss"),
        "source_id": source_id,
        "title": msg.get("title"),
        "body": msg.get("body"),
        "url": msg.get("url"),
        "sentiment_score": msg.get("sentiment_score"),
        "sentiment_label": msg.get("sentiment_label"),
        "location_lat": msg.get("location_lat"),
        "location_lon": msg.get("location_lon"),
        "location_name": msg.get("location_name"),
        "entities": msg.get("entities", []),
        "tags": msg.get("tags", []),
        "raw": msg.get("raw"),
        "source_timestamp": msg.get("source_timestamp", datetime.now(timezone.utc).isoformat()),
    }

    # Parse timestamp if string
    ts = event_row["source_timestamp"]
    if isinstance(ts, str):
        ts = datetime.fromisoformat(ts.replace("Z", "+00:00"))
    if isinstance(ts, datetime) and ts.tzinfo is None:
        ts = ts.replace(tzinfo=timezone.utc)
    event_row["source_timestamp"] = ts

    key = event_dedup_key(event_row)
    async with async_session() as session:
        if key:
            dedup = (
                pg_insert(event_dedup_table)
                .values(url=key)
                .on_conflict_do_nothing(index_elements=["url"])
            )
            claimed = await session.execute(dedup)
            if not claimed.rowcount:
                await session.commit()
                logger.info("skip duplicate event url=%s", key)
                return None
        result = await session.execute(events_table.insert().values(**event_row))
        await session.commit()
        event_id = result.inserted_primary_key[0]  # type: ignore[union-attr]
    logger.info("Ingested event %s from source %s", event_id, msg.get("source_type"))
    return event_id


async def start_nats_consumer():
    """Start NATS JetStream consumer for OSINT events."""
    nc = await nats.connect(NATS_URLS)
    js = nc.jetstream()

    # Create stream if not exists; if it exists, UPDATE its subject list so
    # newly added event types (e.g. events.camera) are actually covered.
    from nats.js.api import StreamConfig
    stream_cfg = StreamConfig(
        name=NATS_STREAM,
        subjects=[
            "events.gdelt", "events.rss", "events.social",
            "events.earthquake", "events.disaster", "events.weather",
            "events.fire", "events.satellite", "events.new", "events.alert",
            "events.camera",
        ],
        retention=nats.js.api.RetentionPolicy.LIMITS,
        max_msgs=1_000_000,
    )
    try:
        await js.add_stream(stream_cfg)
        logger.info("Created NATS stream %s", NATS_STREAM)
    except Exception:
        await js.update_stream(stream_cfg)
        logger.info("Updated NATS stream %s subjects", NATS_STREAM)

    # Create durable consumer
    sub = await js.pull_subscribe(
        subject="events.>",
        durable=NATS_DURABLE,
    )

    logger.info("NATS consumer started, durable=%s", NATS_DURABLE)
    return nc, sub


async def fetch_and_process(batch_size: int = 100):
    """Fetch a batch of messages and process them."""
    nc, sub = await start_nats_consumer()
    js = nc.jetstream()

    msgs = await sub.fetch(batch_size, timeout=5)
    processed = 0

    for msg in msgs:
        try:
            data = json.loads(msg.data)
            await ingest_event(data)
            await msg.ack()
            processed += 1
        except Exception:
            logger.error("Failed to process message: %s", msg.data, exc_info=True)

    await nc.close()
    logger.info("Processed %d messages in batch", processed)
    return processed