CoolBotDataSys/cache.py

"""
Redis 缓存模块 - 修复版
实现爬虫结果缓存、API 响应缓存、去重等功能
"""
import redis
import json
import logging
from typing import Optional, Any
import os
import sys

sys.path.insert(0, os.path.dirname(__file__))

logger = logging.getLogger(__name__)

# 使用 Config 类
from config import config

redis_config = config.redis

class RedisCache:
    _instance = None
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._init_client()
        return cls._instance
    
    def _init_client(self):
        try:
            self.client = redis.Redis(
                host=redis_config.get('host', '127.0.0.1'),
                port=redis_config.get('port', 6379),
                db=redis_config.get('db', 0),
                decode_responses=True,
                socket_timeout=5,
                socket_connect_timeout=5,
                retry_on_timeout=True
            )
            self.client.ping()
            logger.info("Redis connected successfully")
        except redis.ConnectionError as e:
            logger.warning(f"Redis connection failed: {e}, caching disabled")
            self.client = None
    
    @property
    def is_available(self) -> bool:
        if not self.client:
            return False
        try:
            self.client.ping()
            return True
        except:
            return False
    
    def get(self, key: str) -> Optional[Any]:
        if not self.is_available:
            return None
        try:
            val = self.client.get(key)
            if val:
                return json.loads(val)
            return None
        except Exception as e:
            logger.warning(f"Cache get error: {e}")
            return None
    
    def set(self, key: str, value: Any, ttl: int = 300):
        if not self.is_available:
            return False
        try:
            self.client.setex(key, ttl, json.dumps(value, default=str))
            return True
        except Exception as e:
            logger.warning(f"Cache set error: {e}")
            return False
    
    def delete(self, key: str):
        if not self.is_available:
            return False
        try:
            self.client.delete(key)
            return True
        except Exception as e:
            logger.warning(f"Cache delete error: {e}")
            return False
    
    def invalidate_pattern(self, pattern: str):
        if not self.is_available:
            return 0
        count = 0
        try:
            for key in self.client.scan_iter(match=pattern):
                self.client.delete(key)
                count += 1
            return count
        except Exception as e:
            logger.warning(f"Cache invalidate error: {e}")
            return count
    
    def get_post_cached(self, post_id: str) -> Optional[dict]:
        return self.get(f"post:{post_id}")
    
    def set_post_cached(self, post_id: str, data: dict, ttl: int = 300):
        return self.set(f"post:{post_id}", data, ttl)
    
    def is_post_crawled(self, post_id: str) -> bool:
        if not self.is_available:
            return False
        try:
            return self.client.exists(f"crawled:{post_id}") > 0
        except:
            return False
    
    def mark_post_crawled(self, post_id: str, ttl: int = 86400):
        if not self.is_available:
            return False
        try:
            self.client.setex(f"crawled:{post_id}", ttl, "1")
            return True
        except:
            return False
    
    def get_statistics_cached(self) -> Optional[dict]:
        return self.get("yichens:statistics")
    
    def set_statistics_cached(self, data: dict, ttl: int = 60):
        return self.set("yichens:statistics", data, ttl)
    
    def invalidate_statistics(self):
        return self.delete("yichens:statistics")

cache = RedisCache()