ContainerInsights extension - Extend dataCollectionSettings config se…

…ttings with streams field (Azure#232) * extend containerinsights datacollection settings with streams field * bug fix * fix lint issues * fix pr feedback * fix pr feedback * fix lint error
ms-hujia · Jun 30, 2023 · dfd1139 · dfd1139
1 parent 30a373b
commit dfd1139
Showing 1 changed file with 16 additions and 11 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
@@ -31,6 +31,7 @@
     cf_resources, cf_resource_groups, cf_log_analytics)
 
 logger = get_logger(__name__)
+DCR_API_VERSION = "2022-06-01"
 
 
 class ContainerInsights(DefaultExtension):
@@ -100,7 +101,7 @@ def Delete(self, cmd, client, resource_group_name, cluster_name, name, cluster_t
                 if (isinstance(useAADAuthSetting, str) and str(useAADAuthSetting).lower() == "true") or (isinstance(useAADAuthSetting, bool) and useAADAuthSetting):
                     useAADAuth = True
         if useAADAuth:
-            association_url = cmd.cli_ctx.cloud.endpoints.resource_manager + f"{cluster_resource_id}/providers/Microsoft.Insights/dataCollectionRuleAssociations/ContainerInsightsExtension?api-version=2021-04-01"
+            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):
                 try:
                     send_raw_request(cmd.cli_ctx, "GET", association_url,)
@@ -114,7 +115,7 @@ 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 isDCRAExists:
-            association_url = cmd.cli_ctx.cloud.endpoints.resource_manager + f"{cluster_resource_id}/providers/Microsoft.Insights/dataCollectionRuleAssociations/ContainerInsightsExtension?api-version=2021-04-01"
+            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):
                 try:
                     send_raw_request(cmd.cli_ctx, "DELETE", association_url,)
@@ -495,6 +496,10 @@ def _get_container_insights_settings(cmd, cluster_resource_group_name, cluster_r
                 namspaces = dataCollectionSettings["namespaces"]
                 if isinstance(namspaces, list) is False:
                     raise InvalidArgumentValueError('namespaces must be an array type')
+            if 'streams' in dataCollectionSettings.keys():
+                streams = dataCollectionSettings["streams"]
+                if isinstance(streams, list) is False:
+                    raise InvalidArgumentValueError('streams must be an array type')
             extensionSettings["dataCollectionSettings"] = dataCollectionSettings
 
     workspace_resource_id = workspace_resource_id.strip()
@@ -673,9 +678,14 @@ def _ensure_container_insights_dcr_for_monitoring(cmd, subscription_id, cluster_
             if (cluster_region not in region_ids):
                 raise ClientRequestError(f"Data Collection Rule Associations are not supported for cluster region {cluster_region}")
 
-    dcr_url = cmd.cli_ctx.cloud.endpoints.resource_manager + f"{dcr_resource_id}?api-version=2021-04-01"
+    dcr_url = cmd.cli_ctx.cloud.endpoints.resource_manager + f"{dcr_resource_id}?api-version={DCR_API_VERSION}"
     # 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 extensionSettings is not None and 'dataCollectionSettings' in extensionSettings.keys():
+        dataCollectionSettings = extensionSettings["dataCollectionSettings"]
+        if dataCollectionSettings is not None and 'streams' in dataCollectionSettings.keys():
+            streams = dataCollectionSettings["streams"]
 
     # create the DCR
     dcr_creation_body = json.dumps(
@@ -687,20 +697,15 @@ def _ensure_container_insights_dcr_for_monitoring(cmd, subscription_id, cluster_
                     "extensions": [
                         {
                             "name": "ContainerInsightsExtension",
-                            "streams": [
-                                "Microsoft-ContainerInsights-Group-Default"
-                            ],
+                            "streams": streams,
                             "extensionName": "ContainerInsights",
                             "extensionSettings": extensionSettings
                         }
                     ]
                 },
                 "dataFlows": [
                     {
-                        "streams": [
-                            "Microsoft-ContainerInsights-Group-Default"
-
-                        ],
+                        "streams": streams,
                         "destinations": ["la-workspace"],
                     }
                 ],
@@ -735,7 +740,7 @@ def _ensure_container_insights_dcr_for_monitoring(cmd, subscription_id, cluster_
             },
         }
     )
-    association_url = cmd.cli_ctx.cloud.endpoints.resource_manager + f"{cluster_resource_id}/providers/Microsoft.Insights/dataCollectionRuleAssociations/ContainerInsightsExtension?api-version=2021-04-01"
+    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):
         try:
             send_raw_request(cmd.cli_ctx, "PUT", association_url, body=association_body,)