79 lines
3.4 KiB
Python
79 lines
3.4 KiB
Python
# src/dataloader/storage/models.py
|
|
from __future__ import annotations
|
|
|
|
from datetime import datetime
|
|
from typing import Any, Optional
|
|
|
|
from sqlalchemy import BigInteger, DateTime, Text
|
|
from sqlalchemy.dialects.postgresql import ENUM, JSONB, UUID
|
|
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
"""
|
|
Базовый класс для всех ORM моделей.
|
|
"""
|
|
pass
|
|
|
|
|
|
dl_status_enum = ENUM(
|
|
"queued",
|
|
"running",
|
|
"succeeded",
|
|
"failed",
|
|
"canceled",
|
|
"lost",
|
|
name="dl_status",
|
|
create_type=False,
|
|
native_enum=True,
|
|
)
|
|
|
|
|
|
class DLJob(Base):
|
|
"""
|
|
Модель таблицы очереди задач dl_jobs.
|
|
Использует логическое имя схемы 'queue' для поддержки schema_translate_map.
|
|
"""
|
|
__tablename__ = "dl_jobs"
|
|
__table_args__ = {"schema": "queue"}
|
|
|
|
job_id: Mapped[str] = mapped_column(UUID(as_uuid=False), primary_key=True)
|
|
queue: Mapped[str] = mapped_column(Text, nullable=False)
|
|
task: Mapped[str] = mapped_column(Text, nullable=False)
|
|
args: Mapped[dict[str, Any]] = mapped_column(JSONB, default=dict, nullable=False)
|
|
idempotency_key: Mapped[Optional[str]] = mapped_column(Text, unique=True)
|
|
lock_key: Mapped[str] = mapped_column(Text, nullable=False)
|
|
partition_key: Mapped[str] = mapped_column(Text, default="", nullable=False)
|
|
priority: Mapped[int] = mapped_column(nullable=False, default=100)
|
|
available_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
|
|
status: Mapped[str] = mapped_column(dl_status_enum, nullable=False, default="queued")
|
|
attempt: Mapped[int] = mapped_column(nullable=False, default=0)
|
|
max_attempts: Mapped[int] = mapped_column(nullable=False, default=5)
|
|
lease_ttl_sec: Mapped[int] = mapped_column(nullable=False, default=60)
|
|
lease_expires_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True))
|
|
heartbeat_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True))
|
|
cancel_requested: Mapped[bool] = mapped_column(nullable=False, default=False)
|
|
progress: Mapped[dict[str, Any]] = mapped_column(JSONB, default=dict, nullable=False)
|
|
error: Mapped[Optional[str]] = mapped_column(Text)
|
|
producer: Mapped[Optional[str]] = mapped_column(Text)
|
|
consumer_group: Mapped[Optional[str]] = mapped_column(Text)
|
|
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
|
|
started_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True))
|
|
finished_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True))
|
|
|
|
|
|
class DLJobEvent(Base):
|
|
"""
|
|
Модель таблицы журнала событий dl_job_events.
|
|
Использует логическое имя схемы 'queue' для поддержки schema_translate_map.
|
|
"""
|
|
__tablename__ = "dl_job_events"
|
|
__table_args__ = {"schema": "queue"}
|
|
|
|
event_id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True)
|
|
job_id: Mapped[str] = mapped_column(UUID(as_uuid=False), nullable=False)
|
|
queue: Mapped[str] = mapped_column(Text, nullable=False)
|
|
ts: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
|
|
kind: Mapped[str] = mapped_column(Text, nullable=False)
|
|
payload: Mapped[Optional[dict[str, Any]]] = mapped_column(JSONB)
|