160 lines
5.6 KiB
Python
160 lines
5.6 KiB
Python
"""Capa de persistencia SQLite para el servidor de tracking.
|
|
|
|
Guarda el historial completo de lotes de extracción, incluyendo
|
|
métricas de progreso, costos y la ruta del CSV generado.
|
|
"""
|
|
import sqlite3
|
|
import json
|
|
import os
|
|
from typing import Dict, Any, List, Optional
|
|
from datetime import datetime
|
|
|
|
DB_PATH = os.path.join(os.path.dirname(__file__), "..", "tracking.db")
|
|
|
|
|
|
def get_connection() -> sqlite3.Connection:
|
|
conn = sqlite3.connect(DB_PATH)
|
|
conn.row_factory = sqlite3.Row
|
|
return conn
|
|
|
|
|
|
def init_db():
|
|
"""Crea las tablas si no existen."""
|
|
with get_connection() as conn:
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS batch_logs (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
batch_id TEXT NOT NULL,
|
|
timestamp TEXT NOT NULL,
|
|
level TEXT NOT NULL DEFAULT 'info',
|
|
radicado TEXT,
|
|
message TEXT NOT NULL,
|
|
FOREIGN KEY (batch_id) REFERENCES batches(id)
|
|
)
|
|
""")
|
|
conn.execute("CREATE INDEX IF NOT EXISTS idx_batch_logs_batch_id ON batch_logs(batch_id)")
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS batches (
|
|
id TEXT PRIMARY KEY,
|
|
root_path TEXT NOT NULL,
|
|
csv_output_path TEXT,
|
|
hostname TEXT,
|
|
username TEXT,
|
|
total_radicados INTEGER NOT NULL DEFAULT 0,
|
|
config TEXT,
|
|
status TEXT NOT NULL DEFAULT 'processing',
|
|
current_radicado_index INTEGER NOT NULL DEFAULT 0,
|
|
success_count INTEGER NOT NULL DEFAULT 0,
|
|
error_count INTEGER NOT NULL DEFAULT 0,
|
|
total_tokens INTEGER NOT NULL DEFAULT 0,
|
|
total_cost_usd REAL NOT NULL DEFAULT 0.0,
|
|
last_error TEXT,
|
|
current_phase TEXT,
|
|
current_radicado TEXT,
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL
|
|
)
|
|
""")
|
|
conn.commit()
|
|
|
|
|
|
def insert_batch(batch: Dict[str, Any]):
|
|
"""Inserta un nuevo lote en la base de datos."""
|
|
with get_connection() as conn:
|
|
conn.execute("""
|
|
INSERT INTO batches (
|
|
id, root_path, csv_output_path, hostname, username, total_radicados, config,
|
|
status, current_radicado_index, success_count, error_count,
|
|
total_tokens, total_cost_usd, last_error, current_phase, current_radicado,
|
|
created_at, updated_at
|
|
) VALUES (
|
|
:id, :root_path, :csv_output_path, :hostname, :username, :total_radicados, :config,
|
|
:status, :current_radicado_index, :success_count, :error_count,
|
|
:total_tokens, :total_cost_usd, :last_error, :current_phase, :current_radicado,
|
|
:created_at, :updated_at
|
|
)
|
|
""", {
|
|
**batch,
|
|
"config": json.dumps(batch.get("config") or {}),
|
|
"created_at": _to_iso(batch.get("created_at")),
|
|
"updated_at": _to_iso(batch.get("updated_at")),
|
|
"current_phase": batch.get("current_phase"),
|
|
"current_radicado": batch.get("current_radicado"),
|
|
})
|
|
conn.commit()
|
|
|
|
|
|
def update_batch(batch_id: str, fields: Dict[str, Any]):
|
|
"""Actualiza campos de un lote existente."""
|
|
if not fields:
|
|
return
|
|
fields["updated_at"] = _to_iso(datetime.now())
|
|
set_clause = ", ".join(f"{k} = :{k}" for k in fields)
|
|
fields["batch_id"] = batch_id
|
|
with get_connection() as conn:
|
|
conn.execute(
|
|
f"UPDATE batches SET {set_clause} WHERE id = :batch_id",
|
|
fields
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
def get_all_batches() -> List[Dict[str, Any]]:
|
|
"""Retorna todos los lotes ordenados por fecha de creación descendente."""
|
|
with get_connection() as conn:
|
|
rows = conn.execute(
|
|
"SELECT * FROM batches ORDER BY created_at DESC"
|
|
).fetchall()
|
|
return [_row_to_dict(r) for r in rows]
|
|
|
|
|
|
def insert_log(batch_id: str, message: str, level: str = "info", radicado: str = None):
|
|
"""Inserta una entrada de log para un lote."""
|
|
with get_connection() as conn:
|
|
conn.execute(
|
|
"INSERT INTO batch_logs (batch_id, timestamp, level, radicado, message) VALUES (?, ?, ?, ?, ?)",
|
|
(batch_id, datetime.now().isoformat(), level, radicado, message)
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
def get_logs(batch_id: str, limit: int = 200) -> List[Dict[str, Any]]:
|
|
"""Retorna los logs de un lote, del más reciente al más antiguo."""
|
|
with get_connection() as conn:
|
|
rows = conn.execute(
|
|
"SELECT * FROM batch_logs WHERE batch_id = ? ORDER BY id DESC LIMIT ?",
|
|
(batch_id, limit)
|
|
).fetchall()
|
|
return [dict(r) for r in rows]
|
|
|
|
|
|
def get_batch(batch_id: str) -> Optional[Dict[str, Any]]:
|
|
"""Retorna un lote por su ID."""
|
|
with get_connection() as conn:
|
|
row = conn.execute(
|
|
"SELECT * FROM batches WHERE id = ?", (batch_id,)
|
|
).fetchone()
|
|
return _row_to_dict(row) if row else None
|
|
|
|
|
|
def _row_to_dict(row: sqlite3.Row) -> Dict[str, Any]:
|
|
d = dict(row)
|
|
try:
|
|
d["config"] = json.loads(d.get("config") or "{}")
|
|
except Exception:
|
|
d["config"] = {}
|
|
|
|
# Asegurar que hostname y username estén presentes
|
|
if "hostname" not in d: d["hostname"] = None
|
|
if "username" not in d: d["username"] = None
|
|
|
|
return d
|
|
|
|
|
|
def _to_iso(value) -> str:
|
|
if isinstance(value, datetime):
|
|
return value.isoformat()
|
|
if isinstance(value, str):
|
|
return value
|
|
return datetime.now().isoformat()
|