momentum-system/storage/event_store.py

"""Event store implementation for Milestone 2.

Append-only event storage with canonical event types.
"""

from .event_schema import (
    BaseEvent,
    EventType,
    EventMode,
)
from .db import get_connection
from uuid import uuid4
import json
import logging
import threading

logger = logging.getLogger(__name__)


class EventStore:
    """Append-only event store for canonical events."""

    def __init__(self, db_path: str = "momentum"):
        """Initialize event store. ``db_path`` is retained for call-site
        compatibility; the datastore is Postgres (see storage.db_pg)."""
        self.con = get_connection(db_path)
        # serialize DB access: the fast trigger thread and the main loop share
        # this one psycopg2 connection, which is not safe for concurrent use.
        self._lock = threading.RLock()
        self._ensure_schema()

    def _ensure_schema(self):
        """Ensure event tables exist in database."""
        cursor = self.con.cursor()
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS events (
                id VARCHAR PRIMARY KEY,
                timestamp TIMESTAMP,
                mode VARCHAR,
                event_type VARCHAR,
                correlation_id VARCHAR,
                message VARCHAR,
                payload_json TEXT,
                created_at TIMESTAMP DEFAULT current_timestamp
            );
        """)
        # Composite (event_type, timestamp DESC) — THE hot-path index. Almost
        # every projection runs `WHERE event_type=? ORDER BY timestamp DESC
        # LIMIT n` (latest snapshot of a type). With only a plain timestamp
        # index, looking up the latest of a STALE type backward-scans the whole
        # table (measured: 235k rows filtered, ~370ms). This makes it an index
        # range seek (sub-ms) regardless of how long ago that type last fired,
        # and also serves bare `WHERE event_type=?` as the leftmost prefix.
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_events_type_ts
            ON events(event_type, timestamp DESC);
        """)
        # plain timestamp index still earns its keep for type-agnostic range
        # scans (e.g. `WHERE timestamp::date = CURRENT_DATE`).
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_events_timestamp
            ON events(timestamp);
        """)
        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_events_session_id
            ON events(correlation_id);
        """)
        # No index on payload_json: a btree over full JSON text is invalid for
        # large payloads in Postgres and useless for the LIKE-based symbol
        # filter (query_events scans with payload_json LIKE instead).
        self.con.commit()

    def emit(self, event: BaseEvent) -> str:
        """Emit an event to the event store.

        Args:
            event: Event to emit

        Returns:
            Event ID
        """
        event_id = str(uuid4())
        payload_json = json.dumps(event.model_dump(mode="json", exclude_unset=True))

        with self._lock:
            cursor = self.con.cursor()
            cursor.execute(
                """
                INSERT INTO events
                    (id, timestamp, mode, event_type, correlation_id, message, payload_json, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, current_timestamp)
                """,
                [
                    event_id,
                    event.timestamp,
                    event.mode.value,
                    event.event_type.value
                    if isinstance(event.event_type, EventType)
                    else str(event.event_type),
                    event.correlation_id,
                    event.message,
                    payload_json,
                ],
            )
            self.con.commit()

        logger.debug(
            f"Emitted event {event_id}: {event.event_type.value} - {event.message}"
        )
        return event_id

    def query_events(
        self,
        event_type: str | None = None,
        symbol: str | None = None,
        correlation_id: str | None = None,
        session_id: str | None = None,
        since: str | None = None,
        until: str | None = None,
        limit: int | None = 1000,
    ) -> list[dict]:
        """Query events from the store.

        Args:
            event_type: Filter by event type
            symbol: Filter by symbol
            correlation_id: Filter by correlation ID
            session_id: Filter by session ID
            since: Start timestamp (inclusive)
            until: End timestamp (exclusive)
            limit: Maximum number of results

        Returns:
            List of event dictionaries
        """
        conditions = []
        params = []

        if event_type is not None:
            conditions.append("event_type = ?")
            params.append(event_type)

        if symbol is not None:
            conditions.append("payload_json LIKE ?")
            params.append(f"%\"symbol\": \"{symbol}\"%")

        if correlation_id is not None:
            conditions.append("correlation_id = ?")
            params.append(correlation_id)

        if session_id is not None:
            conditions.append("correlation_id = ?")
            params.append(session_id)

        if since is not None:
            conditions.append("timestamp >= ?")
            params.append(since)

        if until is not None:
            conditions.append("timestamp < ?")
            params.append(until)

        where_clause = " AND ".join(conditions) if conditions else "1=1"
        limit_clause = f"LIMIT {limit}" if limit is not None else ""

        query = f"""
            SELECT id, timestamp, mode, event_type, correlation_id, message, payload_json, created_at
            FROM events
            WHERE {where_clause}
            ORDER BY timestamp ASC, created_at ASC
            {limit_clause}
        """

        with self._lock:
            cursor = self.con.cursor()
            cursor.execute(query, params)
            rows = cursor.fetchall()

        results = []
        for row in rows:
            results.append(
                {
                    "id": row[0],
                    "timestamp": str(row[1]),
                    "mode": row[2],
                    "event_type": row[3],
                    "correlation_id": row[4],
                    "message": row[5],
                    "payload_json": row[6],
                    "created_at": str(row[7]),
                }
            )

        return results

    def count_events(self) -> int:
        """Total number of events. Cheap change-detection signal for SSE."""
        with self._lock:
            cursor = self.con.cursor()
            cursor.execute("SELECT COUNT(*) FROM events")
            row = cursor.fetchone()
        return int(row[0]) if row else 0

    def get_symbol_timeline(self, symbol: str) -> list[dict]:
        """Get full timeline for a symbol.

        Args:
            symbol: Ticker symbol

        Returns:
            List of all events for the symbol in chronological order
        """
        return self.query_events(symbol=symbol, limit=None)

    def rebuild_state(self, session_id: str) -> dict:
        """Rebuild state from event store for a session.

        Args:
            session_id: Session ID to rebuild

        Returns:
            Dictionary with session summary and latest state
        """
        events = self.query_events(session_id=session_id, limit=None)

        summary = {
            "session_id": session_id,
            "total_events": len(events),
            "event_types": {},
        }

        for event in events:
            event_type = event["event_type"]
            summary["event_types"][event_type] = (
                summary["event_types"].get(event_type, 0) + 1
            )

        return summary

    def close(self):
        """Close event store connection."""
        if self.con:
            self.con.close()