feature: 订阅对于业务范围的优先级调整 (closed #1944)

apps/backend/subscription/tools.py

@@ -26,6 +26,7 @@
 from django.db.models import Q
 from django.utils import timezone
 
+from apps.backend.constants import InstNodeType
 from apps.backend.subscription import task_tools
 from apps.backend.subscription.commons import get_host_by_inst, list_biz_hosts
 from apps.backend.subscription.constants import SUBSCRIPTION_SCOPE_CACHE_TIME
@@ -673,13 +674,24 @@ def get_host_relation(bk_biz_id, nodes):
     return data
 
 
+def fill_nodes_biz_info(nodes: typing.List[typing.Dict[str, typing.Union[str, int]]]):
+    for node in nodes:
+        bk_biz_id: typing.Optional[int] = node.get("bk_biz_id")
+        bk_obj_id: typing.Optional[str] = node.get("bk_obj_id")
+        bk_inst_id: typing.Optional[int] = node.get("bk_inst_id")
+        if not bk_biz_id and bk_obj_id == InstNodeType.BIZ and bk_inst_id:
+            # 兜底类似 [{"ip": 127.0.0.1, "bk_inst_id": x, "bk_obj_id": biz}] 场景
+            # 填充为 [{"ip": 127.0.0.1, "bk_inst_id": x, "bk_obj_id": biz, "bk_biz_id": x}]
+            node["bk_biz_id"] = bk_inst_id
+
+
 def support_multi_biz(get_instances_by_scope_func):
     """支持scope多范围"""
 
     @wraps(get_instances_by_scope_func)
     def wrapper(scope: Dict[str, Union[Dict, Any]], *args, **kwargs) -> Dict[str, Dict[str, Union[Dict, Any]]]:
-        if scope.get("bk_biz_id") is not None:
-            return get_instances_by_scope_func(scope, **kwargs)
+        nodes: typing.List[typing.Dict[str, typing.Union[str, int]]] = scope["nodes"]
+        fill_nodes_biz_info(nodes=nodes)
         # 兼容只传bk_host_id的情况
         if (
             scope["object_type"] == models.Subscription.ObjectType.HOST
@@ -689,7 +701,7 @@ def wrapper(scope: Dict[str, Union[Dict, Any]], *args, **kwargs) -> Dict[str, Di
                 return get_instances_by_scope_func(scope, **kwargs)
 
         instance_id_info_map = {}
-        nodes = sorted(scope["nodes"], key=lambda node: node["bk_biz_id"])
+        nodes = sorted(scope["nodes"], key=lambda node: node.get("bk_biz_id") or scope.get("bk_biz_id"))
         params_list = [
             {
                 "scope": {
@@ -701,7 +713,7 @@ def wrapper(scope: Dict[str, Union[Dict, Any]], *args, **kwargs) -> Dict[str, Di
                 },
                 **kwargs,
             }
-            for bk_biz_id, nodes in groupby(nodes, key=lambda x: x["bk_biz_id"])
+            for bk_biz_id, nodes in groupby(nodes, key=lambda x: x.get("bk_biz_id") or scope.get("bk_biz_id"))
         ]
         results = request_multi_thread(get_instances_by_scope_func, params_list, get_data=lambda x: [x])
         for result in results:

apps/backend/tests/subscription/test_tools.py

@@ -22,7 +22,7 @@
     CmdbClient,
     list_biz_hosts_without_info_client,
 )
-from apps.node_man import models, constants
+from apps.node_man import constants, models
 
 # 全局使用的mock
 run_task = mock.patch("apps.backend.subscription.tasks.run_subscription_task").start()
@@ -203,3 +203,57 @@ def test_get_empty_list_instance_selector_scope(self):
             }
         )
         self.assertEqual(len(list(instances.keys())), 0)
+
+    def test_sub_biz_priority(self):
+        # 之前订阅优先使用 scope.bk_biz_id 作为整个订阅的业务范围，后面调整为优先使用 scope.nodes 内的业务范围
+        instances = get_instances_by_scope(
+            {
+                "object_type": "HOST",
+                "node_type": "INSTANCE",
+                "bk_biz_id": 2,
+                "nodes": [
+                    {"ip": "127.0.0.1", "bk_cloud_id": 0, "bk_supplier_id": 0, "bk_biz_id": 2},
+                    {"ip": "127.0.0.2", "bk_cloud_id": 0, "bk_supplier_id": 0, "bk_biz_id": 3},
+                ],
+            }
+        )
+
+        self.assertEqual(len(list(instances.keys())), 1)
+        instances = get_instances_by_scope(
+            {
+                "object_type": "HOST",
+                "node_type": "INSTANCE",
+                "nodes": [
+                    {"ip": "127.0.0.1", "bk_cloud_id": 0, "bk_supplier_id": 0, "bk_inst_id": 2, "bk_obj_id": "biz"},
+                ],
+            }
+        )
+        self.assertEqual(len(list(instances.keys())), 1)
+
+        request_handler = mock.MagicMock(
+            return_value=[
+                {"ip": "127.0.0.1", "bk_cloud_id": 0, "bk_supplier_id": 0, "bk_inst_id": 2, "bk_obj_id": "biz"}
+            ]
+        )
+        with mock.patch("apps.backend.subscription.tools.request_multi_thread", request_handler):
+            get_instances_by_scope(
+                {
+                    "object_type": "HOST",
+                    "node_type": "INSTANCE",
+                    "nodes": [
+                        {"ip": "127.0.0.1", "bk_cloud_id": 0, "bk_supplier_id": 0, "bk_inst_id": 2, "bk_obj_id": "biz"},
+                        {"ip": "127.0.0.1", "bk_cloud_id": 0, "bk_supplier_id": 0, "bk_biz_id": 3},
+                        {"ip": "127.0.0.4", "bk_cloud_id": 0, "bk_supplier_id": 0, "bk_biz_id": 3},
+                        {"ip": "127.0.0.5", "bk_cloud_id": 0, "bk_supplier_id": 0, "bk_biz_id": 3},
+                        {"ip": "127.0.0.2", "bk_cloud_id": 0, "bk_supplier_id": 0, "bk_biz_id": 4},
+                        {"ip": "127.0.0.6", "bk_cloud_id": 0, "bk_supplier_id": 0, "bk_biz_id": 4},
+                    ],
+                }
+            )
+
+        self.assertEqual(len(request_handler.call_args[0][1]), 3)
+        for instance in request_handler.call_args[0][1]:
+            if instance["scope"]["bk_biz_id"] == 3:
+                self.assertEqual(len(instance["scope"]["nodes"]), 3)
+            if instance["scope"]["bk_biz_id"] == 4:
+                self.assertEqual(len(instance["scope"]["nodes"]), 2)