init

favicon_app/utils/redis_pool.py

@@ -0,0 +1,143 @@
+# -*- coding: utf-8 -*-
+
+import logging
+from typing import AsyncGenerator, Optional
+
+from redis.asyncio import ConnectionPool, Redis
+
+import setting
+
+logger = logging.getLogger(__name__)
+
+REDIS_URL = setting.REDIS_URL
+
+# Redis键前缀
+FAILED_DOMAINS_PREFIX = "favicon:failed_domain:"  # 失败域名前缀
+TASK_QUEUE_PREFIX = "favicon:task_queue:"  # 任务队列前缀
+PROCESSING_SET_PREFIX = "favicon:processing:"  # 处理中任务集合前缀
+ICON_QUEUE_PREFIX = "favicon:icon_queue:"
+
+pool = ConnectionPool.from_url(
+    REDIS_URL,
+    max_connections=200,
+    decode_responses=True,
+)
+
+
+async def get_redis() -> AsyncGenerator[Redis, None]:
+    async with Redis(connection_pool=pool) as conn:
+        yield conn
+
+
+async def set_cache(key: str, value: [str | int], ttl: int = None) -> None:
+    if not key:
+        return
+
+    try:
+        async for redis in get_redis():
+            await redis.set(key, value, ex=ttl)
+    except Exception as e:
+        logger.error(f"存入redis时出错：{e}")
+
+
+async def get_cache(key: str) -> Optional[str | int]:
+    if not key:
+        return None
+
+    try:
+        async for redis in get_redis():
+            return await redis.get(key)
+    except Exception as e:
+        logger.error(f"读取redis时出错：{e}")
+
+
+async def exist_cache(key: str) -> bool:
+    if not key:
+        return False
+
+    try:
+        async for redis in get_redis():
+            result = await redis.exists(key)
+            return result > 0
+    except Exception as e:
+        logger.error(f"读取redis时出错：{e}")
+        return False
+
+
+async def remove_cache(key: str) -> None:
+    if not key:
+        return
+
+    try:
+        async for redis in get_redis():
+            await redis.delete(key)
+    except Exception as e:
+        logger.error(f"删除redis时出错：{e}")
+
+
+async def get_cache_size(cache_name: str = "default") -> int:
+    try:
+        async for redis in get_redis():
+            return await redis.llen(cache_name)
+    except Exception as e:
+        logger.error(f"获取队列大小时出错：{e}")
+        return 0
+
+
+async def set_failed_domain(domain: str, expire_seconds: int = setting.time_of_7_days) -> None:
+    """将失败的域名存入Redis，并设置过期时间
+
+    Args:
+        domain: 失败的域名
+        expire_seconds: 过期时间（秒），默认为7天
+    """
+    if not domain:
+        return
+
+    try:
+        async for redis in get_redis():
+            redis_key = f"{FAILED_DOMAINS_PREFIX}{domain}"
+            await redis.set(redis_key, domain, ex=expire_seconds)
+            logger.debug(f"已将失败域名 {domain} 存入Redis，过期时间：{expire_seconds}秒")
+    except Exception as e:
+        logger.error(f"将失败域名存入Redis时出错：{e}")
+
+
+async def is_domain_failed(domain: str) -> bool:
+    """检查域名是否在Redis的失败列表中
+
+    Args:
+        domain: 要检查的域名
+
+    Returns:
+        True: 域名在失败列表中；False: 不在或Redis查询失败
+    """
+    if not domain:
+        return False
+
+    try:
+        async for redis in get_redis():
+            redis_key = f"{FAILED_DOMAINS_PREFIX}{domain}"
+            result = await redis.exists(redis_key)
+            return result > 0
+    except Exception as e:
+        logger.error(f"检查域名是否失败时出错：{e}")
+        return False
+
+
+async def delete_failed_domain(domain: str) -> None:
+    """从Redis中删除失败域名记录
+
+    Args:
+        domain: 要删除的域名
+    """
+    if not domain:
+        return
+
+    try:
+        async for redis in get_redis():
+            redis_key = f"{FAILED_DOMAINS_PREFIX}{domain}"
+            await redis.delete(redis_key)
+            logger.debug(f"已从Redis删除失败域名 {domain}")
+    except Exception as e:
+        logger.error(f"从Redis删除失败域名时出错：{e}")