From 74a9bcaa64070357e41f5f161e9257da4eb32bd4 Mon Sep 17 00:00:00 2001
From: "Autumn.home" <autumn@home.com>
Date: Mon, 15 Jul 2024 23:10:26 +0800
Subject: [PATCH] fix bug

---
 api/project.py   | 16 ++++++++++++----
 api/sensitive.py |  2 +-
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/api/project.py b/api/project.py
index 3f357b4..5ac48ed 100644
--- a/api/project.py
+++ b/api/project.py
@@ -21,6 +21,7 @@ router = APIRouter()
 @router.post("/project/data")
 async def get_projects_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token),
                             background_tasks: BackgroundTasks = BackgroundTasks()):
+    background_tasks.add_task(update_project_count)
     search_query = request_data.get("search", "")
     page_index = request_data.get("pageIndex", 1)
     page_size = request_data.get("pageSize", 10)
@@ -57,7 +58,6 @@ async def get_projects_data(request_data: dict, db=Depends(get_mongo_db), _: dic
         results = await cursor.to_list(length=None)
         for result in results:
             result["AssetCount"] = result.get("AssetCount", 0)
-            background_tasks.add_task(update_project_count, id=result["id"])
         return results
 
     fetch_tasks = []
@@ -115,10 +115,14 @@ async def get_projects_all(db=Depends(get_mongo_db), _: dict = Depends(verify_to
         return {"message": "error","code":500}
 
 
-async def update_project_count(id):
-    async for db in get_mongo_db():
+async def update_project_count():
+    db = await get_mongo_db()
+    cursor = db.project.find({}, {"_id": 0, "id": {"$toString": "$_id"}})
+    results = await cursor.to_list(length=None)
+
+    async def update_count(id):
         query = {"project": {"$eq": id}}
-        total_count = await db['asset'].count_documents(query)
+        total_count = await db.asset.count_documents(query)
         update_document = {
             "$set": {
                 "AssetCount": total_count
@@ -126,6 +130,10 @@ async def update_project_count(id):
         }
         await db.project.update_one({"_id": ObjectId(id)}, update_document)
 
+    fetch_tasks = [update_count(r['id']) for r in results]
+
+    await asyncio.gather(*fetch_tasks)
+
 
 @router.post("/project/content")
 async def get_project_content(request_data: dict, db=Depends(get_mongo_db), _: dict = Depends(verify_token)):
diff --git a/api/sensitive.py b/api/sensitive.py
index bb80f34..e7928dd 100644
--- a/api/sensitive.py
+++ b/api/sensitive.py
@@ -33,7 +33,7 @@ async def get_sensitive_data(request_data: dict, db=Depends(get_mongo_db), _: di
         total_count = await db.SensitiveRule.count_documents(query)
 
         # Perform pagination query
-        cursor: AsyncIOMotorCursor = db.SensitiveRule.find(query).skip((page_index - 1) * page_size).limit(page_size).sort([("timestamp", DESCENDING)])
+        cursor: AsyncIOMotorCursor = db.SensitiveRule.find(query).skip((page_index - 1) * page_size).limit(page_size).sort([("_id", DESCENDING)])
         result = await cursor.to_list(length=None)
         if len(result) == 0:
             return {