refactor: remove cl graph lookup + add Instances get_resource_map() (#…

…301) * WIP * WIP * WIP * WIP * WIP
Azure · Aug 8, 2024 · 0528703 · 0528703
1 parent b8ef01f
commit 0528703
Show file tree

Hide file tree

Showing 8 changed files with 132 additions and 43 deletions.
diff --git a/azext_edge/edge/providers/orchestration/common.py b/azext_edge/edge/providers/orchestration/common.py
@@ -27,6 +27,9 @@
 KEYVAULT_DATAPLANE_API_VERSION = "7.4"
 KEYVAULT_CLOUD_API_VERSION = "2022-07-01"
 
+# Custom Locations KPIs
+CUSTOM_LOCATIONS_API_VERSION = "2021-08-31-preview"
+
 
 class MqMode(Enum):
     auto = "auto"

diff --git a/azext_edge/edge/providers/orchestration/resource_map.py b/azext_edge/edge/providers/orchestration/resource_map.py
@@ -42,21 +42,19 @@ def __init__(self):
 
 
 class IoTOperationsResourceMap:
-    def __init__(self, cmd, cluster_name: str, resource_group_name: str):
+    def __init__(self, cmd, cluster_name: str, resource_group_name: str, defer_refresh: bool = False):
         from azure.cli.core.commands.client_factory import get_subscription_id
 
-        self.cmd = cmd
-        self.cluster_name = cluster_name
-        self.resource_group_name = resource_group_name
         self.subscription_id = get_subscription_id(cli_ctx=cmd.cli_ctx)
         self.connected_cluster = ConnectedCluster(
             cmd=cmd,
             subscription_id=self.subscription_id,
-            cluster_name=self.cluster_name,
-            resource_group_name=self.resource_group_name,
+            cluster_name=cluster_name,
+            resource_group_name=resource_group_name,
         )
         self._cluster_container = ClusterContainer()
-        self.refresh_resource_state()
+        if not defer_refresh:
+            self.refresh_resource_state()
 
     @property
     def extensions(self) -> List[IoTOperationsResource]:
@@ -130,7 +128,7 @@ def refresh_resource_state(self):
         self._cluster_container = refreshed_cluster_container
 
     def build_tree(self, category_color: str = "red") -> Tree:
-        tree = Tree(f"[green]{self.cluster_name}")
+        tree = Tree(f"[green]{self.connected_cluster.cluster_name}")
         extensions_node = tree.add(label=f"[{category_color}]extensions")
         [extensions_node.add(ext.display_name) for ext in self.extensions]
 

diff --git a/azext_edge/edge/providers/orchestration/resources/instances.py b/azext_edge/edge/providers/orchestration/resources/instances.py
@@ -10,28 +10,26 @@
 from rich import print
 from rich.console import Console
 
-from ....util.az_client import get_iotops_mgmt_client, parse_resource_id, wait_for_terminal_state
+from ....util.az_client import (
+    get_iotops_mgmt_client,
+    get_resource_client,
+    parse_resource_id,
+    wait_for_terminal_state,
+)
 from ....util.queryable import Queryable
+from ..common import CUSTOM_LOCATIONS_API_VERSION
+from ..resource_map import IoTOperationsResourceMap
 
 logger = get_logger(__name__)
 
 
-QUERIES = {
-    "get_cl_from_instance": """
-        resources
-        | where type =~ 'microsoft.extendedlocation/customlocations'
-        | where id =~ '{resource_id}'
-        | project id, name, properties
-        """
-}
-
-
 class Instances(Queryable):
     def __init__(self, cmd):
         super().__init__(cmd=cmd)
         self.iotops_mgmt_client = get_iotops_mgmt_client(
             subscription_id=self.default_subscription_id,
         )
+        self.resource_client = get_resource_client(subscription_id=self.default_subscription_id)
         self.console = Console()
 
     def show(self, name: str, resource_group_name: str, show_tree: Optional[bool] = None) -> Optional[dict]:
@@ -50,28 +48,25 @@ def list(self, resource_group_name: Optional[str] = None) -> Iterable[dict]:
         return self.iotops_mgmt_client.instance.list_by_subscription()
 
     def _show_tree(self, instance: dict):
-        custom_location = self.get_associated_cl(instance)
-        if not custom_location:
-            logger.warning("Unable to process the resource tree.")
-            return
-
-        resource_id_container = parse_resource_id(custom_location["properties"]["hostResourceId"])
-
-        # Currently resource map will query cluster state upon init
-        # therefore we only use it when necessary to save cycles.
-        from ..resource_map import IoTOperationsResourceMap
-
+        resource_map = self.get_resource_map(instance)
         with self.console.status("Working..."):
-            resource_map = IoTOperationsResourceMap(
-                cmd=self.cmd,
-                cluster_name=resource_id_container.resource_name,
-                resource_group_name=resource_id_container.resource_group_name,
-            )
+            resource_map.refresh_resource_state()
         print(resource_map.build_tree(category_color="cyan"))
 
-    def get_associated_cl(self, instance: dict) -> dict:
-        return self.query(
-            QUERIES["get_cl_from_instance"].format(resource_id=instance["extendedLocation"]["name"]), first=True
+    def _get_associated_cl(self, instance: dict) -> dict:
+        return self.resource_client.resources.get_by_id(
+            resource_id=instance["extendedLocation"]["name"], api_version=CUSTOM_LOCATIONS_API_VERSION
+        ).as_dict()
+
+    def get_resource_map(self, instance: dict) -> IoTOperationsResourceMap:
+        custom_location = self._get_associated_cl(instance)
+        resource_id_container = parse_resource_id(custom_location["properties"]["hostResourceId"])
+
+        return IoTOperationsResourceMap(
+            cmd=self.cmd,
+            cluster_name=resource_id_container.resource_name,
+            resource_group_name=resource_id_container.resource_group_name,
+            defer_refresh=True,
         )
 
     def update(

diff --git a/azext_edge/edge/util/az_client.py b/azext_edge/edge/util/az_client.py
@@ -110,6 +110,7 @@ def get_default_logging_policy() -> HttpLoggingPolicy:
     http_logging_policy.allowed_query_params.add("api-version")
     http_logging_policy.allowed_query_params.add("$filter")
     http_logging_policy.allowed_query_params.add("$expand")
+    http_logging_policy.allowed_header_names.add("x-ms-correlation-request-id")
 
     return http_logging_policy
 

diff --git a/azext_edge/tests/edge/orchestration/resources/conftest.py b/azext_edge/tests/edge/orchestration/resources/conftest.py
@@ -65,3 +65,13 @@ def get_mock_resource(
         },
         "type": QUALIFIED_INSTANCE_TYPE,
     }
+
+
+def get_resource_id(
+    resource_path: str,
+    resource_group_name: str,
+    resource_provider: Optional[str] = RESOURCE_PROVIDER,
+) -> str:
+    return get_base_endpoint(
+        resource_group_name=resource_group_name, resource_path=resource_path, resource_provider=resource_provider
+    ).split("?")[0][len(BASE_URL) :]
diff --git a/azext_edge/tests/edge/orchestration/resources/test_instances_unit.py b/azext_edge/tests/edge/orchestration/resources/test_instances_unit.py
@@ -12,9 +12,14 @@
 import responses
 
 from azext_edge.edge.commands_edge import list_instances, show_instance, update_instance
+from azext_edge.edge.providers.orchestration.resources import Instances
 
 from ....generators import generate_random_string
-from .conftest import get_base_endpoint, get_mock_resource
+from .conftest import get_base_endpoint, get_mock_resource, get_resource_id, BASE_URL
+
+
+CUSTOM_LOCATION_RP = "Microsoft.ExtendedLocation"
+CONNECTED_CLUSTER_RP = "Microsoft.Kubernetes"
 
 
 def get_instance_endpoint(resource_group_name: Optional[str] = None, instance_name: Optional[str] = None) -> str:
@@ -24,6 +29,15 @@ def get_instance_endpoint(resource_group_name: Optional[str] = None, instance_na
     return get_base_endpoint(resource_group_name=resource_group_name, resource_path=resource_path)
 
 
+def get_cl_endpoint(resource_group_name: Optional[str] = None, cl_name: Optional[str] = None) -> str:
+    resource_path = "/customLocations"
+    if cl_name:
+        resource_path += f"/{cl_name}"
+    return get_base_endpoint(
+        resource_group_name=resource_group_name, resource_path=resource_path, resource_provider=CUSTOM_LOCATION_RP
+    )
+
+
 def get_mock_instance_record(name: str, resource_group_name: str) -> dict:
     return get_mock_resource(
         name=name,
@@ -32,6 +46,32 @@ def get_mock_instance_record(name: str, resource_group_name: str) -> dict:
     )
 
 
+def get_mock_cl_record(name: str, resource_group_name: str) -> dict:
+    resource = get_mock_resource(
+        name=name,
+        properties={
+            "hostResourceId": get_resource_id(
+                resource_path="/connectedClusters/mycluster",
+                resource_group_name=resource_group_name,
+                resource_provider=CONNECTED_CLUSTER_RP,
+            ),
+            "namespace": "azure-iot-operations",
+            "displayName": generate_random_string(),
+            "provisioningState": "Succeeded",
+            "clusterExtensionIds": [
+                generate_random_string(),
+                generate_random_string(),
+            ],
+            "authentication": {},
+        },
+        resource_group_name=resource_group_name,
+    )
+    resource.pop("extendedLocation")
+    resource.pop("systemData")
+    resource.pop("resourceGroup")
+    return resource
+
+
 def test_instance_show(mocked_cmd, mocked_responses: responses):
     instance_name = generate_random_string()
     resource_group_name = generate_random_string()
@@ -51,6 +91,35 @@ def test_instance_show(mocked_cmd, mocked_responses: responses):
     assert len(mocked_responses.calls) == 1
 
 
+def test_instance_get_resource_map(mocker, mocked_cmd, mocked_responses: responses):
+    cl_name = generate_random_string()
+    instance_name = generate_random_string()
+    resource_group_name = generate_random_string()
+
+    mock_instance_record = get_mock_instance_record(name=instance_name, resource_group_name=resource_group_name)
+    mock_cl_record = get_mock_cl_record(name=cl_name, resource_group_name=resource_group_name)
+
+    mocked_responses.add(
+        method=responses.GET,
+        url=f"{BASE_URL}{mock_instance_record['extendedLocation']['name']}",
+        json=mock_cl_record,
+        status=200,
+        content_type="application/json",
+    )
+
+    host_resource_id: str = mock_cl_record["properties"]["hostResourceId"]
+    host_resource_parts = host_resource_id.split("/")
+
+    instances = Instances(mocked_cmd)
+    resource_map = instances.get_resource_map(mock_instance_record)
+    assert resource_map.subscription_id == host_resource_parts[2]
+
+    assert resource_map.connected_cluster.subscription_id == host_resource_parts[2]
+    assert resource_map.connected_cluster.resource_group_name == host_resource_parts[4]
+    assert resource_map.connected_cluster.cluster_name == host_resource_parts[-1]
+    assert len(mocked_responses.calls) == 1
+
+
 @pytest.mark.parametrize(
     "resource_group_name",
     [None, generate_random_string()],

diff --git a/azext_edge/tests/edge/orchestration/test_connected_cluster_unit.py b/azext_edge/tests/edge/orchestration/test_connected_cluster_unit.py
@@ -156,6 +156,10 @@ def test_connected_cluster_attr(
         cmd=mocked_cmd, subscription_id=sub, cluster_name=cluster_name, resource_group_name=rg_name
     )
 
+    assert connected_cluster.subscription_id == sub
+    assert connected_cluster.cluster_name == cluster_name
+    assert connected_cluster.resource_group_name == rg_name
+
     resource_id = connected_cluster.resource_id
     assert resource_id == (
         f"/subscriptions/{sub}/resourceGroups/{rg_name}"

diff --git a/azext_edge/tests/edge/orchestration/test_resource_map_unit.py b/azext_edge/tests/edge/orchestration/test_resource_map_unit.py
@@ -42,11 +42,17 @@ def _generate_records(count: int = 1) -> List[dict]:
 
 def _assemble_connected_cluster_mock(
     cluster_mock: Mock,
+    sub: str,
+    cluster_name: str,
+    rg_name: str,
     extensions: Optional[List[dict]],
     custom_locations: Optional[List[dict]],
     resources: Optional[List[dict]],
     sync_rules: Optional[List[dict]],
 ):
+    cluster_mock().subscription_id = sub
+    cluster_mock().cluster_name = cluster_name
+    cluster_mock().resource_group_name = rg_name
     cluster_mock().get_aio_extensions.return_value = extensions
     cluster_mock().get_aio_custom_locations.return_value = custom_locations
     cluster_mock().get_aio_resources.return_value = resources
@@ -76,18 +82,21 @@ def test_resource_map(
         IoTOperationsResourceMap,
     )
 
+    sub = get_zeroed_subscription()
+    cluster_name = generate_random_string()
+    rg_name = generate_random_string()
+
     _assemble_connected_cluster_mock(
         cluster_mock=mocked_connected_cluster,
+        sub=sub,
+        cluster_name=cluster_name,
+        rg_name=rg_name,
         extensions=expected_extensions,
         custom_locations=expected_custom_locations,
         resources=expected_resources,
         sync_rules=expected_resource_sync_rules,
     )
 
-    sub = get_zeroed_subscription()
-    cluster_name = generate_random_string()
-    rg_name = generate_random_string()
-
     resource_map = IoTOperationsResourceMap(cmd=mocked_cmd, cluster_name=cluster_name, resource_group_name=rg_name)
 
     assert resource_map.subscription_id == sub