"""Download-related Celery tasks.""" import asyncio import logging from datetime import datetime, timedelta from typing import Optional from sqlalchemy import select, and_, or_ from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker from sqlalchemy.orm import selectinload from app.config import get_settings from app.tasks.celery_app import celery_app from app.models.base import utcnow from app.models.source import Source from app.models.subscription import Subscription from app.models.download import Download, DownloadStatus, ErrorType as DBErrorType from app.services.gallery_dl import GalleryDLService, SourceConfig, DownloadResult from app.services.credential_manager import CredentialManager from app.events import publish_event logger = logging.getLogger(__name__) # Settings for database connection settings = get_settings() def get_async_session(): """Get async session factory for Celery tasks. Creates a fresh engine and session factory each time to avoid event loop issues when asyncio.run() creates new loops. """ engine = create_async_engine(settings.async_database_url, echo=False) session_factory = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) return session_factory, engine async def _cleanup_engine(engine): """Properly dispose of the async engine.""" await engine.dispose() async def _download_source_async(source_id: int, download_id: int = None) -> dict: """Async implementation of source download. Args: source_id: ID of the source to download download_id: Optional ID of a pre-created Download record (QUEUED state) Returns: Dict with download results """ session_factory, engine = get_async_session() try: async with session_factory() as session: # Get source with subscription result = await session.execute( select(Source) .options(selectinload(Source.subscription)) .where(Source.id == source_id) ) source = result.scalar() if not source: logger.error(f"Source {source_id} not found") return {"source_id": source_id, "status": "error", "error": "Source not found"} if not source.enabled: logger.info(f"Source {source.subscription.name}/{source.platform} is disabled, skipping") # If there's a queued download record, mark it as skipped if download_id: dl_result = await session.execute( select(Download).where(Download.id == download_id) ) dl = dl_result.scalar() if dl: dl.status = DownloadStatus.SKIPPED await session.commit() return {"source_id": source_id, "status": "skipped", "reason": "disabled"} subscription_name = source.subscription.name logger.info(f"Starting download for {subscription_name}/{source.platform}") # Use existing download record or create a new one if download_id: result = await session.execute( select(Download).where(Download.id == download_id) ) download = result.scalar() if download: download.status = DownloadStatus.RUNNING download.started_at = utcnow() else: logger.warning(f"Download {download_id} not found, creating new record") download = None if not download_id or download is None: download = Download( source_id=source.id, url=source.url, status=DownloadStatus.RUNNING, started_at=utcnow(), ) session.add(download) await session.commit() await session.refresh(download) # Broadcast download started event publish_event("download.started", { "download_id": download.id, "source_id": source.id, "subscription_name": subscription_name, "platform": source.platform, "url": source.url, }) try: # Get credentials cred_manager = CredentialManager(session) cookies_path = None auth_token = None # Discord uses token auth, others use cookies if source.platform == "discord": auth_token = await cred_manager.get_token(source.platform) if not auth_token: logger.warning(f"No token for {source.platform}, download may fail") else: cookies_path = await cred_manager.get_cookies_path(source.platform) if not cookies_path and source.platform in ["patreon", "subscribestar"]: logger.warning(f"No credentials for {source.platform}, download may fail") # Build source config from metadata source_config = SourceConfig.from_dict(source.metadata_ or {}) # Execute download gdl_service = GalleryDLService() dl_result = await gdl_service.download( url=source.url, subscription_name=subscription_name, platform=source.platform, source_config=source_config, cookies_path=cookies_path, auth_token=auth_token, ) # Update download record download.completed_at = utcnow() download.file_count = dl_result.files_downloaded download.metadata_ = { "stdout": dl_result.stdout[:10000] if dl_result.stdout else None, # Truncate "stderr": dl_result.stderr[:5000] if dl_result.stderr else None, "duration_seconds": dl_result.duration_seconds, } if dl_result.success: download.status = DownloadStatus.COMPLETED source.last_success = utcnow() source.error_count = 0 logger.info(f"Download completed for {subscription_name}/{source.platform}: {dl_result.files_downloaded} files") # Broadcast download completed event publish_event("download.completed", { "download_id": download.id, "source_id": source.id, "subscription_name": subscription_name, "platform": source.platform, "file_count": dl_result.files_downloaded, "duration_seconds": dl_result.duration_seconds, }) else: download.status = DownloadStatus.FAILED download.error_type = dl_result.error_type.value if dl_result.error_type else None download.error_message = dl_result.error_message source.error_count = (source.error_count or 0) + 1 logger.warning(f"Download failed for {subscription_name}/{source.platform}: {dl_result.error_message}") # Broadcast download failed event publish_event("download.failed", { "download_id": download.id, "source_id": source.id, "subscription_name": subscription_name, "platform": source.platform, "error_type": dl_result.error_type.value if dl_result.error_type else None, "error_message": dl_result.error_message, }) # Update source last_check source.last_check = utcnow() await session.commit() return { "source_id": source_id, "download_id": download.id, "status": "completed" if dl_result.success else "failed", "files_downloaded": dl_result.files_downloaded, "error_type": dl_result.error_type.value if dl_result.error_type else None, "error_message": dl_result.error_message, "duration_seconds": dl_result.duration_seconds, } except Exception as e: logger.exception(f"Unexpected error downloading {subscription_name}/{source.platform}: {e}") download.status = DownloadStatus.FAILED download.error_type = DBErrorType.UNKNOWN_ERROR download.error_message = str(e) download.completed_at = utcnow() source.error_count = (source.error_count or 0) + 1 source.last_check = utcnow() await session.commit() # Broadcast download failed event publish_event("download.failed", { "download_id": download.id, "source_id": source.id, "subscription_name": subscription_name, "platform": source.platform, "error_type": "unknown_error", "error_message": str(e), }) return { "source_id": source_id, "download_id": download.id, "status": "error", "error": str(e), } finally: await _cleanup_engine(engine) async def _scheduled_check_async() -> dict: """Async implementation of scheduled source check. Returns: Dict with check results """ session_factory, engine = get_async_session() try: async with session_factory() as session: now = utcnow() # Find sources due for checking: # - enabled = True # - last_check is NULL OR (now - last_check) > check_interval # Order by subscription priority, then by when last checked query = ( select(Source) .options(selectinload(Source.subscription)) .join(Subscription) .where( and_( Source.enabled == True, Subscription.enabled == True, # Subscription must also be enabled or_( Source.last_check == None, Source.last_check < now - timedelta(seconds=1) # Will be refined below ) ) ) .order_by( Subscription.priority.desc(), # Higher subscription priority first Source.last_check.asc().nullsfirst(), # Never checked first ) ) result = await session.execute(query) sources = result.scalars().all() # Filter by check_interval due_sources = [] for source in sources: if source.last_check is None: due_sources.append(source) else: time_since_check = (now - source.last_check).total_seconds() if time_since_check >= source.check_interval: due_sources.append(source) logger.info(f"Found {len(due_sources)} sources due for checking") # Get source IDs that already have a queued or running download due_source_ids = [s.id for s in due_sources] active_result = await session.execute( select(Download.source_id).where( Download.source_id.in_(due_source_ids), Download.status.in_([DownloadStatus.QUEUED, DownloadStatus.RUNNING]) ) ) active_source_ids = set(active_result.scalars().all()) # Create download records and queue tasks for each eligible source queued = 0 skipped = 0 for source in due_sources: if source.id in active_source_ids: skipped += 1 logger.debug(f"Skipping {source.subscription.name}/{source.platform} - already active") continue download = Download( source_id=source.id, url=source.url, status=DownloadStatus.QUEUED, ) session.add(download) await session.commit() await session.refresh(download) download_source.delay(source.id, download_id=download.id) queued += 1 logger.debug(f"Queued download for {source.subscription.name}/{source.platform}") if skipped: logger.info(f"Skipped {skipped} sources with active downloads") return { "checked": len(sources), "queued": queued, "timestamp": now.isoformat(), } finally: await _cleanup_engine(engine) async def _cleanup_old_downloads_async() -> dict: """Async implementation of download cleanup. Returns: Dict with cleanup results """ session_factory, engine = get_async_session() try: async with session_factory() as session: now = utcnow() # Delete completed downloads older than 30 days completed_cutoff = now - timedelta(days=30) completed_query = select(Download).where( and_( Download.status == DownloadStatus.COMPLETED, Download.created_at < completed_cutoff ) ) completed_result = await session.execute(completed_query) completed_old = completed_result.scalars().all() # Delete failed downloads older than 7 days failed_cutoff = now - timedelta(days=7) failed_query = select(Download).where( and_( Download.status == DownloadStatus.FAILED, Download.created_at < failed_cutoff ) ) failed_result = await session.execute(failed_query) failed_old = failed_result.scalars().all() deleted_count = 0 for download in completed_old + failed_old: await session.delete(download) deleted_count += 1 await session.commit() logger.info(f"Cleaned up {deleted_count} old download records") return {"deleted": deleted_count} finally: await _cleanup_engine(engine) # Celery tasks that wrap async functions @celery_app.task(bind=True, max_retries=3) def download_source(self, source_id: int, download_id: int = None): """Download new content from a single source. Args: source_id: ID of the source to download from download_id: Optional ID of a pre-created Download record (QUEUED state) """ try: return asyncio.run(_download_source_async(source_id, download_id=download_id)) except Exception as e: logger.exception(f"Task failed for source {source_id}: {e}") # Retry with exponential backoff raise self.retry(exc=e, countdown=60 * (2 ** self.request.retries)) @celery_app.task def scheduled_check(): """Check all enabled sources that are due for a check. This is called periodically by Celery Beat. """ return asyncio.run(_scheduled_check_async()) @celery_app.task def cleanup_old_downloads(): """Clean up old download records. Keeps the database from growing indefinitely. """ return asyncio.run(_cleanup_old_downloads_async()) @celery_app.task(bind=True, max_retries=2) def process_download(self, download_id: int): """Process a single download record (for retries). Args: download_id: ID of the download to process """ async def _process(): session_factory, engine = get_async_session() try: async with session_factory() as session: result = await session.execute( select(Download).where(Download.id == download_id) ) download = result.scalar() if not download: return {"download_id": download_id, "status": "error", "error": "Download not found"} if not download.source_id: return {"download_id": download_id, "status": "error", "error": "No source associated"} # Delegate to download_source return await _download_source_async(download.source_id) finally: await _cleanup_engine(engine) try: return asyncio.run(_process()) except Exception as e: logger.exception(f"Task failed for download {download_id}: {e}") raise self.retry(exc=e, countdown=60 * (2 ** self.request.retries))