AIRegulation-DocAnalysis/backend/app/infrastructure/perception/postgres_notification_store.py

"""PostgreSQL-backed notification store.

One row per triggering event, shared by every user; a separate read-receipt
table tracks per-user read state so broadcasting to everyone needs no fan-out
insert per user. See base_notification_store.py for the port contract and the
design doc for why this shape was chosen over per-user subscriptions.
"""

from __future__ import annotations

from contextlib import contextmanager
from typing import Any

import psycopg2
import psycopg2.extras
from psycopg2.pool import ThreadedConnectionPool

from app.config.settings import settings
from app.infrastructure.perception.base_notification_store import BaseNotificationStore

_CREATE_TABLES = """
CREATE TABLE IF NOT EXISTS perception_notifications (
    id           SERIAL PRIMARY KEY,
    event_id     TEXT NOT NULL REFERENCES regulation_events(id) ON DELETE CASCADE,
    kind         TEXT NOT NULL,
    title        TEXT NOT NULL,
    impact_level TEXT,
    summary      TEXT,
    created_at   TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE TABLE IF NOT EXISTS perception_notification_reads (
    notification_id INTEGER NOT NULL REFERENCES perception_notifications(id) ON DELETE CASCADE,
    user_id          TEXT NOT NULL,
    read_at          TIMESTAMPTZ NOT NULL DEFAULT now(),
    PRIMARY KEY (notification_id, user_id)
);
CREATE INDEX IF NOT EXISTS perception_notif_created
    ON perception_notifications (created_at DESC);
"""


def _row_to_dict(row: dict[str, Any]) -> dict:
    """Convert a psycopg2 RealDictRow to a plain dict with an ISO timestamp."""
    d = dict(row)
    if d.get("created_at") is not None:
        d["created_at"] = d["created_at"].isoformat()
    return d


class PostgresNotificationStore(BaseNotificationStore):
    """Notification store backed by PostgreSQL."""

    def __init__(self) -> None:
        """Open a connection pool and ensure both tables exist."""
        self._pool = ThreadedConnectionPool(
            minconn=1,
            maxconn=5,
            host=settings.postgres_host,
            port=settings.postgres_port,
            user=settings.postgres_user,
            password=settings.postgres_password,
            dbname=settings.postgres_db,
        )
        self._ensure_schema()

    def _ensure_schema(self) -> None:
        with self._conn() as conn:
            try:
                with conn.cursor() as cur:
                    cur.execute(_CREATE_TABLES)
                conn.commit()
            except Exception:
                conn.rollback()
                raise

    @contextmanager
    def _conn(self):
        conn = None
        try:
            conn = self._pool.getconn()
            yield conn
        finally:
            if conn is not None:
                self._pool.putconn(conn)

    def create(
        self,
        *,
        event_id: str,
        kind: str,
        title: str,
        impact_level: str | None,
        summary: str | None,
    ) -> None:
        """Insert one notification row for the triggering event."""
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "INSERT INTO perception_notifications "
                    "(event_id, kind, title, impact_level, summary) "
                    "VALUES (%s, %s, %s, %s, %s)",
                    (event_id, kind, title, impact_level, summary),
                )
            conn.commit()

    def list_for_user(self, user_id: str, limit: int = 20) -> list[dict]:
        """Return the newest notifications with this user's read state joined in."""
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute(
                    """
                    SELECT n.*, (r.user_id IS NOT NULL) AS read
                    FROM perception_notifications n
                    LEFT JOIN perception_notification_reads r
                        ON r.notification_id = n.id AND r.user_id = %s
                    ORDER BY n.created_at DESC
                    LIMIT %s
                    """,
                    (user_id, limit),
                )
                return [_row_to_dict(r) for r in cur.fetchall()]

    def unread_count(self, user_id: str) -> int:
        """Count notifications with no read receipt for this user."""
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    SELECT COUNT(*) FROM perception_notifications n
                    WHERE NOT EXISTS (
                        SELECT 1 FROM perception_notification_reads r
                        WHERE r.notification_id = n.id AND r.user_id = %s
                    )
                    """,
                    (user_id,),
                )
                return cur.fetchone()[0]

    def mark_all_read(self, user_id: str) -> int:
        """Insert a read receipt for every notification this user hasn't read."""
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO perception_notification_reads (notification_id, user_id)
                    SELECT n.id, %s FROM perception_notifications n
                    WHERE NOT EXISTS (
                        SELECT 1 FROM perception_notification_reads r
                        WHERE r.notification_id = n.id AND r.user_id = %s
                    )
                    ON CONFLICT (notification_id, user_id) DO NOTHING
                    """,
                    (user_id, user_id),
                )
                marked = cur.rowcount
            conn.commit()
        return marked