165 lines
6.4 KiB
Python
165 lines
6.4 KiB
Python
"""Persistencia Postgres del servicio bot-descarga (DB compartida IPSExtractor)."""
|
|
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,
|
|
Boolean,
|
|
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())
|
|
|
|
|
|
def _iso(dt: datetime | None) -> str | None:
|
|
if dt is None:
|
|
return None
|
|
if dt.tzinfo is None:
|
|
dt = dt.replace(tzinfo=timezone.utc)
|
|
return dt.astimezone().isoformat(timespec="seconds")
|
|
|
|
|
|
class Job(Base):
|
|
"""Contenedor Desktop (lectura / FK)."""
|
|
|
|
__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 DownloadWorkerRow(Base):
|
|
__tablename__ = "download_workers"
|
|
|
|
id: Mapped[str] = mapped_column(String(36), primary_key=True)
|
|
hostname: Mapped[str] = mapped_column(Text, default="")
|
|
max_concurrent: Mapped[int] = mapped_column(Integer, default=1)
|
|
status: Mapped[str] = mapped_column(String(32), index=True)
|
|
last_heartbeat: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
|
|
current_job_id: Mapped[str | None] = mapped_column(String(36), nullable=True)
|
|
|
|
|
|
class DownloadJobRow(Base):
|
|
__tablename__ = "download_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
|
|
)
|
|
task_id: Mapped[str | None] = mapped_column(
|
|
String(36), ForeignKey("tasks.id"), nullable=True, index=True
|
|
)
|
|
worker_id: Mapped[str | None] = mapped_column(String(36), nullable=True, index=True)
|
|
parent_download_job_id: Mapped[str | None] = mapped_column(
|
|
String(36), ForeignKey("download_jobs.id"), nullable=True, index=True
|
|
)
|
|
shard_index: Mapped[int | None] = mapped_column(Integer, nullable=True)
|
|
shard_count: Mapped[int | None] = mapped_column(Integer, nullable=True)
|
|
status: Mapped[str] = mapped_column(String(32), default="pending", index=True)
|
|
title: Mapped[str | None] = mapped_column(Text, nullable=True)
|
|
excel_key: Mapped[str] = mapped_column(Text, default="")
|
|
downloads_prefix: Mapped[str] = mapped_column(Text, default="")
|
|
original_filename: Mapped[str] = mapped_column(Text, default="")
|
|
cancel_requested: Mapped[bool] = mapped_column(Boolean, default=False)
|
|
result: 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)
|
|
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 DownloadJobEventRow(Base):
|
|
__tablename__ = "download_job_events"
|
|
|
|
id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
|
|
download_job_id: Mapped[str] = mapped_column(
|
|
String(36), ForeignKey("download_jobs.id"), index=True
|
|
)
|
|
event_type: Mapped[str] = mapped_column(String(64))
|
|
message: Mapped[str] = mapped_column(Text)
|
|
timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
|
|
radicado: Mapped[str | None] = mapped_column(String(255), nullable=True)
|
|
nit: Mapped[str | None] = mapped_column(String(64), nullable=True)
|
|
idx: Mapped[int | None] = mapped_column(Integer, nullable=True)
|
|
total: Mapped[int | None] = mapped_column(Integer, nullable=True)
|
|
metadata_json: Mapped[dict[str, Any]] = mapped_column("metadata", JSON, default=dict)
|
|
|
|
|
|
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(74512004)"))
|
|
try:
|
|
Base.metadata.create_all(bind=connection)
|
|
finally:
|
|
connection.execute(text("SELECT pg_advisory_unlock(74512004)"))
|
|
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()
|