87 lines
2.5 KiB
Python
87 lines
2.5 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Helper wrapper to integrate monitoring with downloaders
|
|
"""
|
|
|
|
from functools import wraps
|
|
from modules.downloader_monitor import get_monitor
|
|
|
|
|
|
def monitor_download(downloader_name):
|
|
"""
|
|
Decorator to monitor download attempts
|
|
|
|
Usage:
|
|
@monitor_download('fastdl')
|
|
def download_function(username, ...):
|
|
...
|
|
return count
|
|
"""
|
|
def decorator(func):
|
|
@wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
# Extract username from args or kwargs
|
|
username = kwargs.get('username') or (args[0] if args else 'unknown')
|
|
|
|
try:
|
|
# Call the actual download function
|
|
result = func(*args, **kwargs)
|
|
|
|
# Determine success based on result
|
|
if isinstance(result, int):
|
|
count = result
|
|
success = count > 0
|
|
elif isinstance(result, dict):
|
|
count = result.get('count', 0)
|
|
success = result.get('success', count > 0)
|
|
else:
|
|
count = 0
|
|
success = False
|
|
|
|
# Log to monitor
|
|
monitor = get_monitor()
|
|
monitor.log_download_attempt(
|
|
downloader=downloader_name,
|
|
username=username,
|
|
success=success,
|
|
file_count=count,
|
|
error_message=None
|
|
)
|
|
|
|
return result
|
|
|
|
except Exception as e:
|
|
# Log failure
|
|
monitor = get_monitor()
|
|
monitor.log_download_attempt(
|
|
downloader=downloader_name,
|
|
username=username,
|
|
success=False,
|
|
file_count=0,
|
|
error_message=str(e)
|
|
)
|
|
raise
|
|
|
|
return wrapper
|
|
return decorator
|
|
|
|
|
|
def log_download_result(downloader: str, username: str, count: int, error: str = None):
|
|
"""
|
|
Simple function to log download result to monitor
|
|
|
|
Args:
|
|
downloader: Downloader name (fastdl, imginn, etc.)
|
|
username: Username
|
|
count: Number of files downloaded
|
|
error: Error message if failed
|
|
"""
|
|
monitor = get_monitor()
|
|
monitor.log_download_attempt(
|
|
downloader=downloader,
|
|
username=username,
|
|
success=(error is None),
|
|
file_count=count,
|
|
error_message=error
|
|
)
|