[AZINTS-3933] Use ARG to collect existing Log Forwarders

azure/integration_quickstart/src/azure_integration_quickstart/setup.py

@@ -26,7 +26,7 @@
 import traceback
 from typing import Any, Generator, Literal, Optional, TypeVar, TypedDict, Union
 
-from azure_logging_install.existing_lfo import check_existing_lfo
+from azure_logging_install.existing_lfo import find_existing_lfo_control_planes
 
 # General util
 
@@ -439,7 +439,7 @@ def collect_available_scopes(connection: HTTPSConnection, workflow_id: str) -> t
 
 def collect_log_forwarders(subscriptions: list[Scope], step_metadata: dict) -> None:
     scope_id_to_name = { s.id:s.name for s in subscriptions }
-    forwarders = check_existing_lfo(set(scope_id_to_name), scope_id_to_name)
+    forwarders = find_existing_lfo_control_planes(scope_id_to_name)
     step_metadata["log_forwarders"] = [asdict(forwarder) for forwarder in forwarders.values()]
 
 def receive_user_selections(connection: HTTPSConnection, workflow_id: str) -> tuple[Sequence[Scope], dict]:

azure/logging_install/src/azure_logging_install/az_cmd.py

@@ -85,7 +85,7 @@ def list_users_subscriptions() -> dict[str, str]:
     return {sub["id"]: sub["name"] for sub in subs_json}
 
 
-def execute(az_cmd: AzCmd) -> str:
+def execute(az_cmd: AzCmd, can_fail: bool = False) -> str:
     """Run an Azure CLI command and return output or raise error."""
 
     full_command = az_cmd.str()
@@ -97,7 +97,7 @@ def execute(az_cmd: AzCmd) -> str:
             result = subprocess.run(
                 full_command, shell=True, check=True, capture_output=True, text=True
             )
-            if result.returncode != 0:
+            if result.returncode != 0 and not can_fail:
                 log.error(f"Command failed: {full_command}")
                 log.error(result.stderr)
                 raise RuntimeError(f"Command failed: {full_command}")
@@ -132,8 +132,10 @@ def execute(az_cmd: AzCmd) -> str:
                 if error_details:
                     raise AccessError(f"{error_message}: {error_details}") from e
                 raise AccessError(error_message) from e
+            if can_fail:
+                return ""
             log.error(f"Command failed: {full_command}")
             log.error(e.stderr)
             raise RuntimeError(f"Command failed: {full_command}") from e
 
-    raise SystemExit(1)  # unreachable
+    raise SystemExit(1)  # unreachable

azure/logging_install/src/azure_logging_install/existing_lfo.py

@@ -1,84 +1,100 @@
 from dataclasses import dataclass
 from json import JSONDecodeError, loads
 from logging import getLogger
-from typing import Final
+from typing import Final, Optional
 
 from .az_cmd import AzCmd, execute
-from .configuration import Configuration
 
 log = getLogger("installer")
 
 CONTROL_PLANE_RESOURCES_TASK_PREFIX: Final = "resources-task"
 
+@dataclass(frozen=True)
+class LfoControlPlane:
+    subscription: tuple[str, str] # id, name
+    resource_group: str
+    region: str
 
 @dataclass(frozen=True)
 class LfoMetadata:
+    control_plane: LfoControlPlane
     monitored_subs: dict[str, str]
-    control_plane_sub: tuple[str, str]
-    control_plane_rg: str
+
+def find_existing_lfo_control_planes(
+    sub_id_to_name: dict[str, str], subscriptions: Optional[set[str]] = None
+) -> dict[str, LfoControlPlane]:
+    """Find existing lfo control planes in the tenant. If `subscriptions` is specified, search is limited to these subscriptions."""
+    if subscriptions is not None:
+        if len(subscriptions) == 0:
+            return {} # searching empty set of subscriptions
+        subscriptions_clause = " and subscriptionId in ({})".format(", ".join(["'{}'".format(subscription_id) for subscription_id in subscriptions]))
+    else:
+        subscriptions_clause = ""
+
+    # make sure azure resource graph extension is installed
+    if not execute(AzCmd("extension", "show").param("--name", "resource-graph"), can_fail=True):
+        execute(AzCmd("extension", "add").param("--name", "resource-graph").param("--yes", ""))
+
+    func_apps_json = execute(AzCmd("graph", "query").param(
+        "-q",
+        f"\"Resources | where type == 'microsoft.web/sites' and kind contains 'functionapp' and name startswith '{CONTROL_PLANE_RESOURCES_TASK_PREFIX}'{subscriptions_clause} | project name, resourceGroup, subscriptionId, location, properties.state\"",
+    ))
+    try:
+        func_apps_response = loads(func_apps_json)
+    except JSONDecodeError as e:
+        log.error(f"Invalid JSON: {func_apps_json}")
+        log.error(f"Error: {e}")
+        raise
+
+    existing_control_planes: dict[str, LfoControlPlane] = {}
+    for func_app in func_apps_response["data"]:
+        subscription_id = func_app["subscriptionId"]
+        existing_control_planes[func_app["name"]] = LfoControlPlane(
+            (subscription_id, sub_id_to_name[subscription_id]),
+            func_app["resourceGroup"],
+            func_app["location"],
+        )
+    return existing_control_planes
 
 
 def check_existing_lfo(
     subscriptions: set[str], sub_id_to_name: dict[str, str]
 ) -> dict[str, LfoMetadata]:
-    """Check if LFO is already installed"""
+    """Check if LFO is already installed on any of the given subscriptions"""
     log.info(
         "Checking if log forwarding is already installed in this Azure environment..."
     )
 
+    control_planes = find_existing_lfo_control_planes(sub_id_to_name, subscriptions)
     existing_lfos: dict[str, LfoMetadata] = {}  # map control plane ID to metadata
 
-    for sub_id in subscriptions:
-        func_apps_json = execute(
-            AzCmd("functionapp", "list")
-            .param("--subscription", sub_id)
-            .param(
-                "--query",
-                f"\"[?starts_with(name,'{CONTROL_PLANE_RESOURCES_TASK_PREFIX}')].{{resourceGroup:resourceGroup, name:name}}\"",
-            )
-            .param("--output", "json")
+    for resource_task_name, control_plane in control_planes.items():
+        control_plane_id = resource_task_name.split("-")[-1]
+
+        resource_task_monitored_sub_ids = execute(
+            AzCmd("functionapp", "config appsettings list")
+            .param("--subscription", control_plane.subscription[0])
+            .param("--name", resource_task_name)
+            .param("--resource-group", control_plane.resource_group)
+            .param("--query", "\"[?name=='MONITORED_SUBSCRIPTIONS'].value\"")
+            .param("--output", "tsv")
         )
 
+        if not resource_task_monitored_sub_ids:
+            continue
+
         try:
-            resources_task_json = loads(func_apps_json)
+            monitored_sub_ids = loads(resource_task_monitored_sub_ids)
         except JSONDecodeError as e:
-            log.error(f"Invalid JSON: {func_apps_json}")
+            log.error(f"Invalid JSON: {resource_task_monitored_sub_ids}")
             log.error(f"Error: {e}")
             raise
 
-        if not resources_task_json:
-            continue
-
-        for resources_task in resources_task_json:
-            name = resources_task["name"]
-            rg = resources_task["resourceGroup"]
-            control_plane_id = name.split("-")[-1]
-
-            resource_task_monitored_sub_ids = execute(
-                AzCmd("functionapp", "config appsettings list")
-                .param("--subscription", sub_id)
-                .param("--name", name)
-                .param("--resource-group", rg)
-                .param("--query", "\"[?name=='MONITORED_SUBSCRIPTIONS'].value\"")
-                .param("--output", "tsv")
-            )
-
-            if not resource_task_monitored_sub_ids:
-                continue
-
-            try:
-                monitored_sub_ids = loads(resource_task_monitored_sub_ids)
-            except JSONDecodeError as e:
-                log.error(f"Invalid JSON: {resource_task_monitored_sub_ids}")
-                log.error(f"Error: {e}")
-                raise
-
-            existing_lfos[control_plane_id] = LfoMetadata(
-                monitored_subs={
-                    sub_id: sub_id_to_name[sub_id] for sub_id in monitored_sub_ids
-                },
-                control_plane_sub=(sub_id, sub_id_to_name[sub_id]),
-                control_plane_rg=rg,
-            )
+        existing_lfos[control_plane_id] = LfoMetadata(
+            control_plane,
+            monitored_subs={
+                sub_id: sub_id_to_name[sub_id] for sub_id in monitored_sub_ids
+            },
+        )
 
     return existing_lfos

azure/logging_install/tests/test_existing_lfo.py

@@ -48,18 +48,25 @@ def patch(self, path: str, **kwargs):
 
     def test_check_existing_lfo_no_installations(self):
         """Test when no LFO installations exist"""
-        self.execute_mock.return_value = "[]"
+        self.execute_mock.return_value = json.dumps({'data':[]})
 
         result = check_existing_lfo(self.config.all_subscriptions, SUB_ID_TO_NAME)
 
         self.assertEqual(result, {})
         self.assertEqual(
-            self.execute_mock.call_count, len(self.config.all_subscriptions)
+            self.execute_mock.call_count, 1
         )
 
     def test_check_existing_lfo_single_installation(self):
         """Test with a single existing LFO installation"""
-        mock_func_apps = [{"resourceGroup": "lfo-rg", "name": "resources-task-abc123"}]
+        mock_func_apps = {
+            "data": [{
+                "resourceGroup": "lfo-rg",
+                "name": "resources-task-abc123",
+                "location": "eastus",
+                "subscriptionId": "sub-1"
+            }]
+        }
         mock_monitored_subs_json = json.dumps(
             {
                 "sub-1": SUB_ID_TO_NAME["sub-1"],
@@ -69,10 +76,8 @@ def test_check_existing_lfo_single_installation(self):
         )
 
         self.execute_mock.side_effect = [
-            json.dumps(mock_func_apps),  # functionapp list for first subscription
+            json.dumps(mock_func_apps),  # functionapp json for first subscription
             mock_monitored_subs_json,  # appsettings for resources-task-abc123 (TSV returns raw JSON string)
-            "[]",  # functionapp list for second subscription (empty)
-            "[]",  # functionapp list for third subscription (empty)
         ]
 
         result = check_existing_lfo(self.config.all_subscriptions, SUB_ID_TO_NAME)
@@ -89,18 +94,26 @@ def test_check_existing_lfo_single_installation(self):
         }
         self.assertEqual(lfo_metadata.monitored_subs, expected_monitored_subs)
         self.assertIn(CONTROL_PLANE_SUBSCRIPTION, self.config.all_subscriptions)
-        self.assertEqual(lfo_metadata.control_plane_rg, "lfo-rg")
+        self.assertEqual(lfo_metadata.control_plane.resource_group, "lfo-rg")
 
     def test_check_existing_lfo_multiple_installations(self):
         """Test with multiple existing LFO installations"""
-        mock_func_apps_sub1 = [
-            {"resourceGroup": "lfo-rg-1", "name": "resources-task-def456"}
-        ]
-
-        mock_func_apps_sub2 = [
-            {"resourceGroup": "lfo-rg-2", "name": "resources-task-ghi789"}
-        ]
-
+        mock_func_apps = {
+            "data": [
+                {
+                    "resourceGroup": "lfo-rg-1",
+                    "name": "resources-task-def456",
+                    "location": "eastus",
+                    "subscriptionId": "sub-1"
+                },
+                {
+                    "resourceGroup": "lfo-rg-2",
+                    "name": "resources-task-ghi789",
+                    "location": "eastus",
+                    "subscriptionId": "sub-2"
+                },
+            ],
+        }
         mock_monitored_subs_1_json = json.dumps(
             {
                 "sub-1": SUB_ID_TO_NAME["sub-1"],
@@ -115,9 +128,8 @@ def test_check_existing_lfo_multiple_installations(self):
         )
 
         self.execute_mock.side_effect = [
-            json.dumps(mock_func_apps_sub1),  # functionapp list for first subscription
+            json.dumps(mock_func_apps),  # functionapp list for first subscription
             mock_monitored_subs_1_json,  # appsettings for resources-task-def456 (TSV returns raw JSON string)
-            json.dumps(mock_func_apps_sub2),  # functionapp list for second subscription
             mock_monitored_subs_2_json,  # appsettings for resources-task-ghi789 (TSV returns raw JSON string)
             "[]",  # functionapp list for third subscription (empty)
         ]
@@ -134,12 +146,12 @@ def test_check_existing_lfo_multiple_installations(self):
             "sub-2": SUB_ID_TO_NAME["sub-2"],
         }
         self.assertEqual(lfo_1.monitored_subs, expected_lfo_1_subs)
-        self.assertEqual(lfo_1.control_plane_rg, "lfo-rg-1")
+        self.assertEqual(lfo_1.control_plane.resource_group, "lfo-rg-1")
 
         lfo_2 = result["ghi789"]
         expected_lfo_2_subs = {
             "sub-3": SUB_ID_TO_NAME["sub-3"],
             "sub-4": SUB_ID_TO_NAME["sub-4"],
         }
         self.assertEqual(lfo_2.monitored_subs, expected_lfo_2_subs)
-        self.assertEqual(lfo_2.control_plane_rg, "lfo-rg-2")
+        self.assertEqual(lfo_2.control_plane.resource_group, "lfo-rg-2")

azure/logging_install/tests/test_validation.py

@@ -6,6 +6,7 @@
 
 # project
 from azure_logging_install import validation
+from azure_logging_install.existing_lfo import LfoControlPlane
 from azure_logging_install.configuration import Configuration
 from azure_logging_install.constants import REQUIRED_RESOURCE_PROVIDERS
 from azure_logging_install.errors import (
@@ -387,15 +388,13 @@ def test_check_fresh_install_with_existing_lfos(self):
                     "sub-1": SUB_ID_TO_NAME["sub-1"],
                     "sub-2": SUB_ID_TO_NAME["sub-2"],
                 },
-                control_plane_sub=CONTROL_PLANE_SUB_ID_TO_NAME,
-                control_plane_rg="existing-rg",
+                control_plane=LfoControlPlane(CONTROL_PLANE_SUB_ID_TO_NAME, "existing-rg", "eastus")
             ),
             "def456": LfoMetadata(
                 monitored_subs={
                     "sub-3": SUB_ID_TO_NAME["sub-3"],
                 },
-                control_plane_sub=CONTROL_PLANE_SUB_ID_TO_NAME,
-                control_plane_rg="another-rg",
+                control_plane=LfoControlPlane(CONTROL_PLANE_SUB_ID_TO_NAME, "another-rg", "westus")
             ),
         }