optimization: 安装Gse2.0 查询Agent状态优化 (closed TencentBlueKing#1883)

apps/backend/components/collections/agent_new/get_agent_status.py

@@ -9,7 +9,7 @@
 specific language governing permissions and limitations under the License.
 """
 from collections import defaultdict
-from typing import Dict, List, Set, Union
+from typing import Dict, List, Set, Union, Optional
 
 from django.utils.translation import ugettext_lazy as _
 
@@ -67,13 +67,20 @@ def _schedule(self, data, parent_data, callback_data=None):
 
         # 构造 gse 请求参数
         hosts: List[Dict[str, Union[int, str]]] = []
+        need_to_quey_host_agent_id_map: Dict[int, int] = {}
         for host_id in host_ids_need_to_query:
+            sub_inst_id = common_data.host_id__sub_inst_id_map[host_id]
+            sub_inst = common_data.sub_inst_id__sub_inst_obj_map[sub_inst_id]
             host_obj = common_data.host_id_obj_map[host_id]
+
+            bk_agent_id: Optional[str] = sub_inst.instance_info["host"].get("bk_agent_id") or host_obj.bk_agent_id
+            need_to_quey_host_agent_id_map[host_id] = bk_agent_id
+
             hosts.append(
                 {
                     "ip": host_obj.inner_ip or host_obj.inner_ipv6,
                     "bk_cloud_id": host_obj.bk_cloud_id,
-                    "bk_agent_id": host_obj.bk_agent_id,
+                    "bk_agent_id": bk_agent_id,
                 }
             )
 
@@ -93,7 +100,7 @@ def _schedule(self, data, parent_data, callback_data=None):
         }
         for host_id in host_ids_need_to_query:
             host_obj = common_data.host_id_obj_map[host_id]
-            agent_id = common_data.gse_api_helper.get_agent_id(host_obj)
+            agent_id = need_to_quey_host_agent_id_map[host_id]
             # bk_agent_alive 默认值为 None 的背景：expect_status 是可传入的，不能设定一个明确状态作为默认值，不然可能与 expect_status 一致
             # 误判为当前 Agent 状态已符合预期
             agent_state_info = agent_id__agent_state_info_map.get(agent_id, {"version": "", "bk_agent_alive": None})

apps/backend/components/collections/agent_new/upgrade_to_agent_id.py

@@ -9,7 +9,7 @@
 specific language governing permissions and limitations under the License.
 """
 
-from typing import Dict, List, Union
+from typing import Dict, List, Union, Any, Optional
 
 from django.utils.translation import ugettext_lazy as _
 
@@ -18,21 +18,22 @@
 
 class UpgradeToAgentIDService(AgentBaseService):
     def _execute(self, data, parent_data, common_data: AgentCommonData):
-        upgrade_hosts: List[Dict[str, Union[int, str]]] = []
         cloud_ip__sub_inst_id_map: Dict[str, int] = {}
-        # 如果主机有AgentID，则调用 upgrade_to_agent_id 将基于 Host IP 的配置升级到基于 Agent-ID 的配置
-        for host_id, sub_inst_id in common_data.host_id__sub_inst_id_map.items():
-            host = common_data.host_id_obj_map[host_id]
-            upgrade_hosts.append(
-                {"ip": host.inner_ip, "bk_cloud_id": host.bk_cloud_id, "bk_agent_id": host.bk_agent_id}
-            )
+        # 如果主机有AgentID，则调用 upgrade_to_agent_id 将基于 Host IP 的配置升级到基于 Agent-ID 的配
+        upgrade_hosts: List[Dict[str, Union[int, str]]] = []
+
+        for host in common_data.host_id_obj_map.values():
+            sub_inst_id = common_data.host_id__sub_inst_id_map[host.bk_host_id]
+            sub_inst = common_data.sub_inst_id__sub_inst_obj_map[sub_inst_id]
+            bk_agent_id: Optional[str] = sub_inst.instance_info["host"].get("bk_agent_id") or host.bk_agent_id
             cloud_ip__sub_inst_id_map[f"{host.bk_cloud_id}:{host.inner_ip}"] = sub_inst_id
 
-        upgrade_hosts = [
-            {"ip": host.inner_ip, "bk_cloud_id": host.bk_cloud_id, "bk_agent_id": host.bk_agent_id}
-            for host in common_data.host_id_obj_map.values()
-            if host.bk_agent_id
-        ]
+            if bk_agent_id:
+                upgrade_hosts.append({
+                    "ip": host.inner_ip, 
+                    "bk_cloud_id": host.bk_cloud_id, 
+                    "bk_agent_id": bk_agent_id,
+                })
 
         if not upgrade_hosts:
             return True

apps/backend/components/collections/plugin.py

@@ -145,6 +145,7 @@ def get_common_data(cls, data):
             subscription_instance_ids=common_data.subscription_instance_ids,
             sub_inst_id__host_id_map=common_data.sub_inst_id__host_id_map,
             host_id__sub_inst_id_map=common_data.host_id__sub_inst_id_map,
+            sub_inst_id__sub_inst_obj_map=common_data.sub_inst_id__sub_inst_obj_map,
             # Plugin 新增的公共数据
             process_statuses=process_statuses,
             target_host_objs=target_host_objs,