export.py

"""Export collected data as CSV or JSON."""

from __future__ import annotations

import csv
import io
import json
import time
from typing import Any, Literal

import aiosqlite

ExportTable = Literal["events", "tasks", "snapshots", "anomalies"]

# CLI table name -> (SQLite table, timestamp column for --days filter)
EXPORT_TABLES: dict[ExportTable, tuple[str, str]] = {
    "events": ("events", "timestamp_ms"),
    "tasks": ("tasks", "timestamp_ms"),
    "snapshots": ("token_snapshots", "timestamp_ms"),
    "anomalies": ("anomalies", "timestamp_ms"),
}


async def get_table_column_names(
    db: aiosqlite.Connection,
    table: str,
) -> list[str]:
    """Return column names for *table* in schema order."""
    cursor = await db.execute(f"PRAGMA table_info({table})")  # noqa: S608
    rows = await cursor.fetchall()
    return [str(row[1]) for row in rows]


async def fetch_export_rows(
    db: aiosqlite.Connection,
    table: ExportTable,
    *,
    days: int | None = None,
) -> list[dict[str, Any]]:
    """Return all rows for an exportable table, optionally limited to recent days."""
    db_table, ts_col = EXPORT_TABLES[table]
    clauses: list[str] = []
    params: list[int] = []

    if days is not None:
        if days <= 0:
            msg = "--days must be a positive integer"
            raise ValueError(msg)
        cutoff_ms = int((time.time() - days * 86400) * 1000)
        clauses.append(f"{ts_col} >= ?")
        params.append(cutoff_ms)

    where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
    query = (
        f"SELECT * FROM {db_table} {where} ORDER BY {ts_col} ASC"  # noqa: S608
    )
    cursor = await db.execute(query, params)
    rows = await cursor.fetchall()
    return [dict(row) for row in rows]


def format_rows_as_csv(
    rows: list[dict[str, Any]],
    *,
    fieldnames: list[str] | None = None,
) -> str:
    """Serialize *rows* as CSV text."""
    columns = fieldnames or (list(rows[0].keys()) if rows else [])
    buffer = io.StringIO()
    writer = csv.DictWriter(buffer, fieldnames=columns, extrasaction="ignore")
    writer.writeheader()
    for row in rows:
        writer.writerow({key: row.get(key, "") for key in columns})
    return buffer.getvalue()


def format_rows_as_json(rows: list[dict[str, Any]], *, pretty: bool = False) -> str:
    """Serialize *rows* as a JSON array."""
    if pretty:
        return json.dumps(rows, indent=2, default=str) + "\n"
    return json.dumps(rows, default=str) + "\n"