You've already forked favicon-api-async
144 lines
3.8 KiB
Python
144 lines
3.8 KiB
Python
# -*- 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}")
|