Add source

bot/utils/file_cleanup.py

@@ -0,0 +1,147 @@
+"""
+Utilities for cleaning up old files
+"""
+from pathlib import Path
+from datetime import datetime, timedelta
+import logging
+import asyncio
+
+logger = logging.getLogger(__name__)
+
+# Constants for time intervals
+SECONDS_PER_HOUR = 3600
+
+# Maximum file age before deletion (24 hours)
+MAX_FILE_AGE_HOURS = 24
+MAX_FILE_AGE = timedelta(hours=MAX_FILE_AGE_HOURS)
+
+# Default cleanup interval (6 hours)
+DEFAULT_CLEANUP_INTERVAL_HOURS = 6
+
+# Queue of files to delete (files that couldn't be deleted on first attempt)
+_files_to_cleanup: set[str] = set()
+_files_to_cleanup_lock = asyncio.Lock()
+
+
+def add_file_to_cleanup_queue(file_path: str):
+    """
+    Add file to cleanup queue
+    
+    Args:
+        file_path: Path to file
+    """
+    global _files_to_cleanup
+    _files_to_cleanup.add(file_path)
+    logger.debug(f"File added to cleanup queue: {file_path}")
+
+
+async def cleanup_queued_files():
+    """
+    Delete files from queue
+    """
+    global _files_to_cleanup
+    async with _files_to_cleanup_lock:
+        if not _files_to_cleanup:
+            return
+        
+        files_to_remove = list(_files_to_cleanup)
+        _files_to_cleanup.clear()
+    
+    from bot.modules.media_loader.sender import delete_file
+    for file_path in files_to_remove:
+        try:
+            await delete_file(file_path, max_retries=1)  # One attempt, as there were already attempts
+        except Exception as e:
+            logger.warning(f"Failed to delete file from queue: {file_path}: {e}")
+
+
+async def cleanup_old_files(downloads_dir: str = "downloads"):
+    """
+    Clean up old files from downloads/ directory
+    
+    Args:
+        downloads_dir: Path to downloads directory
+    """
+    try:
+        downloads_path = Path(downloads_dir)
+        if not downloads_path.exists():
+            return
+        
+        now = datetime.now()
+        deleted_count = 0
+        total_size = 0
+        
+        for file_path in downloads_path.iterdir():
+            if file_path.is_file():
+                try:
+                    # Get last modification time
+                    mtime = datetime.fromtimestamp(file_path.stat().st_mtime)
+                    age = now - mtime
+                    
+                    # Delete files older than MAX_FILE_AGE
+                    if age > MAX_FILE_AGE:
+                        file_size = file_path.stat().st_size
+                        file_path.unlink()
+                        deleted_count += 1
+                        total_size += file_size
+                        logger.debug(f"Deleted old file: {file_path.name} (age: {age})")
+                except Exception as e:
+                    logger.warning(f"Failed to delete file {file_path}: {e}")
+        
+        if deleted_count > 0:
+            logger.info(f"Cleaned up {deleted_count} old files, freed {total_size / (1024*1024):.2f} MB")
+    
+    except Exception as e:
+        logger.error(f"Error cleaning up old files: {e}", exc_info=True)
+
+
+async def cleanup_files_periodically(
+    downloads_dir: str = "downloads", 
+    interval_hours: int = DEFAULT_CLEANUP_INTERVAL_HOURS
+) -> None:
+    """
+    Periodically clean up old files
+    
+    Args:
+        downloads_dir: Path to downloads directory
+        interval_hours: Interval between cleanups in hours
+    """
+    while True:
+        try:
+            await asyncio.sleep(interval_hours * SECONDS_PER_HOUR)
+            await cleanup_old_files(downloads_dir)
+        except asyncio.CancelledError:
+            logger.info("File cleanup task stopped")
+            break
+        except Exception as e:
+            logger.error(f"Error in file cleanup task: {e}", exc_info=True)
+
+
+def get_downloads_dir_size(downloads_dir: str = "downloads") -> int:
+    """
+    Get total size of downloads/ directory
+    
+    Args:
+        downloads_dir: Path to downloads directory
+    
+    Returns:
+        Size in bytes
+    """
+    try:
+        downloads_path = Path(downloads_dir)
+        if not downloads_path.exists():
+            return 0
+        
+        total_size = 0
+        for file_path in downloads_path.rglob('*'):
+            if file_path.is_file():
+                try:
+                    total_size += file_path.stat().st_size
+                except Exception:
+                    pass
+        
+        return total_size
+    except Exception as e:
+        logger.error(f"Error calculating directory size: {e}")
+        return 0
+