"""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