232 lines
9.3 KiB
Python
232 lines
9.3 KiB
Python
from datetime import datetime, timedelta
|
|
import asyncio
|
|
import httpx
|
|
import json
|
|
from typing import Optional
|
|
|
|
from sqlalchemy import select, update
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
# 引入自定义北京时间工具类
|
|
from app.utils.datetime_util import BEIJING_TZ
|
|
from app.models.user_oauth import UserOAuth
|
|
from app.models.user_oauth_app import UserOAuthApp
|
|
from app.models.base import async_session
|
|
from app.utils.redis import get_redis
|
|
from app.utils.logger import get_logger
|
|
from app.utils.datetime_util import datetime_to_db_tz_str, db_tz_str_to_datetime
|
|
|
|
REDIS_KEY = "douyin:tokens"
|
|
logger = get_logger("token_refresh", "token_refresh")
|
|
|
|
# 配置常量
|
|
REFRESH_THRESHOLD_SECONDS = 800
|
|
CHECK_INTERVAL_MINUTES = 5
|
|
HTTP_TIMEOUT = httpx.Timeout(30.0)
|
|
|
|
|
|
async def _delete_redis_token(oauth_id: str):
|
|
"""【改动1：废弃空值覆盖，直接删除Hash脏缓存】refresh失效则移除字段"""
|
|
redis = get_redis()
|
|
if not redis:
|
|
logger.warning("Redis连接未配置，跳过缓存清理")
|
|
return
|
|
try:
|
|
await redis.hdel(REDIS_KEY, oauth_id)
|
|
logger.info(f"清理失效授权Redis缓存: oauth_id={oauth_id}")
|
|
except Exception as e:
|
|
logger.error(f"删除Redis缓存失败 oauth_id:{oauth_id}, err:{str(e)}")
|
|
|
|
|
|
async def _update_redis_token(oauth_id: str, token: str, expired_at: datetime):
|
|
"""更新Redis缓存中的token（统一北京时间序列化）"""
|
|
redis = get_redis()
|
|
if not redis:
|
|
logger.warning("Redis连接未配置，跳过缓存更新")
|
|
return
|
|
try:
|
|
cache = {
|
|
"token": token,
|
|
"expired_at": datetime_to_db_tz_str(expired_at),
|
|
}
|
|
await redis.hset(REDIS_KEY, oauth_id, json.dumps(cache))
|
|
logger.info(f"Redis缓存已更新: oauth_id={oauth_id}")
|
|
except Exception as e:
|
|
logger.error(f"更新Redis缓存失败 oauth_id:{oauth_id}, err:{str(e)}")
|
|
|
|
|
|
async def refresh_juliang_token(oauth: UserOAuth, app: UserOAuthApp, db: AsyncSession):
|
|
"""刷新巨量引擎token"""
|
|
try:
|
|
async with httpx.AsyncClient(timeout=HTTP_TIMEOUT) as client:
|
|
url = "https://api.oceanengine.com/open_api/oauth2/refresh_token/"
|
|
response = await client.post(
|
|
url,
|
|
json={
|
|
"app_id": app.app_id,
|
|
"secret": app.secret,
|
|
"refresh_token": oauth.refresh_token,
|
|
},
|
|
)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
|
|
if data.get("code") != 0:
|
|
logger.error(f"刷新巨量引擎token失败: oauth_id={oauth.id}, 错误信息: {data}")
|
|
# refresh_token已失效或刷新失败，统一清理token
|
|
if data.get("code") in [40103, 40107, 40000]:
|
|
where_cond = UserOAuth.deleted_at.is_(None)
|
|
if oauth.appid:
|
|
where_cond = where_cond & (UserOAuth.appid == oauth.appid)
|
|
if oauth.account_username:
|
|
where_cond = where_cond & (UserOAuth.account_username == oauth.account_username)
|
|
if oauth.account_userid:
|
|
where_cond = where_cond & (UserOAuth.account_userid == oauth.account_userid)
|
|
|
|
await db.execute(
|
|
update(UserOAuth).where(where_cond).values(
|
|
access_token=None,
|
|
access_token_expired=None,
|
|
refresh_token=None,
|
|
refresh_token_expired=None,
|
|
)
|
|
)
|
|
await db.commit()
|
|
|
|
related_oauth_ids = await db.execute(select(UserOAuth.id).where(where_cond))
|
|
related_oauth_ids = [row[0] for row in related_oauth_ids.all()]
|
|
for related_id in related_oauth_ids:
|
|
await _delete_redis_token(related_id)
|
|
return
|
|
|
|
resp_data = data.get("data", {})
|
|
new_access_token = resp_data.get("access_token", "")
|
|
new_refresh_token = resp_data.get("refresh_token", "")
|
|
|
|
# 【改动3：全部使用北京时间计算过期时间，统一时区】
|
|
now_beijing = datetime.now(tz=BEIJING_TZ)
|
|
expires_in = now_beijing + timedelta(seconds=resp_data.get("expires_in", 0))
|
|
refresh_expires = now_beijing + timedelta(seconds=resp_data.get("refresh_token_expires_in", 0))
|
|
|
|
where_cond = UserOAuth.deleted_at.is_(None)
|
|
if oauth.appid:
|
|
where_cond = where_cond & (UserOAuth.appid == oauth.appid)
|
|
if oauth.account_username:
|
|
where_cond = where_cond & (UserOAuth.account_username == oauth.account_username)
|
|
if oauth.account_userid:
|
|
where_cond = where_cond & (UserOAuth.account_userid == oauth.account_userid)
|
|
|
|
await db.execute(
|
|
update(UserOAuth).where(where_cond).values(
|
|
access_token=new_access_token,
|
|
access_token_expired=expires_in,
|
|
refresh_token=new_refresh_token,
|
|
refresh_token_expired=refresh_expires,
|
|
)
|
|
)
|
|
await db.commit()
|
|
|
|
related_oauth_ids = await db.execute(select(UserOAuth.id).where(where_cond))
|
|
related_oauth_ids = [row[0] for row in related_oauth_ids.all()]
|
|
for related_id in related_oauth_ids:
|
|
await _update_redis_token(related_id, new_access_token, expires_in)
|
|
|
|
logger.info(
|
|
f"成功刷新巨量引擎token: oauth_id={oauth.id}, "
|
|
f"account_id={oauth.account_id}, 关联账户数={len(related_oauth_ids)}"
|
|
)
|
|
except httpx.HTTPError as e:
|
|
logger.error(f"HTTP请求失败: oauth_id={oauth.id}, 错误: {str(e)}")
|
|
except Exception as e:
|
|
logger.error(f"刷新巨量引擎token发生异常: oauth_id={oauth.id}, 错误: {str(e)}")
|
|
|
|
|
|
async def check_and_refresh_tokens():
|
|
"""检查并刷新即将过期的token"""
|
|
async with async_session() as db:
|
|
# 【改动4：统一使用北京时间当前时间，避免时区运算异常】
|
|
now_beijing = datetime.now(tz=BEIJING_TZ)
|
|
|
|
query = select(UserOAuth).where(
|
|
UserOAuth.deleted_at.is_(None),
|
|
UserOAuth.refresh_token.is_not(None),
|
|
UserOAuth.refresh_token_expired.is_not(None),
|
|
UserOAuth.refresh_token_expired > now_beijing,
|
|
)
|
|
|
|
result = await db.execute(query)
|
|
oauth_list = result.scalars().all()
|
|
|
|
refreshed_keys = set()
|
|
|
|
for oauth in oauth_list:
|
|
try:
|
|
if oauth.port_type not in [1]:
|
|
continue
|
|
|
|
# 唯一键优化：过滤空值避免拼接异常
|
|
key_parts = []
|
|
if oauth.appid:
|
|
key_parts.append(oauth.appid)
|
|
if oauth.account_username:
|
|
key_parts.append(oauth.account_username)
|
|
if oauth.account_userid:
|
|
key_parts.append(str(oauth.account_userid))
|
|
login_key = "|".join(key_parts)
|
|
|
|
if login_key in refreshed_keys:
|
|
logger.debug(f"跳过重复刷新: oauth_id={oauth.id}, 同一登录账号已刷新")
|
|
continue
|
|
|
|
# 【改动5：过滤已禁用、已删除应用】
|
|
app_result = await db.execute(
|
|
select(UserOAuthApp).where(
|
|
UserOAuthApp.app_id == oauth.appid,
|
|
UserOAuthApp.status == 1,
|
|
UserOAuthApp.deleted_at.is_(None)
|
|
)
|
|
)
|
|
app = app_result.scalar_one_or_none()
|
|
if not app:
|
|
logger.warning(f"oauth_id:{oauth.id} 对应应用不存在/已禁用，跳过刷新")
|
|
continue
|
|
|
|
need_refresh = False
|
|
if not oauth.access_token or not oauth.access_token_expired:
|
|
need_refresh = True
|
|
else:
|
|
# 同时区时间运算，不会抛异常
|
|
remain_sec = (oauth.access_token_expired - now_beijing).total_seconds()
|
|
if remain_sec < REFRESH_THRESHOLD_SECONDS:
|
|
need_refresh = True
|
|
|
|
if not need_refresh:
|
|
continue
|
|
|
|
await refresh_juliang_token(oauth, app, db)
|
|
refreshed_keys.add(login_key)
|
|
|
|
# 简单限流，防止瞬间大量请求
|
|
await asyncio.sleep(0.2)
|
|
|
|
except Exception as e:
|
|
logger.error(f"oauth_id:{oauth.id} 刷新token异常: {str(e)}", exc_info=True)
|
|
|
|
|
|
async def token_refresh_scheduler():
|
|
"""定时任务调度器"""
|
|
logger.info("开始执行定时Token刷新任务")
|
|
while True:
|
|
try:
|
|
await check_and_refresh_tokens()
|
|
except Exception as e:
|
|
logger.error(f"定时任务token_refresh_scheduler执行失败: {str(e)}", exc_info=True)
|
|
|
|
await asyncio.sleep(CHECK_INTERVAL_MINUTES * 60)
|
|
|
|
|
|
def start_token_refresh_task():
|
|
"""启动token刷新定时任务"""
|
|
logger.info("启动token刷新定时后台任务")
|
|
asyncio.create_task(token_refresh_scheduler())
|