197 lines
6.9 KiB
Python
197 lines
6.9 KiB
Python
"""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",
|
|
"events.camera",
|
|
],
|
|
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
|