296 lines
10 KiB
Python
296 lines
10 KiB
Python
"""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)
|