Bryan Van Deusen
128 lines
3.8 KiB
Python
128 lines
3.8 KiB
Python
"""System maintenance Celery tasks."""
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
|
|
|
|
from app.config import get_settings
|
|
from app.tasks.celery_app import celery_app
|
|
from app.models.setting import Setting, STORAGE_STATS_SETTING
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
settings = get_settings()
|
|
|
|
|
|
def get_async_session():
|
|
"""Get async session factory for Celery tasks."""
|
|
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()
|
|
|
|
|
|
def format_size(size_bytes: int) -> str:
|
|
"""Format bytes as human-readable size."""
|
|
for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
|
|
if size_bytes < 1024:
|
|
return f"{size_bytes:.1f} {unit}"
|
|
size_bytes /= 1024
|
|
return f"{size_bytes:.1f} PB"
|
|
|
|
|
|
def calculate_storage_stats() -> dict:
|
|
"""Calculate storage statistics for the downloads directory.
|
|
|
|
This runs synchronously since it's a filesystem operation.
|
|
"""
|
|
downloads_path = Path(settings.download_path)
|
|
|
|
if not downloads_path.exists():
|
|
return {
|
|
"total_size": 0,
|
|
"total_size_formatted": "0 B",
|
|
"file_count": 0,
|
|
"calculated_at": datetime.utcnow().isoformat(),
|
|
}
|
|
|
|
total_size = 0
|
|
file_count = 0
|
|
|
|
try:
|
|
for root, dirs, files in os.walk(downloads_path):
|
|
for f in files:
|
|
fp = os.path.join(root, f)
|
|
try:
|
|
total_size += os.path.getsize(fp)
|
|
file_count += 1
|
|
except OSError:
|
|
pass
|
|
except Exception as e:
|
|
logger.error(f"Error calculating storage stats: {e}")
|
|
return {
|
|
"total_size": 0,
|
|
"total_size_formatted": "Error",
|
|
"file_count": 0,
|
|
"error": str(e),
|
|
"calculated_at": datetime.utcnow().isoformat(),
|
|
}
|
|
|
|
return {
|
|
"total_size": total_size,
|
|
"total_size_formatted": format_size(total_size),
|
|
"file_count": file_count,
|
|
"calculated_at": datetime.utcnow().isoformat(),
|
|
}
|
|
|
|
|
|
async def _update_storage_stats_async() -> dict:
|
|
"""Async implementation to calculate and store storage stats."""
|
|
session_factory, engine = get_async_session()
|
|
|
|
try:
|
|
# Calculate stats (synchronous filesystem operation)
|
|
stats = calculate_storage_stats()
|
|
logger.info(f"Calculated storage stats: {stats['total_size_formatted']}, {stats['file_count']} files")
|
|
|
|
# Store in database
|
|
async with session_factory() as session:
|
|
# Check if setting exists
|
|
result = await session.execute(
|
|
select(Setting).where(Setting.key == STORAGE_STATS_SETTING)
|
|
)
|
|
setting = result.scalar_one_or_none()
|
|
|
|
if setting:
|
|
setting.value = stats
|
|
else:
|
|
setting = Setting(key=STORAGE_STATS_SETTING, value=stats)
|
|
session.add(setting)
|
|
|
|
await session.commit()
|
|
|
|
return stats
|
|
|
|
finally:
|
|
await _cleanup_engine(engine)
|
|
|
|
|
|
@celery_app.task(name="tasks.update_storage_stats")
|
|
def update_storage_stats() -> dict:
|
|
"""Celery task to calculate and cache storage stats.
|
|
|
|
This task should be scheduled to run periodically (e.g., every hour).
|
|
"""
|
|
logger.info("Starting storage stats calculation...")
|
|
result = asyncio.run(_update_storage_stats_async())
|
|
logger.info(f"Storage stats updated: {result}")
|
|
return result
|