changedetection.io/changedetectionio/notification/task_queue/redis_storage.py

"""
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