IPSExtractor/services/cleaning/api/main.py

"""API REST del microservicio de limpieza."""
from __future__ import annotations

from contextlib import asynccontextmanager
import json
from typing import Any

import pandas as pd
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from sqlalchemy import delete, select

from domain.csv_cleaning import CLEANING_RULES, analyze_dataframe, apply_rules
from infrastructure.database import (
  CleanedRow,
  CleaningJob,
  ExtractionBatch,
  ExtractionRow,
  init_db,
  session_scope,
  utcnow,
)


@asynccontextmanager
async def lifespan(app: FastAPI):
  init_db()
  yield


app = FastAPI(
  title="IPSExtractor Cleaning Service",
  version="1.0.0",
  lifespan=lifespan,
)


class CleaningRequest(BaseModel):
  extraction_batch_id: str
  rules_per_column: dict[str, list[str]] = Field(default_factory=dict)
  cop_mod_fecha_column: dict[str, str] = Field(default_factory=dict)
  cop_mod_config: dict[str, Any] | None = None


def dataframe_to_records(df: pd.DataFrame) -> list[dict[str, Any]]:
  return json.loads(df.to_json(orient="records", date_format="iso"))


def serialize_job(job: CleaningJob) -> dict[str, Any]:
  return {
    "id": job.id,
    "extraction_batch_id": job.extraction_batch_id,
    "status": job.status,
    "request": job.request,
    "log": job.log,
    "error": job.error,
    "created_at": job.created_at.isoformat(),
    "updated_at": job.updated_at.isoformat(),
  }


@app.get("/health")
def health() -> dict[str, str]:
  return {"status": "ok"}


@app.get("/rules")
def list_rules() -> list[dict[str, Any]]:
  return [
    {
      "id": rule["id"],
      "name": rule["name"],
      "description": rule.get("description", ""),
    }
    for rule in CLEANING_RULES
  ]


@app.get("/batches/{batch_id}/analysis")
def analyze_batch(batch_id: str) -> dict[str, Any]:
  with session_scope() as session:
    batch = session.get(ExtractionBatch, batch_id)
    if batch is None:
      raise HTTPException(status_code=404, detail="Lote de extraccion no encontrado")
    rows = session.execute(
      select(ExtractionRow.row_json)
      .where(ExtractionRow.batch_id == batch_id)
      .order_by(ExtractionRow.row_index, ExtractionRow.id)
    ).all()
    records = [row[0] for row in rows]

  df = pd.DataFrame(records)
  return {
    "batch_id": batch_id,
    "columns": [summary.__dict__ for summary in analyze_dataframe(df)],
  }


@app.post("/jobs")
def run_cleaning(request: CleaningRequest) -> dict[str, Any]:
  with session_scope() as session:
    batch = session.get(ExtractionBatch, request.extraction_batch_id)
    if batch is None:
      raise HTTPException(status_code=404, detail="Lote de extraccion no encontrado")

    extracted_rows = session.execute(
      select(ExtractionRow.row_json)
      .where(ExtractionRow.batch_id == request.extraction_batch_id)
      .order_by(ExtractionRow.row_index, ExtractionRow.id)
    ).all()
    records = [row[0] for row in extracted_rows]
    if not records:
      raise HTTPException(status_code=400, detail="El lote no tiene filas extraidas")

    job = CleaningJob(
      extraction_batch_id=request.extraction_batch_id,
      status="running",
      request=request.model_dump(),
      updated_at=utcnow(),
    )
    session.add(job)
    session.flush()
    job_id = job.id

  try:
    df = pd.DataFrame(records)
    result_df, log = apply_rules(
      df,
      request.rules_per_column,
      cop_mod_fecha_column=request.cop_mod_fecha_column,
      cop_mod_config=request.cop_mod_config,
    )
    output_rows = dataframe_to_records(result_df)

    with session_scope() as session:
      session.execute(delete(CleanedRow).where(CleanedRow.cleaning_job_id == job_id))
      for idx, row in enumerate(output_rows, start=1):
        session.add(
          CleanedRow(
            cleaning_job_id=job_id,
            extraction_batch_id=request.extraction_batch_id,
            row_index=idx,
            row_json=row,
          )
        )
      job = session.get(CleaningJob, job_id)
      if job is None:
        raise RuntimeError("No se encontro el job de limpieza recien creado.")
      job.status = "completed"
      job.log = log
      job.updated_at = utcnow()
      session.flush()
      return serialize_job(job)
  except Exception as exc:
    with session_scope() as session:
      job = session.get(CleaningJob, job_id)
      if job is not None:
        job.status = "failed"
        job.error = str(exc)
        job.updated_at = utcnow()
    raise HTTPException(status_code=500, detail=str(exc)) from exc


@app.get("/jobs")
def list_jobs() -> list[dict[str, Any]]:
  with session_scope() as session:
    jobs = session.scalars(
      select(CleaningJob).order_by(CleaningJob.created_at.desc())
    ).all()
    return [serialize_job(job) for job in jobs]


@app.get("/jobs/{job_id}")
def get_job(job_id: str) -> dict[str, Any]:
  with session_scope() as session:
    job = session.get(CleaningJob, job_id)
    if job is None:
      raise HTTPException(status_code=404, detail="Job de limpieza no encontrado")
    return serialize_job(job)


@app.get("/jobs/{job_id}/rows")
def get_cleaned_rows(job_id: str, limit: int = 500, offset: int = 0) -> dict[str, Any]:
  with session_scope() as session:
    job = session.get(CleaningJob, job_id)
    if job is None:
      raise HTTPException(status_code=404, detail="Job de limpieza no encontrado")
    rows = session.scalars(
      select(CleanedRow)
      .where(CleanedRow.cleaning_job_id == job_id)
      .order_by(CleanedRow.row_index, CleanedRow.id)
      .offset(offset)
      .limit(limit)
    ).all()
    return {
      "job_id": job_id,
      "rows": [row.row_json for row in rows],
      "limit": limit,
      "offset": offset,
    }