This commit is contained in:
dgtlmoon
2026-01-05 16:18:15 +01:00
parent b62efa1470
commit ca5bb8fe93
8 changed files with 1127 additions and 330 deletions
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,127 @@
"""
Abstract base class for Huey storage backend task managers.
"""
from abc import ABC, abstractmethod
class HueyTaskManager(ABC):
"""
Abstract base class for Huey storage backend operations.
Provides a polymorphic interface for storage-specific operations like:
- Enumerating results (failed notifications)
- Deleting results
- Counting pending notifications
- Clearing all notifications
Each storage backend (FileStorage, SqliteStorage, RedisStorage) has its own
concrete implementation that knows how to interact with that specific storage.
"""
def __init__(self, storage, storage_path=None):
"""
Initialize task manager with storage instance.
Args:
storage: Huey storage instance
storage_path: Optional path for file-based storage
"""
self.storage = storage
self.storage_path = storage_path
@abstractmethod
def enumerate_results(self):
"""
Enumerate all results from storage.
Returns:
dict: {task_id: result_data} for all stored results
"""
pass
@abstractmethod
def delete_result(self, task_id):
"""
Delete a result from storage.
Args:
task_id: Task ID to delete
Returns:
bool: True if deleted successfully, False otherwise
"""
pass
@abstractmethod
def count_storage_items(self):
"""
Count items in storage (queue + schedule).
Returns:
tuple: (queue_count, schedule_count)
"""
pass
@abstractmethod
def clear_all_notifications(self):
"""
Clear all notifications (queue, schedule, results, metadata).
Returns:
dict: Counts of cleared items by type
"""
pass
@abstractmethod
def store_task_metadata(self, task_id, metadata):
"""
Store task metadata for later retrieval.
Args:
task_id: Task ID
metadata: Metadata dictionary to store
Returns:
bool: True if stored successfully, False otherwise
"""
pass
@abstractmethod
def get_task_metadata(self, task_id):
"""
Retrieve task metadata.
Args:
task_id: Task ID
Returns:
dict: Metadata dictionary or None if not found
"""
pass
@abstractmethod
def delete_task_metadata(self, task_id):
"""
Delete task metadata.
Args:
task_id: Task ID
Returns:
bool: True if deleted successfully, False otherwise
"""
pass
@abstractmethod
def cleanup_old_retry_attempts(self, cutoff_time):
"""
Clean up retry attempt records older than cutoff_time.
Args:
cutoff_time: Unix timestamp - delete records older than this
Returns:
int: Number of retry attempts deleted
"""
pass
@@ -0,0 +1,270 @@
"""
FileStorage backend task manager for Huey notifications.
This is the default backend, optimized for NAS/CIFS compatibility.
"""
from loguru import logger
from .base import HueyTaskManager
import os
class FileStorageTaskManager(HueyTaskManager):
"""Task manager for FileStorage backend (default, NAS-safe)."""
def enumerate_results(self):
"""Enumerate results by walking filesystem directories."""
import os
import pickle
import struct
import time
results = {}
if not self.storage_path:
return results
results_dir = os.path.join(self.storage_path, 'results')
if not os.path.exists(results_dir):
return results
# Walk through all subdirectories to find result files
for root, dirs, files in os.walk(results_dir):
for filename in files:
if filename.startswith('.'):
continue
filepath = os.path.join(root, filename)
try:
# Read and unpickle the result
# Huey FileStorage format: 4-byte length + task_id + pickled data
with open(filepath, 'rb') as f:
# Read the task ID header (length-prefixed)
task_id_len_bytes = f.read(4)
if len(task_id_len_bytes) < 4:
raise EOFError("Incomplete header")
task_id_len = struct.unpack('>I', task_id_len_bytes)[0]
task_id_bytes = f.read(task_id_len)
if len(task_id_bytes) < task_id_len:
raise EOFError("Incomplete task ID")
task_id = task_id_bytes.decode('utf-8')
# Now unpickle the result data
result_data = pickle.load(f)
results[task_id] = result_data
except (pickle.UnpicklingError, EOFError) as e:
# Corrupted or incomplete result file
file_size = os.path.getsize(filepath)
logger.warning(f"Corrupted result file {filename} ({file_size} bytes) - moving to lost-found.")
try:
import shutil
lost_found_dir = os.path.join(self.storage_path, 'lost-found', 'results')
os.makedirs(lost_found_dir, exist_ok=True)
timestamp = int(time.time())
lost_found_path = os.path.join(lost_found_dir, f"{filename}.{timestamp}.corrupted")
shutil.move(filepath, lost_found_path)
logger.info(f"Moved corrupted file to {lost_found_path}")
except Exception as move_err:
logger.error(f"Unable to move corrupted file: {move_err}")
except Exception as e:
logger.debug(f"Unable to read result file {filename}: {e}")
return results
def delete_result(self, task_id):
"""Delete result file from filesystem."""
import hashlib
if not self.storage_path:
return False
results_dir = os.path.join(self.storage_path, 'results')
# Huey uses MD5 hash to create subdirectories
task_id_bytes = task_id.encode('utf-8')
hex_hash = hashlib.md5(task_id_bytes).hexdigest()
# FileStorage creates subdirectories based on first 2 chars of hash
subdir = hex_hash[:2]
result_file = os.path.join(results_dir, subdir, task_id)
if os.path.exists(result_file):
os.remove(result_file)
logger.debug(f"Deleted result file for task {task_id}")
return True
else:
logger.debug(f"Result file not found for task {task_id}")
return False
def count_storage_items(self):
"""Count items by walking filesystem directories."""
queue_count = 0
schedule_count = 0
if not self.storage_path:
return queue_count, schedule_count
try:
# Count queue files
queue_dir = os.path.join(self.storage_path, 'queue')
if os.path.exists(queue_dir):
for root, dirs, files in os.walk(queue_dir):
queue_count += len([f for f in files if not f.startswith('.')])
# Count schedule files
schedule_dir = os.path.join(self.storage_path, 'schedule')
if os.path.exists(schedule_dir):
for root, dirs, files in os.walk(schedule_dir):
schedule_count += len([f for f in files if not f.startswith('.')])
except Exception as e:
logger.debug(f"FileStorage count error: {e}")
return queue_count, schedule_count
def clear_all_notifications(self):
"""Clear all notification files from filesystem."""
cleared = {
'queue': 0,
'schedule': 0,
'results': 0,
'retry_attempts': 0,
'task_metadata': 0
}
if not self.storage_path:
return cleared
# Clear queue
queue_dir = os.path.join(self.storage_path, 'queue')
if os.path.exists(queue_dir):
for root, dirs, files in os.walk(queue_dir):
for f in files:
if not f.startswith('.'):
os.remove(os.path.join(root, f))
cleared['queue'] += 1
# Clear schedule
schedule_dir = os.path.join(self.storage_path, 'schedule')
if os.path.exists(schedule_dir):
for root, dirs, files in os.walk(schedule_dir):
for f in files:
if not f.startswith('.'):
os.remove(os.path.join(root, f))
cleared['schedule'] += 1
# Clear results
results_dir = os.path.join(self.storage_path, 'results')
if os.path.exists(results_dir):
for root, dirs, files in os.walk(results_dir):
for f in files:
if not f.startswith('.'):
os.remove(os.path.join(root, f))
cleared['results'] += 1
# Clear retry attempts
attempts_dir = os.path.join(self.storage_path, 'retry_attempts')
if os.path.exists(attempts_dir):
for f in os.listdir(attempts_dir):
if f.endswith('.json'):
os.remove(os.path.join(attempts_dir, f))
cleared['retry_attempts'] += 1
# Clear task metadata
metadata_dir = os.path.join(self.storage_path, 'task_metadata')
if os.path.exists(metadata_dir):
for f in os.listdir(metadata_dir):
if f.endswith('.json'):
os.remove(os.path.join(metadata_dir, f))
cleared['task_metadata'] += 1
return cleared
def store_task_metadata(self, task_id, metadata):
"""Store task metadata as JSON file."""
import json
import time
if not self.storage_path:
return False
try:
metadata_dir = os.path.join(self.storage_path, 'task_metadata')
os.makedirs(metadata_dir, exist_ok=True)
metadata_file = os.path.join(metadata_dir, f"{task_id}.json")
metadata_with_id = {
'task_id': task_id,
'timestamp': time.time(),
**metadata
}
with open(metadata_file, 'w') as f:
json.dump(metadata_with_id, f, indent=2)
return True
except Exception as e:
logger.debug(f"Unable to store task metadata: {e}")
return False
def get_task_metadata(self, task_id):
"""Retrieve task metadata from JSON file."""
import json
if not self.storage_path:
return None
try:
metadata_dir = os.path.join(self.storage_path, 'task_metadata')
metadata_file = os.path.join(metadata_dir, f"{task_id}.json")
if os.path.exists(metadata_file):
with open(metadata_file, 'r') as f:
return json.load(f)
except Exception as e:
logger.debug(f"Unable to load task metadata for {task_id}: {e}")
return None
def delete_task_metadata(self, task_id):
"""Delete task metadata JSON file."""
if not self.storage_path:
return False
try:
metadata_dir = os.path.join(self.storage_path, 'task_metadata')
metadata_file = os.path.join(metadata_dir, f"{task_id}.json")
if os.path.exists(metadata_file):
os.remove(metadata_file)
return True
return False
except Exception as e:
logger.debug(f"Unable to delete task metadata for {task_id}: {e}")
return False
def cleanup_old_retry_attempts(self, cutoff_time):
"""Clean up old retry attempt files from filesystem."""
if not self.storage_path:
return 0
deleted_count = 0
try:
attempts_dir = os.path.join(self.storage_path, 'retry_attempts')
if os.path.exists(attempts_dir):
for filename in os.listdir(attempts_dir):
if filename.endswith('.json'):
filepath = os.path.join(attempts_dir, filename)
try:
file_mtime = os.path.getmtime(filepath)
if file_mtime < cutoff_time:
os.remove(filepath)
deleted_count += 1
except Exception as fe:
logger.debug(f"Unable to delete old retry attempt file {filename}: {fe}")
except Exception as e:
logger.debug(f"Error cleaning up old retry attempts: {e}")
return deleted_count
@@ -0,0 +1,204 @@
"""
RedisStorage backend task manager for Huey notifications.
For distributed deployments with Redis.
"""
from loguru import logger
from .base import HueyTaskManager
class RedisStorageTaskManager(HueyTaskManager):
"""Task manager for RedisStorage backend (distributed deployments)."""
def enumerate_results(self):
import pickle
"""Enumerate results using Redis commands."""
results = {}
if not hasattr(self.storage, 'conn'):
return results
try:
# Redis stores results with keys like "{name}:result:{task_id}"
name = self.storage.name
pattern = f"{name}:result:*"
# Get all result keys
result_keys = self.storage.conn.keys(pattern)
for key in result_keys:
# Extract task_id from key
task_id = key.decode('utf-8').split(':')[-1]
# Get result data
result_data = self.storage.conn.get(key)
if result_data:
results[task_id] = pickle.loads(result_data)
except Exception as e:
logger.error(f"Error enumerating Redis results: {e}")
return results
def delete_result(self, task_id):
"""Delete result from Redis."""
if not hasattr(self.storage, 'conn'):
return False
try:
name = self.storage.name
result_key = f"{name}:result:{task_id}"
deleted = self.storage.conn.delete(result_key) > 0
logger.debug(f"Deleted result from Redis for task {task_id}: {deleted}")
return deleted
except Exception as e:
logger.error(f"Error deleting Redis result: {e}")
return False
def count_storage_items(self):
"""Count items using Redis commands."""
queue_count = 0
schedule_count = 0
if not hasattr(self.storage, 'conn'):
return queue_count, schedule_count
try:
name = self.storage.name
# Queue is a list
queue_count = self.storage.conn.llen(f"{name}:queue")
# Schedule is a sorted set
schedule_count = self.storage.conn.zcard(f"{name}:schedule")
except Exception as e:
logger.debug(f"Redis count error: {e}")
return queue_count, schedule_count
def clear_all_notifications(self):
"""Clear all notifications from Redis."""
cleared = {
'queue': 0,
'schedule': 0,
'results': 0,
'retry_attempts': 0,
'task_metadata': 0
}
if not hasattr(self.storage, 'conn'):
return cleared
try:
name = self.storage.name
# Clear queue (list)
cleared['queue'] = self.storage.conn.llen(f"{name}:queue")
self.storage.conn.delete(f"{name}:queue")
# Clear schedule (sorted set)
cleared['schedule'] = self.storage.conn.zcard(f"{name}:schedule")
self.storage.conn.delete(f"{name}:schedule")
# Clear results (keys)
result_keys = self.storage.conn.keys(f"{name}:result:*")
if result_keys:
cleared['results'] = len(result_keys)
self.storage.conn.delete(*result_keys)
except Exception as e:
logger.error(f"Error clearing Redis notifications: {e}")
return cleared
def store_task_metadata(self, task_id, metadata):
"""Store task metadata in Redis."""
import json
import time
if not hasattr(self.storage, 'conn'):
return False
try:
name = self.storage.name
metadata_key = f"{name}:metadata:{task_id}"
metadata_with_id = {
'task_id': task_id,
'timestamp': time.time(),
**metadata
}
self.storage.conn.set(metadata_key, json.dumps(metadata_with_id))
return True
except Exception as e:
logger.error(f"Error storing Redis task metadata: {e}")
return False
def get_task_metadata(self, task_id):
"""Retrieve task metadata from Redis."""
import json
if not hasattr(self.storage, 'conn'):
return None
try:
name = self.storage.name
metadata_key = f"{name}:metadata:{task_id}"
data = self.storage.conn.get(metadata_key)
if data:
return json.loads(data.decode('utf-8') if isinstance(data, bytes) else data)
return None
except Exception as e:
logger.debug(f"Error retrieving Redis task metadata: {e}")
return None
def delete_task_metadata(self, task_id):
"""Delete task metadata from Redis."""
if not hasattr(self.storage, 'conn'):
return False
try:
name = self.storage.name
metadata_key = f"{name}:metadata:{task_id}"
deleted = self.storage.conn.delete(metadata_key) > 0
return deleted
except Exception as e:
logger.debug(f"Error deleting Redis task metadata: {e}")
return False
def cleanup_old_retry_attempts(self, cutoff_time):
"""Clean up old retry attempts from Redis."""
if not hasattr(self.storage, 'conn'):
return 0
deleted_count = 0
try:
name = self.storage.name
pattern = f"{name}:retry_attempts:*"
# Get all retry attempt keys
retry_keys = self.storage.conn.keys(pattern)
for key in retry_keys:
try:
# Get the timestamp from the key's data
data = self.storage.conn.get(key)
if data:
import json
attempt_data = json.loads(data.decode('utf-8') if isinstance(data, bytes) else data)
timestamp = attempt_data.get('timestamp', 0)
if timestamp < cutoff_time:
self.storage.conn.delete(key)
deleted_count += 1
except Exception as ke:
logger.debug(f"Error checking retry attempt key: {ke}")
if deleted_count > 0:
logger.info(f"Cleaned up {deleted_count} old retry attempts from Redis")
except Exception as e:
logger.debug(f"Error cleaning up old Redis retry attempts: {e}")
return deleted_count
@@ -0,0 +1,239 @@
"""
SQLiteStorage backend task manager for Huey notifications.
WARNING: Only use on local disk storage, NOT on NFS/CIFS network storage!
"""
from loguru import logger
from .base import HueyTaskManager
class SqliteStorageTaskManager(HueyTaskManager):
"""Task manager for SqliteStorage backend (local disk only)."""
def enumerate_results(self):
import pickle
import sqlite3
"""Enumerate results by querying SQLite database."""
results = {}
if not hasattr(self.storage, 'filename'):
return results
try:
conn = sqlite3.connect(self.storage.filename)
cursor = conn.cursor()
# Query all results from database
cursor.execute("SELECT key, value FROM results")
for row in cursor.fetchall():
task_id = row[0]
result_data = pickle.loads(row[1])
results[task_id] = result_data
conn.close()
except Exception as e:
logger.error(f"Error enumerating SQLite results: {e}")
return results
def delete_result(self, task_id):
"""Delete result from SQLite database."""
if not hasattr(self.storage, 'filename'):
return False
import sqlite3
try:
conn = sqlite3.connect(self.storage.filename)
cursor = conn.cursor()
cursor.execute("DELETE FROM results WHERE key = ?", (task_id,))
conn.commit()
deleted = cursor.rowcount > 0
conn.close()
logger.debug(f"Deleted result from SQLite for task {task_id}: {deleted}")
return deleted
except Exception as e:
logger.error(f"Error deleting SQLite result: {e}")
return False
def count_storage_items(self):
"""Count items by querying SQLite database."""
queue_count = 0
schedule_count = 0
if not hasattr(self.storage, 'filename'):
return queue_count, schedule_count
import sqlite3
try:
conn = sqlite3.connect(self.storage.filename)
cursor = conn.cursor()
cursor.execute("SELECT COUNT(*) FROM queue")
queue_count = cursor.fetchone()[0]
cursor.execute("SELECT COUNT(*) FROM schedule")
schedule_count = cursor.fetchone()[0]
conn.close()
except Exception as e:
logger.debug(f"SQLite count error: {e}")
return queue_count, schedule_count
def clear_all_notifications(self):
"""Clear all notifications from SQLite database."""
cleared = {
'queue': 0,
'schedule': 0,
'results': 0,
'retry_attempts': 0,
'task_metadata': 0
}
if not hasattr(self.storage, 'filename'):
return cleared
import sqlite3
try:
conn = sqlite3.connect(self.storage.filename)
cursor = conn.cursor()
cursor.execute("DELETE FROM queue")
cleared['queue'] = cursor.rowcount
cursor.execute("DELETE FROM schedule")
cleared['schedule'] = cursor.rowcount
cursor.execute("DELETE FROM results")
cleared['results'] = cursor.rowcount
conn.commit()
conn.close()
except Exception as e:
logger.error(f"Error clearing SQLite notifications: {e}")
return cleared
def store_task_metadata(self, task_id, metadata):
"""Store task metadata in SQLite database."""
import sqlite3
import json
import time
if not hasattr(self.storage, 'filename'):
return False
try:
conn = sqlite3.connect(self.storage.filename)
cursor = conn.cursor()
# Create table if it doesn't exist
cursor.execute("""
CREATE TABLE IF NOT EXISTS task_metadata (
task_id TEXT PRIMARY KEY,
timestamp REAL,
metadata TEXT
)
""")
metadata_with_id = {
'task_id': task_id,
'timestamp': time.time(),
**metadata
}
cursor.execute(
"INSERT OR REPLACE INTO task_metadata (task_id, timestamp, metadata) VALUES (?, ?, ?)",
(task_id, time.time(), json.dumps(metadata_with_id))
)
conn.commit()
conn.close()
return True
except Exception as e:
logger.error(f"Error storing SQLite task metadata: {e}")
return False
def get_task_metadata(self, task_id):
"""Retrieve task metadata from SQLite database."""
import sqlite3
import json
if not hasattr(self.storage, 'filename'):
return None
try:
conn = sqlite3.connect(self.storage.filename)
cursor = conn.cursor()
# Check if table exists
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='task_metadata'")
if not cursor.fetchone():
conn.close()
return None
cursor.execute("SELECT metadata FROM task_metadata WHERE task_id = ?", (task_id,))
row = cursor.fetchone()
conn.close()
if row:
return json.loads(row[0])
return None
except Exception as e:
logger.debug(f"Error retrieving SQLite task metadata: {e}")
return None
def delete_task_metadata(self, task_id):
"""Delete task metadata from SQLite database."""
import sqlite3
if not hasattr(self.storage, 'filename'):
return False
try:
conn = sqlite3.connect(self.storage.filename)
cursor = conn.cursor()
# Check if table exists
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='task_metadata'")
if not cursor.fetchone():
conn.close()
return False
cursor.execute("DELETE FROM task_metadata WHERE task_id = ?", (task_id,))
conn.commit()
deleted = cursor.rowcount > 0
conn.close()
return deleted
except Exception as e:
logger.debug(f"Error deleting SQLite task metadata: {e}")
return False
def cleanup_old_retry_attempts(self, cutoff_time):
"""Clean up old retry attempts from SQLite database."""
import sqlite3
if not hasattr(self.storage, 'filename'):
return 0
deleted_count = 0
try:
conn = sqlite3.connect(self.storage.filename)
cursor = conn.cursor()
# Check if retry_attempts table exists
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='retry_attempts'")
if not cursor.fetchone():
conn.close()
return 0
# Delete old retry attempts
cursor.execute("DELETE FROM retry_attempts WHERE timestamp < ?", (cutoff_time,))
deleted_count = cursor.rowcount
conn.commit()
conn.close()
if deleted_count > 0:
logger.info(f"Cleaned up {deleted_count} old retry attempts from SQLite")
except Exception as e:
logger.debug(f"Error cleaning up old SQLite retry attempts: {e}")
return deleted_count