import json import sys from contextlib import asynccontextmanager from pathlib import Path from typing import AsyncGenerator import aiosqlite _DDL = """ PRAGMA journal_mode=WAL; PRAGMA foreign_keys=ON; CREATE TABLE IF NOT EXISTS subscriptions ( id TEXT PRIMARY KEY, folder_id TEXT NOT NULL, folder_name TEXT NOT NULL, expiry_datetime TEXT NOT NULL, notification_url TEXT NOT NULL, client_state TEXT NOT NULL, is_active INTEGER NOT NULL DEFAULT 1, created_at TEXT NOT NULL DEFAULT (datetime('now')), last_renewed_at TEXT ); CREATE TABLE IF NOT EXISTS notifications ( id INTEGER PRIMARY KEY AUTOINCREMENT, received_at TEXT NOT NULL DEFAULT (datetime('now')), subscription_id TEXT, change_type TEXT, resource TEXT, message_id TEXT, email_identifier TEXT, processing_status TEXT NOT NULL DEFAULT 'pending', error_message TEXT, storage_path TEXT ); CREATE TABLE IF NOT EXISTS routing_results ( id INTEGER PRIMARY KEY AUTOINCREMENT, notification_id INTEGER NOT NULL, processed_at TEXT NOT NULL DEFAULT (datetime('now')), email_identifier TEXT, attachment_filename TEXT, document_type TEXT, action TEXT, destination_path TEXT, confidence_scores TEXT, error_message TEXT ); CREATE TABLE IF NOT EXISTS daily_brief ( id INTEGER PRIMARY KEY AUTOINCREMENT, created_at TEXT NOT NULL DEFAULT (datetime('now')), email_identifier TEXT NOT NULL, notification_id INTEGER NOT NULL, sender_email TEXT, subject TEXT, attachment_filename TEXT, reason TEXT NOT NULL, confidence_scores TEXT, error_message TEXT, included_in_report INTEGER NOT NULL DEFAULT 0, report_sent_at TEXT ); CREATE TABLE IF NOT EXISTS client_domains ( domain TEXT PRIMARY KEY, client_name TEXT NOT NULL, created_at TEXT NOT NULL DEFAULT (datetime('now')), last_seen_at TEXT, notes TEXT ); """ async def init_db(db_path: Path) -> None: db_path.parent.mkdir(parents=True, exist_ok=True) async with aiosqlite.connect(db_path) as conn: await conn.executescript(_DDL) await conn.commit() @asynccontextmanager async def get_db(db_path: Path) -> AsyncGenerator[aiosqlite.Connection, None]: async with aiosqlite.connect(db_path) as conn: conn.row_factory = aiosqlite.Row await conn.execute("PRAGMA journal_mode=WAL") await conn.execute("PRAGMA foreign_keys=ON") yield conn async def upsert_subscription( conn: aiosqlite.Connection, sub_id: str, folder_id: str, folder_name: str, expiry_datetime: str, notification_url: str, client_state: str, ) -> None: await conn.execute( """ INSERT INTO subscriptions (id, folder_id, folder_name, expiry_datetime, notification_url, client_state, is_active) VALUES (?, ?, ?, ?, ?, ?, 1) ON CONFLICT(id) DO UPDATE SET expiry_datetime = excluded.expiry_datetime, last_renewed_at = datetime('now'), is_active = 1 """, (sub_id, folder_id, folder_name, expiry_datetime, notification_url, client_state), ) await conn.commit() async def get_active_subscription(conn: aiosqlite.Connection) -> dict | None: cursor = await conn.execute( "SELECT * FROM subscriptions WHERE is_active = 1 ORDER BY created_at DESC LIMIT 1" ) row = await cursor.fetchone() return dict(row) if row else None async def mark_subscription_inactive(conn: aiosqlite.Connection, sub_id: str) -> None: await conn.execute( "UPDATE subscriptions SET is_active = 0 WHERE id = ?", (sub_id,) ) await conn.commit() async def log_notification( conn: aiosqlite.Connection, subscription_id: str | None, change_type: str | None, resource: str | None, message_id: str | None, email_identifier: str | None = None, ) -> int: cursor = await conn.execute( """ INSERT INTO notifications (subscription_id, change_type, resource, message_id, email_identifier) VALUES (?, ?, ?, ?, ?) """, (subscription_id, change_type, resource, message_id, email_identifier or message_id), ) await conn.commit() return cursor.lastrowid async def update_subscription_expiry( conn: aiosqlite.Connection, sub_id: str, expiry_datetime: str ) -> None: await conn.execute( """ UPDATE subscriptions SET expiry_datetime = ?, last_renewed_at = datetime('now') WHERE id = ? """, (expiry_datetime, sub_id), ) await conn.commit() async def get_notification_by_email_identifier( conn: aiosqlite.Connection, email_identifier: str ) -> dict | None: cursor = await conn.execute( "SELECT * FROM notifications WHERE email_identifier = ? LIMIT 1", (email_identifier,), ) row = await cursor.fetchone() return dict(row) if row else None async def update_notification_status( conn: aiosqlite.Connection, notif_id: int, status: str, error_message: str | None = None, storage_path: str | None = None, ) -> None: await conn.execute( """ UPDATE notifications SET processing_status = ?, error_message = ?, storage_path = ? WHERE id = ? """, (status, error_message, storage_path, notif_id), ) await conn.commit() async def get_fetched_notifications(conn: aiosqlite.Connection) -> list[dict]: cursor = await conn.execute( "SELECT * FROM notifications WHERE processing_status = 'fetched' ORDER BY received_at" ) rows = await cursor.fetchall() return [dict(r) for r in rows] async def set_notification_routing_status( conn: aiosqlite.Connection, notif_id: int, status: str ) -> None: await conn.execute( "UPDATE notifications SET processing_status = ? WHERE id = ?", (status, notif_id), ) await conn.commit() async def log_routing_result( conn: aiosqlite.Connection, notification_id: int, filename: str | None, doc_type: str | None, action: str, email_identifier: str | None = None, destination_path: str | None = None, confidence_scores: dict | None = None, error_message: str | None = None, ) -> None: await conn.execute( """ INSERT INTO routing_results (notification_id, email_identifier, attachment_filename, document_type, action, destination_path, confidence_scores, error_message) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, ( notification_id, email_identifier, filename, doc_type, action, destination_path, json.dumps(confidence_scores) if confidence_scores else None, error_message, ), ) await conn.commit() async def insert_daily_brief( conn: aiosqlite.Connection, email_identifier: str, notification_id: int, sender_email: str | None, subject: str | None, filename: str | None, reason: str, confidence_scores: dict | None = None, error_message: str | None = None, ) -> None: await conn.execute( """ INSERT INTO daily_brief (email_identifier, notification_id, sender_email, subject, attachment_filename, reason, confidence_scores, error_message) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, ( email_identifier, notification_id, sender_email, subject, filename, reason, json.dumps(confidence_scores) if confidence_scores else None, error_message, ), ) await conn.commit() async def get_pending_brief_items(conn: aiosqlite.Connection) -> list[dict]: cursor = await conn.execute( "SELECT * FROM daily_brief WHERE included_in_report = 0 ORDER BY created_at" ) rows = await cursor.fetchall() return [dict(r) for r in rows] async def mark_brief_items_sent(conn: aiosqlite.Connection, ids: list[int]) -> None: await conn.execute( f"UPDATE daily_brief SET included_in_report = 1, report_sent_at = datetime('now') " f"WHERE id IN ({','.join('?' * len(ids))})", ids, ) await conn.commit() async def run_retention_cleanup(conn: aiosqlite.Connection) -> None: await conn.executescript(""" DELETE FROM notifications WHERE received_at < datetime('now', '-12 months'); DELETE FROM routing_results WHERE processed_at < datetime('now', '-12 months'); DELETE FROM daily_brief WHERE created_at < datetime('now', '-12 months'); DELETE FROM subscriptions WHERE is_active = 0 AND created_at < datetime('now', '-12 months'); """) await conn.commit() if __name__ == "__main__": import asyncio if "--init" in sys.argv: from app.config import settings asyncio.run(init_db(settings.database_path)) print(f"Database initialized at {settings.database_path}")