IPSExtractor/services/poblacion_apta/infrastructure/database.py

"""Persistencia del servicio Población apta."""
from __future__ import annotations

import os
import uuid
from contextlib import contextmanager
from datetime import datetime, timezone
from typing import Any, Iterator

from sqlalchemy import (
  JSON,
  DateTime,
  ForeignKey,
  Integer,
  String,
  Text,
  create_engine,
  text,
)
from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column, sessionmaker


DATABASE_URL = os.getenv(
  "DATABASE_URL",
  "postgresql+psycopg://ipsextractor:ipsextractor@database:5432/ipsextractor",
)

engine = create_engine(DATABASE_URL, pool_pre_ping=True)
SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False)


class Base(DeclarativeBase):
  pass


def utcnow() -> datetime:
  return datetime.now(timezone.utc)


def new_id() -> str:
  return str(uuid.uuid4())


class Job(Base):
  __tablename__ = "jobs"

  id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_id)
  status: Mapped[str] = mapped_column(String(32), default="pending", index=True)
  title: Mapped[str | None] = mapped_column(Text, nullable=True)
  job_metadata: Mapped[dict[str, Any]] = mapped_column("metadata", JSON, default=dict)
  created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
  updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)


class Task(Base):
  __tablename__ = "tasks"

  id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_id)
  job_id: Mapped[str | None] = mapped_column(
    String(36), ForeignKey("jobs.id"), nullable=True, index=True
  )
  type: Mapped[str] = mapped_column(String(64), index=True)
  status: Mapped[str] = mapped_column(String(32), default="pending", index=True)
  title: Mapped[str | None] = mapped_column(Text, nullable=True)
  request: Mapped[dict[str, Any]] = mapped_column(JSON, default=dict)
  result: Mapped[dict[str, Any]] = mapped_column(JSON, default=dict)
  error: Mapped[str | None] = mapped_column(Text, nullable=True)
  owner_id: Mapped[str | None] = mapped_column(String(36), nullable=True, index=True)
  created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
  updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
  started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
  finished_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)


class CleaningJob(Base):
  __tablename__ = "cleaning_jobs"

  id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_id)
  extraction_batch_id: Mapped[str | None] = mapped_column(String(36), nullable=True, index=True)
  job_id: Mapped[str | None] = mapped_column(
    String(36), ForeignKey("jobs.id"), nullable=True, index=True
  )
  source: Mapped[str] = mapped_column(String(32), default="extraction", index=True)
  status: Mapped[str] = mapped_column(String(32), default="pending", index=True)
  request: Mapped[dict[str, Any]] = mapped_column(JSON, default=dict)
  log: Mapped[list[str]] = mapped_column(JSON, default=list)
  error: Mapped[str | None] = mapped_column(Text, nullable=True)
  created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
  updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)


class CleanedRow(Base):
  __tablename__ = "cleaned_rows"

  id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
  cleaning_job_id: Mapped[str] = mapped_column(ForeignKey("cleaning_jobs.id"), index=True)
  extraction_batch_id: Mapped[str | None] = mapped_column(String(36), nullable=True, index=True)
  row_index: Mapped[int] = mapped_column(Integer, default=0)
  row_json: Mapped[dict[str, Any]] = mapped_column(JSON, default=dict)
  created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)


class PoblacionAptaJob(Base):
  __tablename__ = "poblacion_apta_jobs"

  id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_id)
  job_id: Mapped[str | None] = mapped_column(
    String(36), ForeignKey("jobs.id"), nullable=True, index=True
  )
  status: Mapped[str] = mapped_column(String(32), default="pending", index=True)
  request: Mapped[dict[str, Any]] = mapped_column(JSON, default=dict)
  log: Mapped[list[str]] = mapped_column(JSON, default=list)
  summary: Mapped[dict[str, Any]] = mapped_column(JSON, default=dict)
  error: Mapped[str | None] = mapped_column(Text, nullable=True)
  created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
  updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)


class PoblacionAptaRow(Base):
  __tablename__ = "poblacion_apta_rows"

  id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
  poblacion_apta_job_id: Mapped[str] = mapped_column(
    ForeignKey("poblacion_apta_jobs.id"), index=True
  )
  cleaning_job_id: Mapped[str | None] = mapped_column(String(36), nullable=True, index=True)
  row_index: Mapped[int] = mapped_column(Integer, default=0)
  row_json: Mapped[dict[str, Any]] = mapped_column(JSON, default=dict)
  created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)


def init_db() -> None:
  if DATABASE_URL.startswith("postgresql"):
    with engine.connect().execution_options(isolation_level="AUTOCOMMIT") as connection:
      connection.execute(text("SELECT pg_advisory_lock(74512002)"))
      try:
        Base.metadata.create_all(bind=connection)
      finally:
        connection.execute(text("SELECT pg_advisory_unlock(74512002)"))
    return

  Base.metadata.create_all(bind=engine)


@contextmanager
def session_scope() -> Iterator[Session]:
  session = SessionLocal()
  try:
    yield session
    session.commit()
  except Exception:
    session.rollback()
    raise
  finally:
    session.close()