simple-cost-dashboard/collector/db.py

#!/usr/bin/env python3
"""
db.py — Connexion SQLite + écriture des relevés de coûts.

Utilisé par collect.py (collecte journalière) et import_historical.py
(import ponctuel). Le schéma vit dans db/schema.sql.
"""

from __future__ import annotations

import os
import sqlite3
from pathlib import Path

DEFAULT_DB_PATH = "cost_dashboard.db"
SCHEMA_PATH = Path(__file__).parent.parent / "db" / "schema.sql"


def _resolve_db_path(db_path: str | None = None) -> str:
    return db_path or os.environ.get("COST_DASHBOARD_DB") or DEFAULT_DB_PATH


def init_db(db_path: str | None = None) -> None:
    """Crée le fichier SQLite et applique le schéma s'il n'existe pas déjà."""
    path = _resolve_db_path(db_path)
    Path(path).parent.mkdir(parents=True, exist_ok=True)

    conn = sqlite3.connect(path)
    try:
        conn.executescript(SCHEMA_PATH.read_text())
        conn.commit()
    finally:
        conn.close()


def get_connection(db_path: str | None = None) -> sqlite3.Connection:
    """Ouvre une connexion SQLite (foreign keys activées, row factory par nom)."""
    path = _resolve_db_path(db_path)
    conn = sqlite3.connect(path)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA foreign_keys = ON")
    return conn


def upsert_collection(
    conn: sqlite3.Connection,
    provider: str,
    tenant_name: str,
    granularity: str,
    period_start: str,
    period_end: str,
    collected_at: str,
    status: str,
    data: dict | None = None,
    error_message: str | None = None,
) -> int:
    """Insère ou remplace un relevé (provider, tenant_name, granularity,
    period_start) et sa répartition par service. `data` est le dict retourné
    par providers.get_aws_cost/get_azure_cost (+ convert_result) : total,
    currency, by_service, original_total, original_currency, exchange_rate.
    """
    data = data or {}

    conn.execute(
        """
        INSERT INTO collections (
            provider, tenant_name, granularity, period_start, period_end,
            collected_at, status, error_message,
            total_amount, currency, original_total, original_currency, exchange_rate
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ON CONFLICT(provider, tenant_name, granularity, period_start) DO UPDATE SET
            period_end = excluded.period_end,
            collected_at = excluded.collected_at,
            status = excluded.status,
            error_message = excluded.error_message,
            total_amount = excluded.total_amount,
            currency = excluded.currency,
            original_total = excluded.original_total,
            original_currency = excluded.original_currency,
            exchange_rate = excluded.exchange_rate
        """,
        (
            provider, tenant_name, granularity, period_start, period_end,
            collected_at, status, error_message,
            data.get("total"), data.get("currency"),
            data.get("original_total"), data.get("original_currency"), data.get("exchange_rate"),
        ),
    )

    # last_insert_rowid() n'est pas fiable ici : sur la branche ON CONFLICT DO
    # UPDATE (pas de nouvelle ligne créée), SQLite ne le met à jour que si
    # aucun autre INSERT n'a eu lieu sur la connexion depuis — ce qui n'est
    # pas garanti (ex: service_costs d'un tenant précédent). On relit l'id
    # explicitement via la clé unique.
    collection_id = conn.execute(
        """
        SELECT id FROM collections
        WHERE provider = ? AND tenant_name = ? AND granularity = ? AND period_start = ?
        """,
        (provider, tenant_name, granularity, period_start),
    ).fetchone()["id"]

    conn.execute("DELETE FROM service_costs WHERE collection_id = ?", (collection_id,))
    by_service = data.get("by_service") or {}
    if by_service:
        conn.executemany(
            "INSERT INTO service_costs (collection_id, service_name, amount) VALUES (?, ?, ?)",
            [(collection_id, service, amount) for service, amount in by_service.items()],
        )

    conn.commit()
    return collection_id


def set_resource_collection_status(
    conn: sqlite3.Connection,
    provider: str,
    tenant_name: str,
    status: str,
    collected_at: str,
    message: str | None = None,
) -> None:
    conn.execute(
        """
        INSERT INTO resource_collection_status (provider, tenant_name, status, message, collected_at)
        VALUES (?, ?, ?, ?, ?)
        ON CONFLICT(provider, tenant_name) DO UPDATE SET
            status = excluded.status,
            message = excluded.message,
            collected_at = excluded.collected_at
        """,
        (provider, tenant_name, status, message, collected_at),
    )
    conn.commit()


def upsert_resource_costs(
    conn: sqlite3.Connection,
    provider: str,
    tenant_name: str,
    period_start: str,
    period_end: str,
    services: dict,
) -> None:
    """Remplace le cliché de coûts par ressource pour ce tenant (voir
    resource_costs dans le schéma : c'est une photo glissante, pas un
    historique — on supprime l'ancien cliché avant d'écrire le nouveau).
    `services` : {service_name: [{resource_id, resource_name, amount}, ...]},
    la forme renvoyée par providers.get_*_resource_costs.
    """
    conn.execute(
        "DELETE FROM resource_costs WHERE provider = ? AND tenant_name = ?",
        (provider, tenant_name),
    )
    rows = [
        (provider, tenant_name, service, r["resource_id"], r["resource_name"], period_start, period_end, r["amount"])
        for service, resources in services.items()
        for r in resources
    ]
    if rows:
        conn.executemany(
            """
            INSERT INTO resource_costs (
                provider, tenant_name, service_name, resource_id, resource_name,
                period_start, period_end, amount
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """,
            rows,
        )
    conn.commit()