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 config import NATS_URL

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"],
        )
    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
    if isinstance(event_row["source_timestamp"], str):
        event_row["source_timestamp"] = datetime.fromisoformat(event_row["source_timestamp"])

    async with async_session() as session:
        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
    try:
        await js.add_stream(
            name=NATS_STREAM,
            subjects=[
                "events.gdelt", "events.rss", "events.social",
                "events.earthquake", "events.disaster", "events.weather",
                "events.fire", "events.satellite", "events.new", "events.alert",
            ],
            retention=nats.js.api.RetentionPolicy.LIMITS,
            max_msgs=1_000_000,
        )
        logger.info("Created NATS stream %s", NATS_STREAM)
    except Exception:
        logger.debug("Stream %s already exists", 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