128 lines
4.2 KiB
Python
128 lines
4.2 KiB
Python
"""Persistencia del servicio sobre la base de datos de integracion."""
|
|
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,
|
|
select,
|
|
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 ExtractionBatch(Base):
|
|
__tablename__ = "extraction_batches"
|
|
|
|
id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_id)
|
|
status: Mapped[str] = mapped_column(String(32), default="pending", index=True)
|
|
root_path: Mapped[str | None] = mapped_column(Text, nullable=True)
|
|
request: 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 ExtractionRow(Base):
|
|
__tablename__ = "extraction_rows"
|
|
|
|
id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
|
|
batch_id: Mapped[str] = mapped_column(ForeignKey("extraction_batches.id"), index=True)
|
|
radicado: Mapped[str | None] = mapped_column(String(255), nullable=True, index=True)
|
|
source_pdf: Mapped[str | None] = mapped_column(Text, nullable=True)
|
|
row_index: Mapped[int] = mapped_column(Integer, default=0)
|
|
row_json: Mapped[dict[str, Any]] = mapped_column(JSON, default=dict)
|
|
raw_json: Mapped[dict[str, Any]] = mapped_column(JSON, default=dict)
|
|
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
|
|
|
|
|
|
class CleaningJob(Base):
|
|
__tablename__ = "cleaning_jobs"
|
|
|
|
id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_id)
|
|
extraction_batch_id: Mapped[str] = mapped_column(ForeignKey("extraction_batches.id"), 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] = mapped_column(ForeignKey("extraction_batches.id"), 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.begin() as connection:
|
|
connection.execute(text("SELECT pg_advisory_lock(74512001)"))
|
|
try:
|
|
Base.metadata.create_all(bind=connection)
|
|
finally:
|
|
connection.execute(text("SELECT pg_advisory_unlock(74512001)"))
|
|
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()
|
|
|
|
|
|
def get_extraction_rows(session: Session, batch_id: str) -> list[ExtractionRow]:
|
|
return list(
|
|
session.scalars(
|
|
select(ExtractionRow)
|
|
.where(ExtractionRow.batch_id == batch_id)
|
|
.order_by(ExtractionRow.row_index, ExtractionRow.id)
|
|
)
|
|
)
|