@simbimbo/brainstem 0.0.1 → 0.0.3

package/brainstem/storage.py

@@ -0,0 +1,428 @@
+from __future__ import annotations
+
+import json
+import sqlite3
+from pathlib import Path
+from typing import Any, Iterable, List
+
+from .models import Candidate, Event, RawInputEnvelope, Signature
+
+
+def default_db_path() -> Path:
+    return Path('.brainstem-state') / 'brainstem.sqlite3'
+
+
+RAW_ENVELOPE_STATUSES = ("received", "canonicalized", "parse_failed", "unsupported")
+RAW_ENVELOPE_FAILURE_STATUSES = ("parse_failed", "unsupported")
+
+
+def _validate_canonicalization_status(status: str) -> None:
+    if status not in RAW_ENVELOPE_STATUSES:
+        raise ValueError(f"unsupported canonicalization_status: {status}")
+
+
+def connect(db_path: str | None = None) -> sqlite3.Connection:
+    path = Path(db_path) if db_path else default_db_path()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    conn = sqlite3.connect(path)
+    conn.row_factory = sqlite3.Row
+    return conn
+
+
+def init_db(db_path: str | None = None) -> None:
+    conn = connect(db_path)
+    try:
+        conn.executescript(
+            '''
+            CREATE TABLE IF NOT EXISTS events (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                tenant_id TEXT NOT NULL,
+                source_type TEXT NOT NULL,
+                timestamp TEXT NOT NULL,
+                host TEXT,
+                service TEXT,
+                severity TEXT,
+                asset_id TEXT,
+                source_path TEXT,
+                facility TEXT,
+                message_raw TEXT NOT NULL,
+                structured_fields_json TEXT NOT NULL,
+                correlation_keys_json TEXT NOT NULL
+            );
+
+            CREATE TABLE IF NOT EXISTS raw_envelopes (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                tenant_id TEXT NOT NULL,
+                source_type TEXT NOT NULL,
+                source_id TEXT,
+                source_name TEXT,
+                timestamp TEXT NOT NULL,
+                host TEXT,
+                service TEXT,
+                severity TEXT,
+                asset_id TEXT,
+                source_path TEXT,
+                facility TEXT,
+                message_raw TEXT NOT NULL,
+                structured_fields_json TEXT NOT NULL,
+                correlation_keys_json TEXT NOT NULL,
+                metadata_json TEXT NOT NULL,
+                canonicalization_status TEXT NOT NULL DEFAULT 'received',
+                failure_reason TEXT
+            );
+
+            CREATE TABLE IF NOT EXISTS signatures (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                signature_key TEXT NOT NULL UNIQUE,
+                event_family TEXT NOT NULL,
+                normalized_pattern TEXT NOT NULL,
+                service TEXT,
+                metadata_json TEXT NOT NULL,
+                occurrence_count INTEGER NOT NULL DEFAULT 0
+            );
+
+            CREATE TABLE IF NOT EXISTS candidates (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                candidate_type TEXT NOT NULL,
+                title TEXT NOT NULL,
+                summary TEXT NOT NULL,
+                score_total REAL NOT NULL,
+                score_breakdown_json TEXT NOT NULL,
+                decision_band TEXT NOT NULL,
+                source_signature_ids_json TEXT NOT NULL,
+                source_event_ids_json TEXT NOT NULL,
+                confidence REAL NOT NULL,
+                metadata_json TEXT NOT NULL
+            );
+            '''
+        )
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def store_raw_envelopes(raw_envelopes: Iterable[RawInputEnvelope], db_path: str | None = None) -> List[int]:
+    conn = connect(db_path)
+    raw_ids: List[int] = []
+    try:
+        for raw in raw_envelopes:
+            cursor = conn.execute(
+                '''
+                INSERT INTO raw_envelopes (
+                    tenant_id, source_type, source_id, source_name, timestamp, host, service, severity,
+                    asset_id, source_path, facility, message_raw,
+                    structured_fields_json, correlation_keys_json, metadata_json,
+                    canonicalization_status, failure_reason
+                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                ''',
+                (
+                    raw.tenant_id,
+                    raw.source_type,
+                    raw.source_id,
+                    raw.source_name,
+                    raw.timestamp,
+                    raw.host,
+                    raw.service,
+                    raw.severity,
+                    raw.asset_id,
+                    raw.source_path,
+                    raw.facility,
+                    raw.message_raw,
+                    json.dumps(raw.structured_fields, ensure_ascii=False),
+                    json.dumps(raw.correlation_keys, ensure_ascii=False),
+                    json.dumps(raw.metadata, ensure_ascii=False),
+                    "received",
+                    None,
+                ),
+            )
+            raw_ids.append(int(cursor.lastrowid))
+        conn.commit()
+        return raw_ids
+    finally:
+        conn.close()
+
+
+def set_raw_envelope_status(
+    raw_envelope_id: int,
+    status: str,
+    db_path: str | None = None,
+    *,
+    failure_reason: str | None = None,
+) -> None:
+    _validate_canonicalization_status(status)
+    conn = connect(db_path)
+    try:
+        conn.execute(
+            '''
+            UPDATE raw_envelopes
+            SET canonicalization_status = ?, failure_reason = ?
+            WHERE id = ?
+            ''',
+            (status, failure_reason, raw_envelope_id),
+        )
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def get_raw_envelope_by_id(raw_envelope_id: int, db_path: str | None = None) -> sqlite3.Row | None:
+    conn = connect(db_path)
+    try:
+        return conn.execute(
+            "SELECT * FROM raw_envelopes WHERE id = ?",
+            (raw_envelope_id,),
+        ).fetchone()
+    finally:
+        conn.close()
+
+
+def _recent_raw_envelopes_query(
+    canonicalization_status: str | None,
+    *,
+    failures_only: bool,
+) -> tuple[str, tuple[str, ...], bool]:
+    if canonicalization_status is None and failures_only:
+        return "WHERE canonicalization_status IN (?, ?)", RAW_ENVELOPE_FAILURE_STATUSES, True
+    if canonicalization_status is None and not failures_only:
+        return "", (), False
+    _validate_canonicalization_status(canonicalization_status)
+    return "WHERE canonicalization_status = ?", (canonicalization_status,), False
+
+
+def list_recent_raw_envelopes(
+    db_path: str | None = None,
+    status: str | None = None,
+    limit: int = 20,
+    *,
+    failures_only: bool = False,
+) -> List[sqlite3.Row]:
+    conn = connect(db_path)
+    try:
+        where_clause, status_args, _ = _recent_raw_envelopes_query(status, failures_only=failures_only)
+        prefix = f"{where_clause} " if where_clause else ""
+        rows = conn.execute(
+            f"""
+            SELECT * FROM raw_envelopes
+            {prefix}
+            ORDER BY id DESC
+            LIMIT ?
+            """,
+            (*status_args, max(1, limit)),
+        ).fetchall()
+        return rows
+    finally:
+        conn.close()
+
+
+def list_recent_failed_raw_envelopes(
+    db_path: str | None = None,
+    *,
+    status: str | None = None,
+    limit: int = 20,
+) -> List[sqlite3.Row]:
+    return list_recent_raw_envelopes(
+        db_path=db_path,
+        status=status,
+        limit=limit,
+        failures_only=status is None,
+    )
+
+
+def store_events(events: Iterable[Event], db_path: str | None = None) -> int:
+    conn = connect(db_path)
+    count = 0
+    try:
+        for event in events:
+            conn.execute(
+                '''
+                INSERT INTO events (
+                    tenant_id, source_type, timestamp, host, service, severity,
+                    asset_id, source_path, facility, message_raw,
+                    structured_fields_json, correlation_keys_json
+                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                ''',
+                (
+                    event.tenant_id,
+                    event.source_type,
+                    event.timestamp,
+                    event.host,
+                    event.service,
+                    event.severity,
+                    event.asset_id,
+                    event.source_path,
+                    event.facility,
+                    event.message_raw,
+                    json.dumps(event.structured_fields, ensure_ascii=False),
+                    json.dumps(event.correlation_keys, ensure_ascii=False),
+                ),
+            )
+            count += 1
+        conn.commit()
+        return count
+    finally:
+        conn.close()
+
+
+SOURCE_SUMMARY_DIMENSIONS = (
+    "source_type",
+    "source_path",
+    "source_id",
+    "source_name",
+    "host",
+    "service",
+)
+
+
+def _summarize_raw_envelopes_by_dimension(
+    conn: sqlite3.Connection,
+    dimension: str,
+    limit: int = 20,
+) -> List[dict[str, int | str]]:
+    if dimension not in SOURCE_SUMMARY_DIMENSIONS:
+        raise ValueError(f"unsupported dimension: {dimension}")
+    return [
+        {"value": row["value"], "count": int(row["count"])}
+        for row in conn.execute(
+            f"""
+            SELECT {dimension} AS value, COUNT(*) AS count
+            FROM raw_envelopes
+            WHERE COALESCE(TRIM({dimension}), '') <> ''
+            GROUP BY {dimension}
+            ORDER BY COUNT(*) DESC, value ASC
+            LIMIT ?
+            """,
+            (max(1, limit),),
+        ).fetchall()
+    ]
+
+
+def get_source_dimension_summaries(
+    db_path: str | None = None,
+    *,
+    limit: int = 20,
+) -> dict[str, List[dict[str, int | str]]]:
+    init_db(db_path)
+    conn = connect(db_path)
+    try:
+        return _get_source_dimension_summaries_from_conn(conn, limit=limit)
+    finally:
+        conn.close()
+
+
+def _get_source_dimension_summaries_from_conn(
+    conn: sqlite3.Connection,
+    *,
+    limit: int = 20,
+) -> dict[str, List[dict[str, int | str]]]:
+    return {
+        "source_type": _summarize_raw_envelopes_by_dimension(conn, "source_type", limit=limit),
+        "source_path": _summarize_raw_envelopes_by_dimension(conn, "source_path", limit=limit),
+        "source_id": _summarize_raw_envelopes_by_dimension(conn, "source_id", limit=limit),
+        "source_name": _summarize_raw_envelopes_by_dimension(conn, "source_name", limit=limit),
+        "host": _summarize_raw_envelopes_by_dimension(conn, "host", limit=limit),
+        "service": _summarize_raw_envelopes_by_dimension(conn, "service", limit=limit),
+    }
+
+
+def _query_count(conn: sqlite3.Connection, query: str) -> int:
+    return int(conn.execute(query).fetchone()[0])
+
+
+def get_ingest_stats(db_path: str | None = None) -> dict[str, Any]:
+    init_db(db_path)
+    conn = connect(db_path)
+    try:
+        return {
+            "received": _query_count(conn, "SELECT COUNT(*) FROM raw_envelopes"),
+            "canonicalized": _query_count(
+                conn,
+                "SELECT COUNT(*) FROM raw_envelopes WHERE canonicalization_status = 'canonicalized'",
+            ),
+            "parse_failed": _query_count(
+                conn,
+                "SELECT COUNT(*) FROM raw_envelopes WHERE canonicalization_status = 'parse_failed'",
+            ),
+            "unsupported": _query_count(
+                conn,
+                "SELECT COUNT(*) FROM raw_envelopes WHERE canonicalization_status = 'unsupported'",
+            ),
+            "candidates_generated": _query_count(conn, "SELECT COUNT(*) FROM candidates"),
+            "source_summaries": _get_source_dimension_summaries_from_conn(conn),
+        }
+    finally:
+        conn.close()
+
+
+def store_signatures(signatures: Iterable[Signature], db_path: str | None = None) -> int:
+    conn = connect(db_path)
+    count = 0
+    try:
+        for signature in signatures:
+            conn.execute(
+                '''
+                INSERT INTO signatures (
+                    signature_key, event_family, normalized_pattern, service, metadata_json, occurrence_count
+                ) VALUES (?, ?, ?, ?, ?, 1)
+                ON CONFLICT(signature_key) DO UPDATE SET
+                    occurrence_count = occurrence_count + 1,
+                    metadata_json = excluded.metadata_json,
+                    service = excluded.service
+                ''',
+                (
+                    signature.signature_key,
+                    signature.event_family,
+                    signature.normalized_pattern,
+                    signature.service,
+                    json.dumps(signature.metadata, ensure_ascii=False),
+                ),
+            )
+            count += 1
+        conn.commit()
+        return count
+    finally:
+        conn.close()
+
+
+def store_candidates(candidates: Iterable[Candidate], db_path: str | None = None) -> int:
+    conn = connect(db_path)
+    count = 0
+    try:
+        for candidate in candidates:
+            conn.execute(
+                '''
+                INSERT INTO candidates (
+                    candidate_type, title, summary, score_total, score_breakdown_json,
+                    decision_band, source_signature_ids_json, source_event_ids_json,
+                    confidence, metadata_json
+                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                ''',
+                (
+                    candidate.candidate_type,
+                    candidate.title,
+                    candidate.summary,
+                    candidate.score_total,
+                    json.dumps(candidate.score_breakdown, ensure_ascii=False),
+                    candidate.decision_band,
+                    json.dumps(candidate.source_signature_ids, ensure_ascii=False),
+                    json.dumps(candidate.source_event_ids, ensure_ascii=False),
+                    candidate.confidence,
+                    json.dumps(candidate.metadata, ensure_ascii=False),
+                ),
+            )
+            count += 1
+        conn.commit()
+        return count
+    finally:
+        conn.close()
+
+
+def list_candidates(db_path: str | None = None, limit: int = 20) -> List[sqlite3.Row]:
+    conn = connect(db_path)
+    try:
+        rows = conn.execute(
+            'SELECT * FROM candidates ORDER BY score_total DESC, id DESC LIMIT ?',
+            (max(1, limit),),
+        ).fetchall()
+        return rows
+    finally:
+        conn.close()