From 3abd1acef13f2820e77471cf0cacb881141b2e8a Mon Sep 17 00:00:00 2001
From: longwan <longwan@microsoft.com>
Date: Tue, 15 Jul 2025 07:05:35 +0000
Subject: [PATCH 1/7] Extend ContainerInsights Extension for high log scale
 mode support

---
 src/k8s-extension/HISTORY.rst                 |   4 +
 .../partner_extensions/ContainerInsights.py   | 144 +++++++++++++++++-
 src/k8s-extension/setup.py                    |   2 +-
 testing/pipeline/k8s-custom-pipelines.yml     |   4 +
 .../public/AzureMonitorHighScale.Tests.ps1    | 126 +++++++++++++++
 5 files changed, 277 insertions(+), 3 deletions(-)
 create mode 100644 testing/test/extensions/public/AzureMonitorHighScale.Tests.ps1

diff --git a/src/k8s-extension/HISTORY.rst b/src/k8s-extension/HISTORY.rst
index 09a13348d04..94df11be9ce 100644
--- a/src/k8s-extension/HISTORY.rst
+++ b/src/k8s-extension/HISTORY.rst
@@ -3,6 +3,10 @@
 Release History
 ===============
 
+1.6.6
+++++++++
+* microsoft.azuremonitor.containers: Extend ContainerInsights Extension for high log scale mode support.
+
 1.6.5
 ++++++++++++++++++
 * microsoft.dataprotection.kubernetes: Add support for 'DisableInformerCache' configuration.
diff --git a/src/k8s-extension/azext_k8s_extension/partner_extensions/ContainerInsights.py b/src/k8s-extension/azext_k8s_extension/partner_extensions/ContainerInsights.py
index b059850d7b4..5cf74793808 100644
--- a/src/k8s-extension/azext_k8s_extension/partner_extensions/ContainerInsights.py
+++ b/src/k8s-extension/azext_k8s_extension/partner_extensions/ContainerInsights.py
@@ -33,6 +33,21 @@
 logger = get_logger(__name__)
 DCR_API_VERSION = "2022-06-01"
 
+ContainerInsightsStreams = [
+    "Microsoft-ContainerLog",
+    "Microsoft-ContainerLogV2-HighScale",
+    "Microsoft-KubeEvents",
+    "Microsoft-KubePodInventory",
+    "Microsoft-KubeNodeInventory",
+    "Microsoft-KubePVInventory",
+    "Microsoft-KubeServices",
+    "Microsoft-KubeMonAgentEvents",
+    "Microsoft-InsightsMetrics",
+    "Microsoft-ContainerInventory",
+    "Microsoft-ContainerNodeInventory",
+    "Microsoft-Perf",
+]
+
 
 class ContainerInsights(DefaultExtension):
     def Create(self, cmd, client, resource_group_name, cluster_name, name, cluster_type, cluster_rp,
@@ -83,6 +98,7 @@ def Delete(self, cmd, client, resource_group_name, cluster_name, name, cluster_t
         # Delete DCR-A if it exists incase of MSI Auth
         useAADAuth = False
         isDCRAExists = False
+        enable_high_log_scale_mode = False
         cluster_rp, _ = get_cluster_rp_api_version(cluster_type=cluster_type, cluster_rp=cluster_rp)
         try:
             extension = client.get(resource_group_name, cluster_rp, cluster_type, cluster_name, name)
@@ -95,6 +111,7 @@ def Delete(self, cmd, client, resource_group_name, cluster_name, name, cluster_t
                 return
 
         subscription_id = get_subscription_id(cmd.cli_ctx)
+        resources = cf_resources(cmd.cli_ctx, subscription_id)
         # handle cluster type here
         cluster_resource_id = '/subscriptions/{0}/resourceGroups/{1}/providers/{2}/{3}/{4}'.format(subscription_id, resource_group_name, cluster_rp, cluster_type, cluster_name)
         if (extension is not None) and (extension.configuration_settings is not None):
@@ -108,6 +125,14 @@ def Delete(self, cmd, client, resource_group_name, cluster_name, name, cluster_t
                 useAADAuthSetting = configSettings['amalogs.useAADAuth']
                 if (isinstance(useAADAuthSetting, str) and str(useAADAuthSetting).lower() == "true") or (isinstance(useAADAuthSetting, bool) and useAADAuthSetting):
                     useAADAuth = True
+            
+            # Check if high log scale mode was enabled
+            if useAADAuth and 'amalogs.enableHighLogScaleMode' in configSettings:
+                highLogScaleSetting = configSettings['amalogs.enableHighLogScaleMode']
+                if isinstance(highLogScaleSetting, str):
+                    enable_high_log_scale_mode = (highLogScaleSetting.lower() == "true")
+                elif isinstance(highLogScaleSetting, bool):
+                    enable_high_log_scale_mode = highLogScaleSetting
         if useAADAuth:
             association_url = cmd.cli_ctx.cloud.endpoints.resource_manager + f"{cluster_resource_id}/providers/Microsoft.Insights/dataCollectionRuleAssociations/ContainerInsightsExtension?api-version={DCR_API_VERSION}"
             for _ in range(3):
@@ -131,6 +156,28 @@ def Delete(self, cmd, client, resource_group_name, cluster_name, name, cluster_t
                 except Exception:
                     pass  # its OK to ignore the exception since MSI auth in preview
 
+        if useAADAuth:
+            resource = resources.get_by_id(cluster_resource_id, '2020-01-01-preview')
+            cluster_location = resource.location.lower()
+            dcr_name = f"MSCI-{cluster_location}-{cluster_name}"
+            dcr_name = dcr_name[0:64]
+            
+            dcr_resource_id = f"/subscriptions/{subscription_id}/resourceGroups/{resource_group_name}/providers/Microsoft.Insights/dataCollectionRules/{dcr_name}"
+            dcr_url = cmd.cli_ctx.cloud.endpoints.resource_manager + f"{dcr_resource_id}?api-version={DCR_API_VERSION}"
+            response = send_raw_request(cmd.cli_ctx, "GET", dcr_url)
+            dcr_config = json.loads(response.text)
+            # Delete the DCR
+            for _ in range(3):
+                try:
+                    send_raw_request(cmd.cli_ctx, "DELETE", dcr_url,)
+                    logger.info(f"Successfully deleted DCR: {dcr_name}")
+                    break
+                except Exception as ex:
+                    logger.warning(f"Error deleting DCR: {str(ex)}")
+                    pass
+
+            if enable_high_log_scale_mode:
+                _delete_dce_for_dcr(cmd, subscription_id, resource_group_name, dcr_config)
 
 # Custom Validation Logic for Container Insights
 
@@ -464,6 +511,7 @@ def _get_container_insights_settings(cmd, cluster_resource_group_name, cluster_r
     subscription_id = get_subscription_id(cmd.cli_ctx)
     workspace_resource_id = ''
     useAADAuth = True
+    enableHighLogScaleMode = False  # Default value
     if 'amalogs.useAADAuth' not in configuration_settings:
         configuration_settings['amalogs.useAADAuth'] = "true"
     extensionSettings = {}
@@ -520,6 +568,16 @@ def _get_container_insights_settings(cmd, cluster_resource_group_name, cluster_r
                     raise InvalidArgumentValueError('streams must be an array type')
             extensionSettings["dataCollectionSettings"] = dataCollectionSettings
 
+        if useAADAuth and 'amalogs.enableHighLogScaleMode' in configuration_settings:
+            enableHighLogScaleMode = configuration_settings['amalogs.enableHighLogScaleMode']
+            if isinstance(enableHighLogScaleMode, str):
+                enableHighLogScaleMode_str = enableHighLogScaleMode.lower()
+                if enableHighLogScaleMode_str not in ["true", "false"]:
+                    raise InvalidArgumentValueError('amalogs.enableHighLogScaleMode value MUST be either true or false')
+                enableHighLogScaleMode = (enableHighLogScaleMode_str == "true")
+            elif not isinstance(enableHighLogScaleMode, bool):
+                raise InvalidArgumentValueError('amalogs.enableHighLogScaleMode value MUST be either true or false')
+
     workspace_resource_id = workspace_resource_id.strip()
 
     if configuration_protected_settings is not None:
@@ -548,7 +606,7 @@ def _get_container_insights_settings(cmd, cluster_resource_group_name, cluster_r
     if is_ci_extension_type:
         if useAADAuth:
             logger.info("creating data collection rule and association")
-            _ensure_container_insights_dcr_for_monitoring(cmd, subscription_id, cluster_resource_group_name, cluster_rp, cluster_type, cluster_name, workspace_resource_id, extensionSettings)
+            _ensure_container_insights_dcr_for_monitoring(cmd, subscription_id, cluster_resource_group_name, cluster_rp, cluster_type, cluster_name, workspace_resource_id, extensionSettings, enableHighLogScaleMode)
         elif not _is_container_insights_solution_exists(cmd, workspace_resource_id):
             logger.info("Creating ContainerInsights solution resource, since it doesn't exist and it is using legacy authentication")
             _ensure_container_insights_for_monitoring(cmd, workspace_resource_id).result()
@@ -597,6 +655,37 @@ def _get_container_insights_settings(cmd, cluster_resource_group_name, cluster_r
         configuration_settings['amalogs.domain'] = 'opinsights.azure.microsoft.scloud'
 
 
+def _delete_dce_for_dcr(cmd, subscription_id, cluster_resource_group_name, dcr_config):
+    """Delete Data Collection Endpoint associated with a DCR if it exists"""
+    try:
+        if ("properties" in dcr_config and 
+            "dataCollectionEndpointId" in dcr_config["properties"] and 
+            dcr_config["properties"]["dataCollectionEndpointId"]):
+            
+            dce_id = dcr_config["properties"]["dataCollectionEndpointId"]
+            dce_parts = dce_id.split('/')
+
+            if len(dce_parts) > 0:
+                dce_name = dce_parts[-1]
+                dce_resource_id = f"/subscriptions/{subscription_id}/resourceGroups/{cluster_resource_group_name}/providers/Microsoft.Insights/dataCollectionEndpoints/{dce_name}"
+                dce_url = cmd.cli_ctx.cloud.endpoints.resource_manager + f"{dce_resource_id}?api-version=2022-06-01"              
+                # Try to delete up to 3 times
+                for retry in range(3):
+                    try:
+                        send_raw_request(cmd.cli_ctx, "DELETE", dce_url)
+                        logger.info("Successfully deleted DCE: %s", dce_name)
+                        return True
+                    except CLIError as e:
+                        if "ResourceNotFound" in str(e):
+                            return True
+                        if retry == 2:
+                            logger.warning("Failed to delete DCE: %s - %s", dce_name, str(e))
+                            return False
+                        logger.info("Retrying DCE deletion after error: %s", str(e))
+    except CLIError:
+        pass
+    return True
+
 def get_existing_container_insights_extension_dcr_tags(cmd, dcr_url):
     tags = {}
     _MAX_RETRY_TIMES = 3
@@ -617,7 +706,7 @@ def get_existing_container_insights_extension_dcr_tags(cmd, dcr_url):
     return tags
 
 
-def _ensure_container_insights_dcr_for_monitoring(cmd, subscription_id, cluster_resource_group_name, cluster_rp, cluster_type, cluster_name, workspace_resource_id, extensionSettings):
+def _ensure_container_insights_dcr_for_monitoring(cmd, subscription_id, cluster_resource_group_name, cluster_rp, cluster_type, cluster_name, workspace_resource_id, extensionSettings, enable_high_log_scale_mode):
     from azure.core.exceptions import HttpResponseError
 
     cluster_region = ''
@@ -652,6 +741,18 @@ def _ensure_container_insights_dcr_for_monitoring(cmd, subscription_id, cluster_
     dataCollectionRuleName = dataCollectionRuleName[0:64]
     dcr_resource_id = f"/subscriptions/{subscription_id}/resourceGroups/{cluster_resource_group_name}/providers/Microsoft.Insights/dataCollectionRules/{dataCollectionRuleName}"
 
+    # ingestion DCE MUST be in workspace region
+    ingestionDataCollectionEndpointName = f"MSCI-ingest-{workspace_region}-{cluster_name}"
+    # Max length of the DCE name is 43 chars
+    ingestionDataCollectionEndpointName = _trim_suffix_if_needed(ingestionDataCollectionEndpointName[0:43])
+    ingestion_dce_resource_id = None
+
+    # create ingestion DCE if high log scale mode enabled
+    if enable_high_log_scale_mode:
+        ingestion_dce_resource_id = create_data_collection_endpoint(
+            cmd, subscription_id, cluster_resource_group_name, workspace_region, ingestionDataCollectionEndpointName
+        )
+
     # first get the association between region display names and region IDs (because for some reason
     # the "which RPs are available in which regions" check returns region display names)
     region_names_to_id = {}
@@ -677,6 +778,8 @@ def _ensure_container_insights_dcr_for_monitoring(cmd, subscription_id, cluster_
     # get existing tags on the container insights extension DCR if the customer added any
     existing_tags = get_existing_container_insights_extension_dcr_tags(cmd, dcr_url)
     streams = ["Microsoft-ContainerInsights-Group-Default"]
+    if enable_high_log_scale_mode:
+        streams = ContainerInsightsStreams
     if extensionSettings is None:
         extensionSettings = {}
     if 'dataCollectionSettings' in extensionSettings.keys():
@@ -691,6 +794,11 @@ def _ensure_container_insights_dcr_for_monitoring(cmd, subscription_id, cluster_
         }
     extensionSettings["dataCollectionSettings"] = dataCollectionSettings
 
+    if enable_high_log_scale_mode:
+        for i, v in enumerate(streams):
+            if v == "Microsoft-ContainerLogV2":
+                streams[i] = "Microsoft-ContainerLogV2-HighScale"
+
     # create the DCR
     dcr_creation_body = json.dumps(
         {
@@ -722,6 +830,7 @@ def _ensure_container_insights_dcr_for_monitoring(cmd, subscription_id, cluster_
                         }
                     ]
                 },
+                "dataCollectionEndpointId": ingestion_dce_resource_id
             },
         }
     )
@@ -755,3 +864,34 @@ def _ensure_container_insights_dcr_for_monitoring(cmd, subscription_id, cluster_
             error = e
         else:
             raise error
+
+
+def create_data_collection_endpoint(cmd, subscription_id, cluster_resource_group_name, workspace_region, ingestionDataCollectionEndpointName):
+    # create the ingestion DCE
+    ingestion_dce_resource_id = f"/subscriptions/{subscription_id}/resourceGroups/{cluster_resource_group_name}/providers/Microsoft.Insights/dataCollectionEndpoints/{ingestionDataCollectionEndpointName}"
+    ingestion_dce_url = cmd.cli_ctx.cloud.endpoints.resource_manager + f"{ingestion_dce_resource_id}?api-version=2022-06-01"
+    ingestion_dce_creation_body = json.dumps({
+        "location": workspace_region,
+        "kind": "Linux",
+        "properties": {
+            "networkAcls": {
+                "publicNetworkAccess": "Enabled"
+            }
+        }
+    })
+    error = None
+    for _ in range(3):
+        try:
+            send_raw_request(cmd.cli_ctx, "PUT", ingestion_dce_url, body=ingestion_dce_creation_body)
+            return ingestion_dce_resource_id
+        except AzCLIError as e:
+            error = e
+    if error:
+        raise error
+    return ingestion_dce_resource_id
+
+
+def _trim_suffix_if_needed(s, suffix="-"):
+    if s.endswith(suffix):
+        s = s[:-len(suffix)]
+    return s
diff --git a/src/k8s-extension/setup.py b/src/k8s-extension/setup.py
index c46e1aca890..bdee2be8241 100644
--- a/src/k8s-extension/setup.py
+++ b/src/k8s-extension/setup.py
@@ -33,7 +33,7 @@
 # TODO: Add any additional SDK dependencies here
 DEPENDENCIES = []
 
-VERSION = "1.6.5"
+VERSION = "1.6.6"
 
 with open("README.rst", "r", encoding="utf-8") as f:
     README = f.read()
diff --git a/testing/pipeline/k8s-custom-pipelines.yml b/testing/pipeline/k8s-custom-pipelines.yml
index 0cc9895583a..8fe39c8728d 100644
--- a/testing/pipeline/k8s-custom-pipelines.yml
+++ b/testing/pipeline/k8s-custom-pipelines.yml
@@ -35,6 +35,10 @@ stages:
     parameters:
       jobName: AzureMonitor
       path: ./test/extensions/public/AzureMonitor.Tests.ps1
+  - template: ./templates/run-test.yml
+    parameters:
+      jobName: AzureMonitorHighScale
+      path: ./test/extensions/public/AzureMonitorHighScale.Tests.ps1
   - template: ./templates/run-test.yml
     parameters:
       jobName: AzurePolicy
diff --git a/testing/test/extensions/public/AzureMonitorHighScale.Tests.ps1 b/testing/test/extensions/public/AzureMonitorHighScale.Tests.ps1
new file mode 100644
index 00000000000..183bfb6b476
--- /dev/null
+++ b/testing/test/extensions/public/AzureMonitorHighScale.Tests.ps1
@@ -0,0 +1,126 @@
+Describe 'Azure Monitor High Scale Mode Testing' {
+    BeforeAll {
+        $extensionType = "microsoft.azuremonitor.containers"
+        $extensionName = "azuremonitor-containers"
+        $extensionAgentName = "omsagent"
+        $extensionAgentNamespace = "kube-system"
+        
+        . $PSScriptRoot/../../helper/Constants.ps1
+        . $PSScriptRoot/../../helper/Helper.ps1
+    }
+
+    It 'Creates the extension with high log scale mode and verifies DCE creation' {
+        # Create extension with high scale mode enabled
+        az $Env:K8sExtensionName create -c $($ENVCONFIG.arcClusterName) -g $($ENVCONFIG.resourceGroup) `
+            --cluster-type connectedClusters --extension-type $extensionType -n $extensionName `
+            --configuration-settings "amalogs.enableHighLogScaleMode=true" --no-wait
+        $? | Should -BeTrue
+
+        # Verify extension creation
+        $output = az $Env:K8sExtensionName show -c $($ENVCONFIG.arcClusterName) -g $($ENVCONFIG.resourceGroup) --cluster-type connectedClusters -n $extensionName
+        $? | Should -BeTrue
+
+        $extension = ($output | ConvertFrom-Json)
+        $extension | Should -Not -BeNullOrEmpty
+        
+        # Verify high scale mode configuration
+        $settings = $extension.configurationSettings
+        $settings.'amalogs.enableHighLogScaleMode' | Should -Be "true"
+        $settings.'amalogs.useAADAuth' | Should -Be "true"
+
+        # Wait for extension to install
+        $n = 0
+        do {
+            if (Has-ExtensionData $extensionName) {
+                break
+            }
+            Start-Sleep -Seconds 10
+            $n += 1
+        } while ($n -le $MAX_RETRY_ATTEMPTS)
+        $n | Should -BeLessOrEqual $MAX_RETRY_ATTEMPTS
+
+        # Verify DCE creation
+        $clusterLocation = (az resource show -g $($ENVCONFIG.resourceGroup) -n $($ENVCONFIG.arcClusterName) --resource-type "Microsoft.Kubernetes/connectedClusters" --query location -o tsv).ToLower()
+        $dceName = "MSCI-ingest-$clusterLocation-$($ENVCONFIG.arcClusterName)"
+        if ($dceName.Length -gt 43) {
+            $dceName = $dceName.Substring(0, 43)
+            # Remove trailing hyphen if present
+            if ($dceName.EndsWith("-")) {
+                $dceName = $dceName.Substring(0, $dceName.Length - 1)
+            }
+        }
+        
+        $dce = az monitor data-collection endpoint show -g $($ENVCONFIG.resourceGroup) -n $dceName
+        $? | Should -BeTrue
+        $dce | Should -Not -BeNullOrEmpty
+        
+        # Verify DCE configuration
+        $dceObj = ($dce | ConvertFrom-Json)
+        $dceObj.kind | Should -Be "Linux"
+        $dceObj.properties.networkAcls.publicNetworkAccess | Should -Be "Enabled"
+    }
+
+    It "Verifies Data Collection Rule configuration for high scale mode" {
+        # Get the DCR name
+        $clusterLocation = (az resource show -g $($ENVCONFIG.resourceGroup) -n $($ENVCONFIG.arcClusterName) --resource-type "Microsoft.Kubernetes/connectedClusters" --query location -o tsv).ToLower()
+        $dcrName = "MSCI-$clusterLocation-$($ENVCONFIG.arcClusterName)"
+        if ($dcrName.Length -gt 64) {
+            $dcrName = $dcrName.Substring(0, 64)
+        }
+
+        # Get the DCR
+        $dcr = az monitor data-collection rule show -g $($ENVCONFIG.resourceGroup) -n $dcrName
+        $? | Should -BeTrue
+        $dcr | Should -Not -BeNullOrEmpty
+
+        # Verify high scale mode streams configuration
+        $dcrObj = ($dcr | ConvertFrom-Json)
+        $streams = $dcrObj.properties.dataSources.extensions[0].streams
+        $streams | Should -Contain "Microsoft-ContainerLogV2-HighScale"
+        $streams | Should -Contain "Microsoft-KubeEvents"
+        $streams | Should -Contain "Microsoft-KubePodInventory"
+        $streams | Should -Contain "Microsoft-KubeNodeInventory"
+        $streams | Should -Contain "Microsoft-KubePVInventory"
+        $streams | Should -Contain "Microsoft-KubeServices"
+        $streams | Should -Contain "Microsoft-KubeMonAgentEvents"
+        $streams | Should -Contain "Microsoft-InsightsMetrics"
+        $streams | Should -Contain "Microsoft-ContainerInventory"
+        $streams | Should -Contain "Microsoft-ContainerNodeInventory"
+        $streams | Should -Contain "Microsoft-Perf"
+    }
+
+    It "Deletes the extension and verifies DCE cleanup" {
+        # Get DCE name before deletion
+        $clusterLocation = (az resource show -g $($ENVCONFIG.resourceGroup) -n $($ENVCONFIG.arcClusterName) --resource-type "Microsoft.Kubernetes/connectedClusters" --query location -o tsv).ToLower()
+        $dceName = "MSCI-ingest-$clusterLocation-$($ENVCONFIG.arcClusterName)"
+        if ($dceName.Length -gt 43) {
+            $dceName = $dceName.Substring(0, 43)
+            # Remove trailing hyphen if present
+            if ($dceName.EndsWith("-")) {
+                $dceName = $dceName.Substring(0, $dceName.Length - 1)
+            }
+        }
+
+        # Delete the extension
+        $output = az $Env:K8sExtensionName delete -c $($ENVCONFIG.arcClusterName) -g $($ENVCONFIG.resourceGroup) --cluster-type connectedClusters -n $extensionName --force
+        $? | Should -BeTrue
+
+        # Verify extension is deleted
+        $output = az $Env:K8sExtensionName show -c $($ENVCONFIG.arcClusterName) -g $($ENVCONFIG.resourceGroup) --cluster-type connectedClusters -n $extensionName
+        $? | Should -BeFalse
+        $output | Should -BeNullOrEmpty
+
+        # Verify DCE is deleted
+        $dce = az monitor data-collection endpoint show -g $($ENVCONFIG.resourceGroup) -n $dceName
+        $? | Should -BeFalse
+    }
+
+    It "Performs another list after the delete" {
+        $output = az $Env:K8sExtensionName list -c $($ENVCONFIG.arcClusterName) -g $($ENVCONFIG.resourceGroup) --cluster-type connectedClusters
+        $? | Should -BeTrue
+        $output | Should -Not -BeNullOrEmpty
+        
+        $extensionExists = $output | ConvertFrom-Json | Where-Object { $_.extensionType -eq $extensionName }
+        $extensionExists | Should -BeNullOrEmpty
+    }
+}

From a7efafd3c2f87c2ceb16096d405590c52eaa1281 Mon Sep 17 00:00:00 2001
From: longwan <longwan@microsoft.com>
Date: Wed, 16 Jul 2025 22:56:46 +0000
Subject: [PATCH 2/7] update test with retry

---
 .../public/AzureMonitorHighScale.Tests.ps1    | 50 +++++++++++++++----
 1 file changed, 39 insertions(+), 11 deletions(-)

diff --git a/testing/test/extensions/public/AzureMonitorHighScale.Tests.ps1 b/testing/test/extensions/public/AzureMonitorHighScale.Tests.ps1
index 183bfb6b476..79aae55b21a 100644
--- a/testing/test/extensions/public/AzureMonitorHighScale.Tests.ps1
+++ b/testing/test/extensions/public/AzureMonitorHighScale.Tests.ps1
@@ -13,7 +13,7 @@ Describe 'Azure Monitor High Scale Mode Testing' {
         # Create extension with high scale mode enabled
         az $Env:K8sExtensionName create -c $($ENVCONFIG.arcClusterName) -g $($ENVCONFIG.resourceGroup) `
             --cluster-type connectedClusters --extension-type $extensionType -n $extensionName `
-            --configuration-settings "amalogs.enableHighLogScaleMode=true" --no-wait
+            --configuration-settings "amalogs.useAADAuth=true" "amalogs.enableHighLogScaleMode=true" --no-wait
         $? | Should -BeTrue
 
         # Verify extension creation
@@ -50,12 +50,24 @@ Describe 'Azure Monitor High Scale Mode Testing' {
             }
         }
         
-        $dce = az monitor data-collection endpoint show -g $($ENVCONFIG.resourceGroup) -n $dceName
-        $? | Should -BeTrue
-        $dce | Should -Not -BeNullOrEmpty
+        # Wait for DCE to be fully created and configured
+        $maxRetries = 30  # 5 minutes total
+        $retryCount = 0
+        $dceObj = $null
         
-        # Verify DCE configuration
-        $dceObj = ($dce | ConvertFrom-Json)
+        while ($retryCount -lt $maxRetries) {
+            $dce = az monitor data-collection endpoint show -g $($ENVCONFIG.resourceGroup) -n $dceName 2>$null
+            if ($?) {
+                $dceObj = ($dce | ConvertFrom-Json)
+                if ($dceObj.properties.networkAcls.publicNetworkAccess -eq "Enabled") {
+                    break
+                }
+            }
+            Start-Sleep -Seconds 10
+            $retryCount++
+        }
+
+        $dceObj | Should -Not -BeNullOrEmpty
         $dceObj.kind | Should -Be "Linux"
         $dceObj.properties.networkAcls.publicNetworkAccess | Should -Be "Enabled"
     }
@@ -68,13 +80,29 @@ Describe 'Azure Monitor High Scale Mode Testing' {
             $dcrName = $dcrName.Substring(0, 64)
         }
 
-        # Get the DCR
-        $dcr = az monitor data-collection rule show -g $($ENVCONFIG.resourceGroup) -n $dcrName
-        $? | Should -BeTrue
-        $dcr | Should -Not -BeNullOrEmpty
+        # Wait for DCR to be fully created and configured
+        $maxRetries = 30  # 5 minutes total
+        $retryCount = 0
+        $dcrObj = $null
+        
+        while ($retryCount -lt $maxRetries) {
+            $dcr = az monitor data-collection rule show -g $($ENVCONFIG.resourceGroup) -n $dcrName 2>$null
+            if ($?) {
+                $dcrObj = ($dcr | ConvertFrom-Json)
+                if ($dcrObj.properties.dataSources.extensions -and 
+                    $dcrObj.properties.dataSources.extensions[0].streams) {
+                    break
+                }
+            }
+            Start-Sleep -Seconds 10
+            $retryCount++
+        }
+
+        $dcrObj | Should -Not -BeNullOrEmpty
+        $dcrObj.properties.dataSources.extensions | Should -Not -BeNullOrEmpty
+        $dcrObj.properties.dataSources.extensions[0].streams | Should -Not -BeNullOrEmpty
 
         # Verify high scale mode streams configuration
-        $dcrObj = ($dcr | ConvertFrom-Json)
         $streams = $dcrObj.properties.dataSources.extensions[0].streams
         $streams | Should -Contain "Microsoft-ContainerLogV2-HighScale"
         $streams | Should -Contain "Microsoft-KubeEvents"

From c26d0d02ef72f17ec0b837bf956912acca285a01 Mon Sep 17 00:00:00 2001
From: longwan <longwan@microsoft.com>
Date: Wed, 16 Jul 2025 23:53:49 +0000
Subject: [PATCH 3/7] update test with retry

---
 .../public/AzureMonitorHighScale.Tests.ps1    | 122 ++++--------------
 1 file changed, 22 insertions(+), 100 deletions(-)

diff --git a/testing/test/extensions/public/AzureMonitorHighScale.Tests.ps1 b/testing/test/extensions/public/AzureMonitorHighScale.Tests.ps1
index 79aae55b21a..10ce9fb6f0b 100644
--- a/testing/test/extensions/public/AzureMonitorHighScale.Tests.ps1
+++ b/testing/test/extensions/public/AzureMonitorHighScale.Tests.ps1
@@ -9,28 +9,28 @@ Describe 'Azure Monitor High Scale Mode Testing' {
         . $PSScriptRoot/../../helper/Helper.ps1
     }
 
-    It 'Creates the extension with high log scale mode and verifies DCE creation' {
-        # Create extension with high scale mode enabled
+    It 'Creates the extension with high log scale mode and checks that it onboards correctly' {
         az $Env:K8sExtensionName create -c $($ENVCONFIG.arcClusterName) -g $($ENVCONFIG.resourceGroup) `
             --cluster-type connectedClusters --extension-type $extensionType -n $extensionName `
-            --configuration-settings "amalogs.useAADAuth=true" "amalogs.enableHighLogScaleMode=true" --no-wait
-        $? | Should -BeTrue
+            --configuration-settings "amalogs.enableHighLogScaleMode=true" --no-wait
+        $? | Should -BeTrue        
 
-        # Verify extension creation
         $output = az $Env:K8sExtensionName show -c $($ENVCONFIG.arcClusterName) -g $($ENVCONFIG.resourceGroup) --cluster-type connectedClusters -n $extensionName
         $? | Should -BeTrue
 
         $extension = ($output | ConvertFrom-Json)
-        $extension | Should -Not -BeNullOrEmpty
-        
+        $isAutoUpgradeMinorVersion = $extension.autoUpgradeMinorVersion 
+        $isAutoUpgradeMinorVersion.ToString() -eq "True" | Should -BeTrue
+
         # Verify high scale mode configuration
         $settings = $extension.configurationSettings
         $settings.'amalogs.enableHighLogScaleMode' | Should -Be "true"
         $settings.'amalogs.useAADAuth' | Should -Be "true"
 
-        # Wait for extension to install
+        # Loop and retry until the extension installs
         $n = 0
-        do {
+        do 
+        {
             if (Has-ExtensionData $extensionName) {
                 break
             }
@@ -38,109 +38,31 @@ Describe 'Azure Monitor High Scale Mode Testing' {
             $n += 1
         } while ($n -le $MAX_RETRY_ATTEMPTS)
         $n | Should -BeLessOrEqual $MAX_RETRY_ATTEMPTS
-
-        # Verify DCE creation
-        $clusterLocation = (az resource show -g $($ENVCONFIG.resourceGroup) -n $($ENVCONFIG.arcClusterName) --resource-type "Microsoft.Kubernetes/connectedClusters" --query location -o tsv).ToLower()
-        $dceName = "MSCI-ingest-$clusterLocation-$($ENVCONFIG.arcClusterName)"
-        if ($dceName.Length -gt 43) {
-            $dceName = $dceName.Substring(0, 43)
-            # Remove trailing hyphen if present
-            if ($dceName.EndsWith("-")) {
-                $dceName = $dceName.Substring(0, $dceName.Length - 1)
-            }
-        }
-        
-        # Wait for DCE to be fully created and configured
-        $maxRetries = 30  # 5 minutes total
-        $retryCount = 0
-        $dceObj = $null
-        
-        while ($retryCount -lt $maxRetries) {
-            $dce = az monitor data-collection endpoint show -g $($ENVCONFIG.resourceGroup) -n $dceName 2>$null
-            if ($?) {
-                $dceObj = ($dce | ConvertFrom-Json)
-                if ($dceObj.properties.networkAcls.publicNetworkAccess -eq "Enabled") {
-                    break
-                }
-            }
-            Start-Sleep -Seconds 10
-            $retryCount++
-        }
-
-        $dceObj | Should -Not -BeNullOrEmpty
-        $dceObj.kind | Should -Be "Linux"
-        $dceObj.properties.networkAcls.publicNetworkAccess | Should -Be "Enabled"
     }
 
-    It "Verifies Data Collection Rule configuration for high scale mode" {
-        # Get the DCR name
-        $clusterLocation = (az resource show -g $($ENVCONFIG.resourceGroup) -n $($ENVCONFIG.arcClusterName) --resource-type "Microsoft.Kubernetes/connectedClusters" --query location -o tsv).ToLower()
-        $dcrName = "MSCI-$clusterLocation-$($ENVCONFIG.arcClusterName)"
-        if ($dcrName.Length -gt 64) {
-            $dcrName = $dcrName.Substring(0, 64)
-        }
-
-        # Wait for DCR to be fully created and configured
-        $maxRetries = 30  # 5 minutes total
-        $retryCount = 0
-        $dcrObj = $null
-        
-        while ($retryCount -lt $maxRetries) {
-            $dcr = az monitor data-collection rule show -g $($ENVCONFIG.resourceGroup) -n $dcrName 2>$null
-            if ($?) {
-                $dcrObj = ($dcr | ConvertFrom-Json)
-                if ($dcrObj.properties.dataSources.extensions -and 
-                    $dcrObj.properties.dataSources.extensions[0].streams) {
-                    break
-                }
-            }
-            Start-Sleep -Seconds 10
-            $retryCount++
-        }
+    It "Performs a show on the extension" {
+        $output = az $Env:K8sExtensionName show -c $($ENVCONFIG.arcClusterName) -g $($ENVCONFIG.resourceGroup) --cluster-type connectedClusters -n $extensionName
+        $? | Should -BeTrue
+        $output | Should -Not -BeNullOrEmpty
+    }
 
-        $dcrObj | Should -Not -BeNullOrEmpty
-        $dcrObj.properties.dataSources.extensions | Should -Not -BeNullOrEmpty
-        $dcrObj.properties.dataSources.extensions[0].streams | Should -Not -BeNullOrEmpty
+    It "Lists the extensions on the cluster" {
+        $output = az $Env:K8sExtensionName list -c $($ENVCONFIG.arcClusterName) -g $($ENVCONFIG.resourceGroup) --cluster-type connectedClusters
+        $? | Should -BeTrue
 
-        # Verify high scale mode streams configuration
-        $streams = $dcrObj.properties.dataSources.extensions[0].streams
-        $streams | Should -Contain "Microsoft-ContainerLogV2-HighScale"
-        $streams | Should -Contain "Microsoft-KubeEvents"
-        $streams | Should -Contain "Microsoft-KubePodInventory"
-        $streams | Should -Contain "Microsoft-KubeNodeInventory"
-        $streams | Should -Contain "Microsoft-KubePVInventory"
-        $streams | Should -Contain "Microsoft-KubeServices"
-        $streams | Should -Contain "Microsoft-KubeMonAgentEvents"
-        $streams | Should -Contain "Microsoft-InsightsMetrics"
-        $streams | Should -Contain "Microsoft-ContainerInventory"
-        $streams | Should -Contain "Microsoft-ContainerNodeInventory"
-        $streams | Should -Contain "Microsoft-Perf"
+        $output | Should -Not -BeNullOrEmpty
+        $extensionExists = $output | ConvertFrom-Json | Where-Object { $_.extensionType -eq $extensionType }
+        $extensionExists | Should -Not -BeNullOrEmpty
     }
 
-    It "Deletes the extension and verifies DCE cleanup" {
-        # Get DCE name before deletion
-        $clusterLocation = (az resource show -g $($ENVCONFIG.resourceGroup) -n $($ENVCONFIG.arcClusterName) --resource-type "Microsoft.Kubernetes/connectedClusters" --query location -o tsv).ToLower()
-        $dceName = "MSCI-ingest-$clusterLocation-$($ENVCONFIG.arcClusterName)"
-        if ($dceName.Length -gt 43) {
-            $dceName = $dceName.Substring(0, 43)
-            # Remove trailing hyphen if present
-            if ($dceName.EndsWith("-")) {
-                $dceName = $dceName.Substring(0, $dceName.Length - 1)
-            }
-        }
-
-        # Delete the extension
+    It "Deletes the extension from the cluster" {
         $output = az $Env:K8sExtensionName delete -c $($ENVCONFIG.arcClusterName) -g $($ENVCONFIG.resourceGroup) --cluster-type connectedClusters -n $extensionName --force
         $? | Should -BeTrue
 
-        # Verify extension is deleted
+        # Extension should not be found on the cluster
         $output = az $Env:K8sExtensionName show -c $($ENVCONFIG.arcClusterName) -g $($ENVCONFIG.resourceGroup) --cluster-type connectedClusters -n $extensionName
         $? | Should -BeFalse
         $output | Should -BeNullOrEmpty
-
-        # Verify DCE is deleted
-        $dce = az monitor data-collection endpoint show -g $($ENVCONFIG.resourceGroup) -n $dceName
-        $? | Should -BeFalse
     }
 
     It "Performs another list after the delete" {

From 93e986b4a27cd579997fe0b3c35d8d4b24701979 Mon Sep 17 00:00:00 2001
From: longwan <longwan@microsoft.com>
Date: Wed, 16 Jul 2025 23:58:52 +0000
Subject: [PATCH 4/7] clean dcr and dce when delete

---
 testing/test/extensions/public/AzureMonitorHighScale.Tests.ps1 | 1 -
 1 file changed, 1 deletion(-)

diff --git a/testing/test/extensions/public/AzureMonitorHighScale.Tests.ps1 b/testing/test/extensions/public/AzureMonitorHighScale.Tests.ps1
index 10ce9fb6f0b..8a5f4dc2f2d 100644
--- a/testing/test/extensions/public/AzureMonitorHighScale.Tests.ps1
+++ b/testing/test/extensions/public/AzureMonitorHighScale.Tests.ps1
@@ -25,7 +25,6 @@ Describe 'Azure Monitor High Scale Mode Testing' {
         # Verify high scale mode configuration
         $settings = $extension.configurationSettings
         $settings.'amalogs.enableHighLogScaleMode' | Should -Be "true"
-        $settings.'amalogs.useAADAuth' | Should -Be "true"
 
         # Loop and retry until the extension installs
         $n = 0

From 09a303494b74cf08761c9ff3d944131d08b3d687 Mon Sep 17 00:00:00 2001
From: longwan <longwan@microsoft.com>
Date: Thu, 17 Jul 2025 19:52:27 +0000
Subject: [PATCH 5/7] update revert version bump

---
 src/k8s-extension/HISTORY.rst | 4 ----
 src/k8s-extension/setup.py    | 2 +-
 2 files changed, 1 insertion(+), 5 deletions(-)

diff --git a/src/k8s-extension/HISTORY.rst b/src/k8s-extension/HISTORY.rst
index 94df11be9ce..09a13348d04 100644
--- a/src/k8s-extension/HISTORY.rst
+++ b/src/k8s-extension/HISTORY.rst
@@ -3,10 +3,6 @@
 Release History
 ===============
 
-1.6.6
-++++++++
-* microsoft.azuremonitor.containers: Extend ContainerInsights Extension for high log scale mode support.
-
 1.6.5
 ++++++++++++++++++
 * microsoft.dataprotection.kubernetes: Add support for 'DisableInformerCache' configuration.
diff --git a/src/k8s-extension/setup.py b/src/k8s-extension/setup.py
index bdee2be8241..c46e1aca890 100644
--- a/src/k8s-extension/setup.py
+++ b/src/k8s-extension/setup.py
@@ -33,7 +33,7 @@
 # TODO: Add any additional SDK dependencies here
 DEPENDENCIES = []
 
-VERSION = "1.6.6"
+VERSION = "1.6.5"
 
 with open("README.rst", "r", encoding="utf-8") as f:
     README = f.read()

From d90834acb39265d7655054afdd1877bce03406a8 Mon Sep 17 00:00:00 2001
From: longwan <longwan@microsoft.com>
Date: Fri, 18 Jul 2025 17:49:44 +0000
Subject: [PATCH 6/7] add check

---
 .../azext_k8s_extension/partner_extensions/ContainerInsights.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/k8s-extension/azext_k8s_extension/partner_extensions/ContainerInsights.py b/src/k8s-extension/azext_k8s_extension/partner_extensions/ContainerInsights.py
index 5cf74793808..70c249099c3 100644
--- a/src/k8s-extension/azext_k8s_extension/partner_extensions/ContainerInsights.py
+++ b/src/k8s-extension/azext_k8s_extension/partner_extensions/ContainerInsights.py
@@ -133,6 +133,8 @@ def Delete(self, cmd, client, resource_group_name, cluster_name, name, cluster_t
                     enable_high_log_scale_mode = (highLogScaleSetting.lower() == "true")
                 elif isinstance(highLogScaleSetting, bool):
                     enable_high_log_scale_mode = highLogScaleSetting
+                else:
+                    raise InvalidArgumentValueError('amalogs.enableHighLogScaleMode value MUST be either true/false or boolean type')
         if useAADAuth:
             association_url = cmd.cli_ctx.cloud.endpoints.resource_manager + f"{cluster_resource_id}/providers/Microsoft.Insights/dataCollectionRuleAssociations/ContainerInsightsExtension?api-version={DCR_API_VERSION}"
             for _ in range(3):

From 832c52e48585709a02edb6f71a627a091e7a183c Mon Sep 17 00:00:00 2001
From: longwan <longwan@microsoft.com>
Date: Mon, 21 Jul 2025 19:49:50 +0000
Subject: [PATCH 7/7] update delete logic using workspace_region

---
 .../partner_extensions/ContainerInsights.py   | 23 ++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/src/k8s-extension/azext_k8s_extension/partner_extensions/ContainerInsights.py b/src/k8s-extension/azext_k8s_extension/partner_extensions/ContainerInsights.py
index 70c249099c3..2b490f78005 100644
--- a/src/k8s-extension/azext_k8s_extension/partner_extensions/ContainerInsights.py
+++ b/src/k8s-extension/azext_k8s_extension/partner_extensions/ContainerInsights.py
@@ -114,8 +114,12 @@ def Delete(self, cmd, client, resource_group_name, cluster_name, name, cluster_t
         resources = cf_resources(cmd.cli_ctx, subscription_id)
         # handle cluster type here
         cluster_resource_id = '/subscriptions/{0}/resourceGroups/{1}/providers/{2}/{3}/{4}'.format(subscription_id, resource_group_name, cluster_rp, cluster_type, cluster_name)
+        workspace_resource_id = None
         if (extension is not None) and (extension.configuration_settings is not None):
             configSettings = extension.configuration_settings
+            # Extract workspace resource ID if present
+            if 'logAnalyticsWorkspaceResourceID' in configSettings:
+                workspace_resource_id = configSettings['logAnalyticsWorkspaceResourceID']
             # omsagent is being renamed to ama-logs. Check for both for compatibility
             if 'omsagent.useAADAuth' in configSettings:
                 useAADAuthSetting = configSettings['omsagent.useAADAuth']
@@ -159,9 +163,22 @@ def Delete(self, cmd, client, resource_group_name, cluster_name, name, cluster_t
                     pass  # its OK to ignore the exception since MSI auth in preview
 
         if useAADAuth:
-            resource = resources.get_by_id(cluster_resource_id, '2020-01-01-preview')
-            cluster_location = resource.location.lower()
-            dcr_name = f"MSCI-{cluster_location}-{cluster_name}"
+            # Get the workspace region if workspace_resource_id is available
+            workspace_region = None
+            if workspace_resource_id:
+                try:
+                    workspace_resource = resources.get_by_id(workspace_resource_id, '2015-11-01-preview')
+                    workspace_region = workspace_resource.location.replace(" ", "").lower()
+                except Exception as ex:
+                    logger.warning("Skipping DCR and DCE deletion due to inability to determine workspace region")
+                    return
+            # If workspace_region still couldn't be determined, skip deletion
+            if not workspace_region:
+                logger.warning("Workspace region could not be determined. Skipping DCR and DCE deletion.")
+                return
+
+            # Use workspace_region for DCR name to match creation logic
+            dcr_name = f"MSCI-{workspace_region}-{cluster_name}"
             dcr_name = dcr_name[0:64]
             
             dcr_resource_id = f"/subscriptions/{subscription_id}/resourceGroups/{resource_group_name}/providers/Microsoft.Insights/dataCollectionRules/{dcr_name}"