Compare commits

...

2 Commits

+49 -11
View File
@@ -1,5 +1,5 @@
import httpx import httpx
from fastapi import APIRouter, Depends, Request, HTTPException from fastapi import APIRouter, Depends, Request, HTTPException, BackgroundTasks
from fastapi.responses import RedirectResponse from fastapi.responses import RedirectResponse
import os import os
import secrets import secrets
@@ -19,6 +19,38 @@ GOOGLE_AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"
GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token" GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
GOOGLE_USERINFO_URL = "https://www.googleapis.com/oauth2/v2/userinfo" GOOGLE_USERINFO_URL = "https://www.googleapis.com/oauth2/v2/userinfo"
# 定义独立的后台任务函数,用于处理头像转存到 S3 的耗时操作
async def async_upload_avatar_task(google_avatar_url: str, user_id: str):
"""
后台任务函数:下载 Google 头像并上传到 S3,然后更新用户记录。
因为主接口的 db 已经关闭,这里需要独立管理 db 的生命周期。
"""
from app.utils.s3 import upload_google_avatar_to_s3
# 如果 Google 本身就没给头像,无需转存
if not google_avatar_url:
return
try:
# 执行耗时的网络 I/O:下载并上传到 S3
s3_avatar_url = await upload_google_avatar_to_s3(google_avatar_url, user_id)
# 成功后,开启一个独立的数据库连接写回数据
db_task = SessionLocal()
try:
user = db_task.query(User).filter(User.id == user_id).first()
if user:
user.avatar = s3_avatar_url
db_task.commit()
print(f"成功在后台为用户 {user_id} 转存头像到 S3")
finally:
db_task.close() # 必须手动关闭连接
except Exception as s3_err:
# 如果 S3 依然失败,打印错误日志,但此时绝不会打扰到已经登录的用户
print(f"Background S3 Upload failed for user {user_id}: {s3_err}")
@router.get("/login/google") @router.get("/login/google")
def login_google(): def login_google():
state = secrets.token_urlsafe(16) state = secrets.token_urlsafe(16)
@@ -43,7 +75,7 @@ def login_google():
return response return response
@router.get("/auth/callback") @router.get("/auth/callback")
async def auth_callback(request: Request, code: str, state: str): async def auth_callback(request: Request, code: str, state: str, background_tasks: BackgroundTasks):
cookie_state = request.cookies.get("oauth_state") cookie_state = request.cookies.get("oauth_state")
if not cookie_state or cookie_state != state: if not cookie_state or cookie_state != state:
@@ -86,11 +118,18 @@ async def auth_callback(request: Request, code: str, state: str):
random_suffix = secrets.token_hex(3) # 生成 6 位随机字符 random_suffix = secrets.token_hex(3) # 生成 6 位随机字符
suggested_handle = f"{base_handle}_{random_suffix}" suggested_handle = f"{base_handle}_{random_suffix}"
# 如果 Google 没有头像,使用 DiceBear 头像生成器作为兜底
google_avatar = user_info.get("picture")
if not google_avatar:
default_avatar_url = f"https://api.dicebear.com/7.x/lorelei/svg?seed={base_handle}"
else:
default_avatar_url = google_avatar
user = User( user = User(
google_id=user_info["id"], google_id=user_info["id"],
email=user_info["email"], email=user_info["email"],
name=user_info["name"], name=user_info["name"],
avatar=user_info["picture"], avatar=default_avatar_url,
handle=suggested_handle, handle=suggested_handle,
) )
db.add(user) db.add(user)
@@ -98,14 +137,13 @@ async def auth_callback(request: Request, code: str, state: str):
db.refresh(user) db.refresh(user)
# 异步/同步将头像转存到 S3 # 异步/同步将头像转存到 S3
from app.utils.s3 import upload_google_avatar_to_s3 # FastAPI 会在把下方的 return 响应发送给浏览器之后,立刻在后台启动这个任务
try: if google_avatar:
s3_avatar_url = await upload_google_avatar_to_s3(user_info["picture"], str(user.id)) background_tasks.add_task(
# 将更新后的 S3 URL 写回数据库 async_upload_avatar_task,
user.avatar = s3_avatar_url google_avatar,
db.commit() str(user.id)
except Exception as s3_err: )
print(f"S3 Upload failed: {s3_err}")
# 这里不需要再次 refresh,除非后面还要用到更新后的 user 对象 # 这里不需要再次 refresh,除非后面还要用到更新后的 user 对象
# db.refresh(user) # db.refresh(user)