update for mcp

2026-08-06 11:08:46 +08:00
parent 31bbf80aeb
commit b2feaeddb4
40 changed files with 1986 additions and 202 deletions
@@ -0,0 +1,157 @@
+"""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