From 3538a048708208cf86eafc8adb8e23de112f8aaa Mon Sep 17 00:00:00 2001
From: jxxghp <jxxghp@gmail.com>
Date: Wed, 5 Jul 2023 11:25:05 +0800
Subject: [PATCH] fix sse

---
 app/api/endpoints/site.py   | 22 ++++++----------------
 app/api/endpoints/system.py |  4 ++--
 app/helper/message.py       |  4 +++-
 3 files changed, 11 insertions(+), 19 deletions(-)

diff --git a/app/api/endpoints/site.py b/app/api/endpoints/site.py
index 5630b9fc..fae72e27 100644
--- a/app/api/endpoints/site.py
+++ b/app/api/endpoints/site.py
@@ -23,16 +23,6 @@ def start_cookiecloud_sync():
     CookieCloudChain().process(manual=True)
 
 
-def start_site_cookie(site_info: Site, username: str, password: str):
-    """
-    后台启动站点Cookie更新
-    """
-    SiteChain().update_cookie(site_info=site_info,
-                              username=username,
-                              password=password,
-                              manual=True)
-
-
 @router.get("/", summary="所有站点", response_model=List[schemas.Site])
 async def read_sites(db: Session = Depends(get_db),
                      _: schemas.TokenPayload = Depends(verify_token)) -> List[dict]:
@@ -103,7 +93,7 @@ async def cookie_cloud_sync(background_tasks: BackgroundTasks,
     return schemas.Response(success=True, message="CookieCloud同步任务已启动！")
 
 
-@router.get("/cookie/{site_id}", summary="更新站点Cookie&UA", response_model=schemas.Response)
+@router.put("/cookie/{site_id}", summary="更新站点Cookie&UA", response_model=schemas.Response)
 async def update_cookie(
         site_id: int,
         username: str,
@@ -122,11 +112,11 @@ async def update_cookie(
             detail=f"站点 {site_id} 不存在！",
         )
     # 更新Cookie
-    background_tasks.add_task(start_site_cookie,
-                              site_info=site_info,
-                              username=username,
-                              password=password)
-    return schemas.Response(success=True, message="站点Cookie更新任务已启动！")
+    state, message = SiteChain().update_cookie(site_info=site_info,
+                                               username=username,
+                                               password=password,
+                                               manual=True)
+    return schemas.Response(success=state, message=message)
 
 
 @router.get("/test/{site_id}", summary="连接测试", response_model=schemas.Response)
diff --git a/app/api/endpoints/system.py b/app/api/endpoints/system.py
index 9fa47a95..9b560295 100644
--- a/app/api/endpoints/system.py
+++ b/app/api/endpoints/system.py
@@ -51,7 +51,7 @@ async def get_progress(token: str):
     async def event_generator():
         while True:
             detail = message.get()
-            yield 'data: %s\n\n' % detail
-            await asyncio.sleep(1)
+            yield 'data: %s\n\n' % (detail or '')
+            await asyncio.sleep(3)
 
     return StreamingResponse(event_generator(), media_type="text/event-stream")
diff --git a/app/helper/message.py b/app/helper/message.py
index 27157cf8..d344faad 100644
--- a/app/helper/message.py
+++ b/app/helper/message.py
@@ -14,4 +14,6 @@ class MessageHelper(metaclass=Singleton):
         self.queue.put(message)
 
     def get(self):
-        return self.queue.get(block=False)
+        if not self.queue.empty():
+            return self.queue.get(block=False)
+        return None