Files
FabledCurator/backend/app/tasks/download.py
T

89 lines
3.2 KiB
Python

"""download_source Celery task — runs DownloadService for one source."""
import asyncio
from pathlib import Path
from sqlalchemy import select
from sqlalchemy.exc import DBAPIError, OperationalError
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from ..celery_app import celery
from ..config import get_config
from ..models import ImportSettings
from ..services.credential_crypto import CredentialCrypto
from ..services.credential_service import CredentialService
from ..services.download_service import DownloadService
from ..services.gallery_dl import GalleryDLService
from ..services.importer import Importer
from ..services.thumbnailer import Thumbnailer
from .import_file import _sync_session_factory
IMAGES_ROOT = Path("/images")
_KEY_PATH = IMAGES_ROOT / "secrets" / "credential_key.b64"
def _async_session_factory():
cfg = get_config()
engine = create_async_engine(cfg.database_url, future=True, pool_pre_ping=True)
return async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False), engine
@celery.task(
name="backend.app.tasks.download.download_source",
bind=True,
acks_late=True,
autoretry_for=(OperationalError, DBAPIError, OSError),
retry_backoff=10,
retry_backoff_max=120,
retry_jitter=True,
max_retries=3,
soft_time_limit=900,
time_limit=1200,
)
def download_source(self, source_id: int) -> int:
"""Returns the DownloadEvent.id."""
async def _run():
async_factory, async_engine = _async_session_factory()
SyncFactory = _sync_session_factory()
try:
with SyncFactory() as sync_session:
settings = sync_session.execute(
select(ImportSettings).where(ImportSettings.id == 1)
).scalar_one()
rate_limit = settings.download_rate_limit_seconds
validate_files = settings.download_validate_files
gdl = GalleryDLService(
images_root=IMAGES_ROOT,
rate_limit=rate_limit,
validate_files=validate_files,
)
crypto = CredentialCrypto(_KEY_PATH)
async with async_factory() as async_session:
cred_service = CredentialService(async_session, crypto)
with SyncFactory() as sync_session:
sync_settings = sync_session.execute(
select(ImportSettings).where(ImportSettings.id == 1)
).scalar_one()
importer = Importer(
session=sync_session,
images_root=IMAGES_ROOT,
import_root=IMAGES_ROOT,
thumbnailer=Thumbnailer(images_root=IMAGES_ROOT),
settings=sync_settings,
)
svc = DownloadService(
async_session=async_session,
sync_session=sync_session,
gdl=gdl,
importer=importer,
cred_service=cred_service,
)
return await svc.download_source(source_id)
finally:
await async_engine.dispose()
return asyncio.run(_run())