background task for s3 avatar uploading

app/auth.py

@@ -1,5 +1,5 @@
 import httpx
-from fastapi import APIRouter, Depends, Request, HTTPException
+from fastapi import APIRouter, Depends, Request, HTTPException, BackgroundTasks
 from fastapi.responses import RedirectResponse
 import os
 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_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")
 def login_google():
     state = secrets.token_urlsafe(16)
@@ -43,7 +75,7 @@ def login_google():
     return response
 
 @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")
 
     if not cookie_state or cookie_state != state:
@@ -85,12 +117,19 @@ async def auth_callback(request: Request, code: str, state: str):
             import secrets
             random_suffix = secrets.token_hex(3) # 生成 6 位随机字符
             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(
                 google_id=user_info["id"],
                 email=user_info["email"],
                 name=user_info["name"],
-                avatar=user_info["picture"],
+                avatar=default_avatar_url,
                 handle=suggested_handle,
             )
             db.add(user)
@@ -98,14 +137,13 @@ async def auth_callback(request: Request, code: str, state: str):
             db.refresh(user)
             
             # 异步/同步将头像转存到 S3
-            from app.utils.s3 import upload_google_avatar_to_s3
-            try:
-                s3_avatar_url = await upload_google_avatar_to_s3(user_info["picture"], str(user.id))
-                # 将更新后的 S3 URL 写回数据库
-                user.avatar = s3_avatar_url
-                db.commit()
-            except Exception as s3_err:
-                print(f"S3 Upload failed: {s3_err}")
+            # FastAPI 会在把下方的 return 响应发送给浏览器之后，立刻在后台启动这个任务
+            if google_avatar:
+                background_tasks.add_task(
+                    async_upload_avatar_task, 
+                    google_avatar, 
+                    str(user.id)
+                )
             
             # 这里不需要再次 refresh，除非后面还要用到更新后的 user 对象
             # db.refresh(user)