169 lines
7.3 KiB
Diff
169 lines
7.3 KiB
Diff
From 8e19d92b6a484ddcf7ca7bf666ce21baa56ab326 Mon Sep 17 00:00:00 2001
|
|
From: rearcher <123781007@qq.com>
|
|
Date: Wed, 20 Dec 2023 17:21:22 +0800
|
|
Subject: [PATCH] fix TimedCorrectTask
|
|
|
|
---
|
|
apollo/cron/timed_correct_manager.py | 12 ++++--
|
|
apollo/database/proxy/task/base.py | 50 ++++++++++++++++++-----
|
|
apollo/database/proxy/task/timed_proxy.py | 8 ++++
|
|
apollo/tests/database/test_task.py | 2 +-
|
|
4 files changed, 57 insertions(+), 15 deletions(-)
|
|
|
|
diff --git a/apollo/cron/timed_correct_manager.py b/apollo/cron/timed_correct_manager.py
|
|
index db0b4c1..ae3a1e3 100644
|
|
--- a/apollo/cron/timed_correct_manager.py
|
|
+++ b/apollo/cron/timed_correct_manager.py
|
|
@@ -41,11 +41,15 @@ class TimedCorrectTask(TimedTask):
|
|
"""
|
|
Start the correct after the specified time of day.
|
|
"""
|
|
- LOGGER.info("Begin to correct the whole host in %s.", str(datetime.datetime.now()))
|
|
+ LOGGER.info(
|
|
+ "Begin to correct the status of timeout tasks and scan timeout host in %s.",
|
|
+ str(datetime.datetime.now()))
|
|
abnormal_task_ids, abnormal_host_ids = self.get_abnormal_task()
|
|
- self._update_host_status(abnormal_host_ids)
|
|
- with TimedProxy() as proxy:
|
|
- proxy.timed_correct_error_task_status(abnormal_task_ids)
|
|
+ if len(abnormal_host_ids) != 0:
|
|
+ self._update_host_status(abnormal_host_ids)
|
|
+ if len(abnormal_task_ids) != 0:
|
|
+ with TimedProxy() as proxy:
|
|
+ proxy.timed_correct_error_task_status(abnormal_task_ids)
|
|
|
|
@staticmethod
|
|
def _abnormal_task(tasks):
|
|
diff --git a/apollo/database/proxy/task/base.py b/apollo/database/proxy/task/base.py
|
|
index 840c140..a5ddede 100644
|
|
--- a/apollo/database/proxy/task/base.py
|
|
+++ b/apollo/database/proxy/task/base.py
|
|
@@ -861,17 +861,17 @@ class TaskProxy(TaskMysqlProxy, TaskEsProxy):
|
|
|
|
raise EsOperationError("Delete task from elasticsearch failed due to internal error.")
|
|
|
|
- def get_running_task_form_task_cve_host(self) -> list:
|
|
+ def get_running_task_form_hotpatch_remove_task(self) -> list:
|
|
"""
|
|
- Get all CVE repair tasks with running status under Username
|
|
+ Get all hotpatch remove tasks with running status under Username
|
|
|
|
Returns:
|
|
list: task id list
|
|
"""
|
|
- task_cve_query = (
|
|
- self.session.query(HotpatchRemoveTask).filter(HotpatchRemoveTask.status == TaskStatus.RUNNING).all()
|
|
+ hotpatch_remove_query = (
|
|
+ self.session.query(HotpatchRemoveTask.task_id).filter(HotpatchRemoveTask.status == TaskStatus.RUNNING).all()
|
|
)
|
|
- task_id_list = [task.task_id for task in task_cve_query]
|
|
+ task_id_list = [task.task_id for task in hotpatch_remove_query]
|
|
return task_id_list
|
|
|
|
def get_running_task_form_task_host_repo(self) -> list:
|
|
@@ -882,13 +882,39 @@ class TaskProxy(TaskMysqlProxy, TaskEsProxy):
|
|
list: task id list
|
|
"""
|
|
host_repo_query = (
|
|
- self.session.query(TaskHostRepoAssociation)
|
|
+ self.session.query(TaskHostRepoAssociation.task_id)
|
|
.filter(TaskHostRepoAssociation.status == TaskStatus.RUNNING)
|
|
.all()
|
|
)
|
|
task_id_list = [task.task_id for task in host_repo_query]
|
|
return task_id_list
|
|
|
|
+ def get_running_task_form_cve_fix_task(self) -> list:
|
|
+ """
|
|
+ Get all CVE fix tasks with running status
|
|
+
|
|
+ Returns:
|
|
+ list: task id list
|
|
+ """
|
|
+ cve_fix_query = (
|
|
+ self.session.query(CveFixTask.task_id).filter(CveFixTask.status == TaskStatus.RUNNING).all()
|
|
+ )
|
|
+ task_id_list = [task.task_id for task in cve_fix_query]
|
|
+ return task_id_list
|
|
+
|
|
+ def get_running_task_form_cve_rollback_task(self) -> list:
|
|
+ """
|
|
+ Get all CVE rollback tasks with running status
|
|
+
|
|
+ Returns:
|
|
+ list: task id list
|
|
+ """
|
|
+ cve_rollback_query = (
|
|
+ self.session.query(CveRollbackTask.task_id).filter(CveRollbackTask.status == TaskStatus.RUNNING).all()
|
|
+ )
|
|
+ task_id_list = [task.task_id for task in cve_rollback_query]
|
|
+ return task_id_list
|
|
+
|
|
def get_scanning_status_and_time_from_host(self) -> list:
|
|
"""
|
|
Get all host id and time with scanning status from the host table
|
|
@@ -907,13 +933,17 @@ class TaskProxy(TaskMysqlProxy, TaskEsProxy):
|
|
Returns:
|
|
list: Each element is a task information, including the task ID, task type, creation time
|
|
"""
|
|
- task_cve_id_list = self.get_running_task_form_task_cve_host()
|
|
- task_repo_id_list = self.get_running_task_form_task_host_repo()
|
|
host_info_list = self.get_scanning_status_and_time_from_host()
|
|
- task_id_list = task_cve_id_list + task_repo_id_list
|
|
+
|
|
+ task_cve_id_list = self.get_running_task_form_hotpatch_remove_task()
|
|
+ task_repo_id_list = self.get_running_task_form_task_host_repo()
|
|
+ task_cve_fix_list = self.get_running_task_form_cve_fix_task()
|
|
+ task_cve_rollback_list = self.get_running_task_form_cve_rollback_task()
|
|
+
|
|
+ task_id_list = task_cve_id_list + task_repo_id_list + task_cve_fix_list + task_cve_rollback_list
|
|
|
|
task_query = self.session.query(Task).filter(Task.task_id.in_(task_id_list)).all()
|
|
- running_task_list = [(task.task_id, task.create_time) for task in task_query]
|
|
+ running_task_list = [(task.task_id, task.latest_execute_time) for task in task_query]
|
|
return running_task_list, host_info_list
|
|
|
|
def validate_cves(self, cve_id: list) -> bool:
|
|
diff --git a/apollo/database/proxy/task/timed_proxy.py b/apollo/database/proxy/task/timed_proxy.py
|
|
index 436c3bd..fd396d1 100644
|
|
--- a/apollo/database/proxy/task/timed_proxy.py
|
|
+++ b/apollo/database/proxy/task/timed_proxy.py
|
|
@@ -22,6 +22,8 @@ from apollo.conf.constant import TaskStatus
|
|
from apollo.database.table import (
|
|
HotpatchRemoveTask,
|
|
TaskHostRepoAssociation,
|
|
+ CveFixTask,
|
|
+ CveRollbackTask,
|
|
)
|
|
|
|
|
|
@@ -42,6 +44,12 @@ class TimedProxy(MysqlProxy):
|
|
self.session.query(TaskHostRepoAssociation).filter(TaskHostRepoAssociation.task_id.in_(task_ids)).update(
|
|
{TaskHostRepoAssociation.status: TaskStatus.UNKNOWN}, synchronize_session=False
|
|
)
|
|
+ self.session.query(CveFixTask).filter(CveFixTask.task_id.in_(task_ids)).update(
|
|
+ {CveFixTask.status: TaskStatus.UNKNOWN}, synchronize_session=False
|
|
+ )
|
|
+ self.session.query(CveRollbackTask).filter(CveRollbackTask.task_id.in_(task_ids)).update(
|
|
+ {CveRollbackTask.status: TaskStatus.UNKNOWN}, synchronize_session=False
|
|
+ )
|
|
self.session.commit()
|
|
except SQLAlchemyError as error:
|
|
self.session.rollback()
|
|
diff --git a/apollo/tests/database/test_task.py b/apollo/tests/database/test_task.py
|
|
index 35d923b..ceb84ab 100644
|
|
--- a/apollo/tests/database/test_task.py
|
|
+++ b/apollo/tests/database/test_task.py
|
|
@@ -401,7 +401,7 @@ class TestTaskMysqlFirst(DatabaseTestCase):
|
|
|
|
def test_get_running_task_form_task_cve_host(self):
|
|
self.assertEqual(
|
|
- self.task_database.get_running_task_form_task_cve_host(),
|
|
+ self.task_database.get_running_task_form_hotpatch_remove_task(),
|
|
["1111111111poiuytrewqasdfghjklmnb"],
|
|
)
|
|
|
|
--
|
|
2.33.0
|
|
|