Bryan Van Deusen
394 lines
15 KiB
Python
394 lines
15 KiB
Python
"""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.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) -> dict:
|
|
"""Async implementation of source download.
|
|
|
|
Args:
|
|
source_id: ID of the source to download
|
|
|
|
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")
|
|
return {"source_id": source_id, "status": "skipped", "reason": "disabled"}
|
|
|
|
subscription_name = source.subscription.name
|
|
logger.info(f"Starting download for {subscription_name}/{source.platform}")
|
|
|
|
# Create download record
|
|
download = Download(
|
|
source_id=source.id,
|
|
url=source.url,
|
|
status=DownloadStatus.RUNNING,
|
|
started_at=datetime.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 = datetime.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 = datetime.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 = datetime.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 = datetime.utcnow()
|
|
source.error_count = (source.error_count or 0) + 1
|
|
source.last_check = datetime.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 = datetime.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")
|
|
|
|
# Queue download tasks for each source
|
|
queued = 0
|
|
for source in due_sources:
|
|
download_source.delay(source.id)
|
|
queued += 1
|
|
logger.debug(f"Queued download for {source.subscription.name}/{source.platform}")
|
|
|
|
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 = datetime.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 new content from a single source.
|
|
|
|
Args:
|
|
source_id: ID of the source to download from
|
|
"""
|
|
try:
|
|
return asyncio.run(_download_source_async(source_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))
|