scripts/bootstrap_extract.py

#!/usr/bin/env python3
"""Bootstrap da fundação de dados: baixa o ano corrente inteiro da Apoena.

Meses fechados (jan..mes_atual-1) -> pivot_<base>_<MM>.
Mês corrente -> current_<base> (swap atômico).

Uso:
  ./scripts/bootstrap_extract.py            # ano corrente
  ./scripts/bootstrap_extract.py --year 2026 --months 1-6
"""
from __future__ import annotations

import argparse
import logging
import os
import sys
from datetime import datetime
from pathlib import Path
from zoneinfo import ZoneInfo

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

BRT = ZoneInfo("America/Sao_Paulo")
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", datefmt="%H:%M:%S")
log = logging.getLogger("bootstrap_extract")


def load_env_local(path: Path) -> None:
    if not path.exists():
        return
    for line in path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        k, v = line.split("=", 1)
        os.environ.setdefault(k.strip(), v.strip().strip('"').strip("'"))


def main() -> int:
    load_env_local(ROOT / ".env.local")

    # importa só depois de carregar env (extractor lê REPORT_BASE no import-time)
    import duckdb
    from dashboard_apn.extractor import F2MExtractor, PIVOT_BASES

    now = datetime.now(BRT)
    p = argparse.ArgumentParser()
    p.add_argument("--db", default=os.environ.get("F2M_DB_PATH", str(ROOT / "f2m_local.duckdb")))
    p.add_argument("--year", type=int, default=now.year)
    p.add_argument("--months", help="Ex: 1-6 (default: jan..mes_atual)")
    args = p.parse_args()

    db_path = Path(args.db)
    email = os.environ.get("F2M_EMAIL", "")
    password = os.environ.get("F2M_PASSWORD", "")
    if not email or not password:
        log.error("F2M_EMAIL/F2M_PASSWORD ausentes")
        return 1

    mes_atual = now.month if args.year == now.year else 12
    if args.months:
        a, b = (args.months.split("-", 1) + [None])[:2]
        months = list(range(int(a), int(b) + 1)) if b else [int(a)]
    else:
        months = list(range(1, mes_atual + 1))

    log.info("Bases: %s", list(PIVOT_BASES.values()))
    log.info("DB: %s | ano=%d | meses=%s | mes_corrente=%d", db_path, args.year, months, mes_atual)

    # cria DB vazio se não existir
    duckdb.connect(str(db_path)).close()

    ex = F2MExtractor(str(db_path), email, password)

    for mes in months:
        if args.year == now.year and mes == mes_atual:
            log.info("Mês corrente %02d -> current_*", mes)
            ex.extract_current_month()
        else:
            log.info("Mês fechado %02d -> pivot_*_%02d", mes, mes)
            ex.extract_month(args.year, mes, year_prefix=(args.year != now.year))

    # relatório de contagens
    con = duckdb.connect(str(db_path), read_only=True)
    try:
        tables = sorted(r[0] for r in con.execute("SHOW TABLES").fetchall())
        log.info("=== Tabelas (%d) ===", len(tables))
        for t in tables:
            n = con.execute(f'SELECT COUNT(*) FROM "{t}"').fetchone()[0]
            log.info("  %-40s %8d", t, n)
    finally:
        con.close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())