5302e261d7
A ~90k-row INSERT ... ON CONFLICT dies with InterfaceError, so fire_loop fails every poll and /api/health stays degraded (fires age > 30 min). Chunk 2000 rows per statement, still one commit per poll.
280 lines
10 KiB
Python
280 lines
10 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 models import event_dedup as event_dedup_table
|
||
from config import NATS_URL
|
||
from sources import event_dedup_key
|
||
|
||
logger = logging.getLogger("osint.ingestor")
|
||
|
||
# asyncpg rejects statements with >32767 bind params. A FIRMS poll is ~90k
|
||
# rows × 14 columns. Chunk inserts; still one transaction / one commit.
|
||
FIRE_ROW_BIND_PARAMS = 14
|
||
FIRE_INSERT_CHUNK = 2000
|
||
|
||
# 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
|
||
inserted = 0
|
||
async with async_session() as session:
|
||
for i in range(0, len(rows), FIRE_INSERT_CHUNK):
|
||
chunk = rows[i:i + FIRE_INSERT_CHUNK]
|
||
stmt = (
|
||
pg_insert(fires_table)
|
||
.values(chunk)
|
||
.on_conflict_do_nothing(constraint="pk_fires_natural_key")
|
||
)
|
||
result = await session.execute(stmt)
|
||
inserted += int(result.rowcount or 0)
|
||
await session.commit()
|
||
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
|