fix 计划任务scan更新报错

This commit is contained in:
Autumn.home 2024-07-16 20:07:57 +08:00
parent 74a9bcaa64
commit afeffa6d3f
3 changed files with 14 additions and 4 deletions

View File

@ -5,6 +5,7 @@
# @time : 2024/4/14 17:14 # @time : 2024/4/14 17:14
# ------------------------------------------- # -------------------------------------------
import json import json
import traceback
from bson import ObjectId from bson import ObjectId
from fastapi import APIRouter, Depends from fastapi import APIRouter, Depends
@ -74,7 +75,7 @@ async def asset_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Dep
.skip((page_index - 1) * page_size) .skip((page_index - 1) * page_size)
.limit(page_size)) .limit(page_size))
.sort([("timestamp", DESCENDING)])) .sort([("timestamp", DESCENDING)]))
result = await cursor.to_list(length=None) result = cursor.to_list(length=None)
result_list = [] result_list = []
for r in result: for r in result:
tmp = {} tmp = {}
@ -127,6 +128,7 @@ async def asset_data(request_data: dict, db=Depends(get_mongo_db), _: dict = Dep
} }
except Exception as e: except Exception as e:
logger.error(str(e)) logger.error(str(e))
logger.error(traceback.format_exc())
# Handle exceptions as needed # Handle exceptions as needed
return {"message": "error", "code": 500} return {"message": "error", "code": 500}

View File

@ -5,6 +5,7 @@
# @time : 2024/6/16 16:11 # @time : 2024/6/16 16:11
# ------------------------------------------- # -------------------------------------------
import os import os
import traceback
from bson import ObjectId from bson import ObjectId
from fastapi import APIRouter, Depends, BackgroundTasks from fastapi import APIRouter, Depends, BackgroundTasks
@ -150,7 +151,7 @@ async def fetch_data(db, collection, query, quantity, project_list):
{"$project": {"_id": 0, "vulnid": 0}} {"$project": {"_id": 0, "vulnid": 0}}
] ]
cursor = db[collection].aggregate(pipeline) cursor = await db[collection].aggregate(pipeline)
return cursor return cursor
@ -300,6 +301,7 @@ async def export_data_from_mongodb(quantity, query, file_name, index):
await db.export.update_one({"file_name": file_name}, update_document) await db.export.update_one({"file_name": file_name}, update_document)
except Exception as e: except Exception as e:
logger.error(str(e)) logger.error(str(e))
logger.error(traceback.format_exc())
update_document = { update_document = {
"$set": { "$set": {
"state": 2, "state": 2,

View File

@ -4,6 +4,8 @@
# @version: # @version:
import asyncio import asyncio
import json import json
import traceback
from loguru import logger from loguru import logger
from bson import ObjectId from bson import ObjectId
from fastapi import APIRouter, Depends, BackgroundTasks from fastapi import APIRouter, Depends, BackgroundTasks
@ -305,9 +307,13 @@ async def update_task_data(request_data: dict, db=Depends(get_mongo_db), _: dict
await db.ScheduledTasks.update_one({"id": task_id}, {"$set": {'state': False}}) await db.ScheduledTasks.update_one({"id": task_id}, {"$set": {'state': False}})
if newScheduledTasks: if newScheduledTasks:
if hour != old_hour: if hour != old_hour:
await db.ScheduledTasks.update_one({"id": task_id}, {"$set": {'hour': hour}})
job = scheduler.get_job(task_id) job = scheduler.get_job(task_id)
if job is not None: if job is not None:
scheduler.remove_job(job) scheduler.remove_job(task_id)
scheduler.add_job(scheduler_scan_task, 'interval', hours=hour, args=[task_id],
id=str(task_id), jobstore='mongo')
else:
scheduler.add_job(scheduler_scan_task, 'interval', hours=hour, args=[task_id], scheduler.add_job(scheduler_scan_task, 'interval', hours=hour, args=[task_id],
id=str(task_id), jobstore='mongo') id=str(task_id), jobstore='mongo')
@ -324,7 +330,7 @@ async def update_task_data(request_data: dict, db=Depends(get_mongo_db), _: dict
except Exception as e: except Exception as e:
logger.error(str(e)) logger.error(str(e))
# Handle exceptions as needed logger.error(traceback.format_exc())
return {"message": "error", "code": 500} return {"message": "error", "code": 500}