"""Persistence for analysis jobs, mails, and weekly reports.""" from __future__ import annotations import logging from datetime import date from typing import Any from psycopg.rows import dict_row from psycopg.types.json import Jsonb from db.pool import get_pool, pool_enabled from mail_parser import ParsedMail from queries import Jobs, Mails, Reports, Schema logger = logging.getLogger("mail_summary") def _row_to_job(row: Any) -> dict[str, Any]: return { "id": row["id"], "status": row["status"], "mailbox_scope": row["mailbox_scope"], "date_from": row["date_from"].isoformat() if row["date_from"] else None, "date_to": row["date_to"].isoformat() if row["date_to"] else None, "model": row["model"], "inbox_zip_name": row["inbox_zip_name"], "sent_zip_name": row["sent_zip_name"], "inbox_count": row["inbox_count"], "sent_count": row["sent_count"], "filtered_inbox_count": row["filtered_inbox_count"], "filtered_sent_count": row["filtered_sent_count"], "progress_pct": row.get("progress_pct") or 0, "progress_label": row.get("progress_label") or "", "progress_detail": row.get("progress_detail") or "", "error_message": row["error_message"], "created_at": row["created_at"].isoformat() if row["created_at"] else None, "started_at": row["started_at"].isoformat() if row["started_at"] else None, "completed_at": row["completed_at"].isoformat() if row["completed_at"] else None, } def _row_to_report(row: Any) -> dict[str, Any]: return { "id": row["id"], "job_id": row["job_id"], "model": row["model"], "scope_note": row["scope_note"], "report": row["report"] or {}, "report_markdown": row["report_markdown"], "created_at": row["created_at"].isoformat() if row["created_at"] else None, "updated_at": row["updated_at"].isoformat() if row["updated_at"] else None, } async def ensure_tables() -> None: if not pool_enabled(): return pool = get_pool() async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute(Schema.ENSURE_ANALYSIS_JOBS) await cur.execute(Schema.ADD_PROGRESS_PCT) await cur.execute(Schema.ADD_PROGRESS_LABEL) await cur.execute(Schema.ADD_PROGRESS_DETAIL) await cur.execute(Schema.IDX_ANALYSIS_JOBS_CREATED_AT) await cur.execute(Schema.IDX_ANALYSIS_JOBS_STATUS) await cur.execute(Schema.ENSURE_MAILS) await cur.execute(Schema.IDX_MAILS_JOB_ID) await cur.execute(Schema.IDX_MAILS_JOB_MAILED_AT) await cur.execute(Schema.IDX_MAILS_THREAD_KEY) await cur.execute(Schema.IDX_MAILS_MESSAGE_ID) await cur.execute(Schema.ENSURE_WEEKLY_REPORTS) await cur.execute(Schema.IDX_WEEKLY_REPORTS_CREATED_AT) await conn.commit() logger.info("Ensured mail_summary tables") async def create_job( *, mailbox_scope: str, date_from: date, date_to: date, model: str, inbox_zip_name: str | None, sent_zip_name: str | None, ) -> dict[str, Any]: if not pool_enabled(): raise RuntimeError("database_unavailable") pool = get_pool() async with pool.connection() as conn: async with conn.cursor(row_factory=dict_row) as cur: await cur.execute( Jobs.INSERT, { "mailbox_scope": mailbox_scope, "date_from": date_from, "date_to": date_to, "model": model, "inbox_zip_name": inbox_zip_name, "sent_zip_name": sent_zip_name, }, ) row = await cur.fetchone() await conn.commit() return _row_to_job(row) async def get_job(job_id: int) -> dict[str, Any] | None: if not pool_enabled(): return None pool = get_pool() async with pool.connection() as conn: async with conn.cursor(row_factory=dict_row) as cur: await cur.execute(Jobs.GET_BY_ID, {"job_id": job_id}) row = await cur.fetchone() return _row_to_job(row) if row else None async def list_jobs(limit: int = 20) -> list[dict[str, Any]]: if not pool_enabled(): return [] pool = get_pool() async with pool.connection() as conn: async with conn.cursor(row_factory=dict_row) as cur: await cur.execute(Jobs.LIST_RECENT, {"limit": limit}) rows = await cur.fetchall() return [_row_to_job(row) for row in rows] async def update_job( job_id: int, *, status: str | None = None, inbox_count: int | None = None, sent_count: int | None = None, filtered_inbox_count: int | None = None, filtered_sent_count: int | None = None, progress_pct: int | None = None, progress_label: str | None = None, progress_detail: str | None = None, error_message: str | None = None, started: bool = False, completed: bool = False, ) -> dict[str, Any] | None: if not pool_enabled(): return None pool = get_pool() async with pool.connection() as conn: async with conn.cursor(row_factory=dict_row) as cur: await cur.execute( Jobs.UPDATE, { "job_id": job_id, "status": status, "inbox_count": inbox_count, "sent_count": sent_count, "filtered_inbox_count": filtered_inbox_count, "filtered_sent_count": filtered_sent_count, "progress_pct": progress_pct, "progress_label": progress_label, "progress_detail": progress_detail, "error_message": error_message, "set_started": started, "set_completed": completed, }, ) row = await cur.fetchone() await conn.commit() return _row_to_job(row) if row else None async def insert_mails(job_id: int, mails: list[ParsedMail]) -> int: if not pool_enabled() or not mails: return 0 pool = get_pool() count = 0 async with pool.connection() as conn: async with conn.cursor() as cur: for mail in mails: await cur.execute( Mails.INSERT, { "job_id": job_id, "mailbox": mail.mailbox, "message_id": mail.message_id, "thread_key": mail.thread_key, "subject": mail.subject, "from_addr": mail.from_addr, "to_addrs": mail.to_addrs, "cc_addrs": mail.cc_addrs, "mailed_at": mail.mailed_at, "attachment_names": Jsonb(mail.attachment_names), "body_text": mail.body_text, }, ) count += 1 await conn.commit() return count async def list_mails_for_job(job_id: int) -> list[dict[str, Any]]: if not pool_enabled(): return [] pool = get_pool() async with pool.connection() as conn: async with conn.cursor(row_factory=dict_row) as cur: await cur.execute(Mails.LIST_BY_JOB, {"job_id": job_id}) rows = await cur.fetchall() result: list[dict[str, Any]] = [] for row in rows: result.append( { "id": row["id"], "job_id": row["job_id"], "mailbox": row["mailbox"], "message_id": row["message_id"], "thread_key": row["thread_key"], "subject": row["subject"], "from_addr": row["from_addr"], "to_addrs": row["to_addrs"], "cc_addrs": row["cc_addrs"], "mailed_at": row["mailed_at"].isoformat() if row["mailed_at"] else None, "attachment_names": row["attachment_names"] or [], "body_text": row["body_text"], "body_purged": row["body_purged"], "summary": row["summary"], } ) return result async def update_mail_summary(mail_id: int, summary: dict[str, Any]) -> None: if not pool_enabled(): return pool = get_pool() async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute( Mails.UPDATE_SUMMARY, { "mail_id": mail_id, "summary": Jsonb(summary), }, ) await conn.commit() async def purge_mail_bodies(job_id: int) -> None: if not pool_enabled(): return pool = get_pool() async with pool.connection() as conn: async with conn.cursor() as cur: await cur.execute(Mails.PURGE_BODIES, {"job_id": job_id}) await conn.commit() async def get_report(job_id: int) -> dict[str, Any] | None: if not pool_enabled(): return None pool = get_pool() async with pool.connection() as conn: async with conn.cursor(row_factory=dict_row) as cur: await cur.execute(Reports.GET_BY_JOB, {"job_id": job_id}) row = await cur.fetchone() return _row_to_report(row) if row else None async def save_weekly_report( *, job_id: int, model: str, scope_note: str, report: dict[str, Any], report_markdown: str, ) -> dict[str, Any]: if not pool_enabled(): raise RuntimeError("database_unavailable") pool = get_pool() async with pool.connection() as conn: async with conn.cursor(row_factory=dict_row) as cur: await cur.execute( Reports.UPSERT, { "job_id": job_id, "model": model, "scope_note": scope_note, "report": Jsonb(report), "report_markdown": report_markdown, }, ) row = await cur.fetchone() await conn.commit() return _row_to_report(row)