itqop
/
dataloader
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: fixes

This commit is contained in:
itqop 2025-11-05 02:34:15 +03:00
parent fbcdbac6a0
commit 1789270d17
2 changed files with 17 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
src/dataloader/context.py
View File

@ -8,7 +8,7 @@ from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
from dataloader.base import Singleton
from dataloader.config import APP_CONFIG, Secrets
from dataloader.logger import ContextVarsContainer, LoggerConfigurator
from sqlalchemy import event
from sqlalchemy import event, select, func, text
class AppContext(metaclass=Singleton):
@ -100,6 +100,11 @@ class AppContext(metaclass=Singleton):
cur.close()
self._sessionmaker = async_sessionmaker(self._engine, expire_on_commit=False, class_=AsyncSession)
async with self._sessionmaker() as s:
await s.execute(select(func.count()).select_from(text("dl_jobs")))
await s.execute(select(func.count()).select_from(text("dl_job_events")))
self.logger.info("All connections checked. Application is up and ready.")
async def on_shutdown(self) -> None:

23
src/dataloader/storage/repositories.py
View File

@ -5,8 +5,8 @@ from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Any, Optional
from sqlalchemy import BigInteger, String, Text, select, func, update
from sqlalchemy.dialects.postgresql import JSONB, ENUM
from sqlalchemy import BigInteger, String, Text, select, func, update, DateTime
from sqlalchemy.dialects.postgresql import JSONB, ENUM, UUID
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import Mapped, mapped_column
@ -32,7 +32,7 @@ class DLJob(Base):
"""
__tablename__ = "dl_jobs"
job_id: Mapped[str] = mapped_column(String(36), primary_key=True)
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)
@ -40,21 +40,21 @@ class DLJob(Base):
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(nullable=False)
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(nullable=True)
heartbeat_at: Mapped[Optional[datetime]] = mapped_column(nullable=True)
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(nullable=False)
started_at: Mapped[Optional[datetime]] = mapped_column(nullable=True)
finished_at: Mapped[Optional[datetime]] = mapped_column(nullable=True)
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):
@ -64,9 +64,9 @@ class DLJobEvent(Base):
__tablename__ = "dl_job_events"
event_id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True)
job_id: Mapped[str] = mapped_column(String(36), nullable=False)
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(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)
@ -314,7 +314,6 @@ class QueueRepository:
job.lease_expires_at = None
ids.append(job.job_id)
await self._append_event(job.job_id, job.queue, "requeue_lost", None)
await self.s.commit()
return ids
async def _get(self, job_id: str) -> Optional[DLJob]: