Update create.py

Author: bragi92

src/k8s-extension/azext_k8s_extension/partner_extensions/azuremonitormetrics/recordingrules/create.py

@@ -3,8 +3,9 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 import json
-from ..constants import ALERTS_API, RULES_API
 from knack.util import CLIError
+from ..constants import ALERTS_API, RULES_API
+from azure.cli.command_modules.acs.azuremonitormetrics.recordingrules.common import truncate_rule_group_name
 
 
 # pylint: disable=line-too-long
@@ -15,7 +16,18 @@ def get_recording_rules_template(cmd, azure_monitor_workspace_resource_id):
     url = f"{armendpoint}{azure_monitor_workspace_resource_id}/providers/microsoft.alertsManagement/alertRuleRecommendations?api-version={ALERTS_API}"
     r = send_raw_request(cmd.cli_ctx, "GET", url, headers=headers)
     data = json.loads(r.text)
-    return data['value']
+
+    filtered_templates = [
+        template for template in data.get('value', [])
+        if template.get("properties", {}).get("alertRuleType", "").lower() == "microsoft.alertsmanagement/prometheusrulegroups" and isinstance(template.get("properties", {}).get("rulesArmTemplate", {}).get("resources"), list) and all(
+            isinstance(rule, dict) and "record" in rule and "expression" in rule
+            for resource in template["properties"]["rulesArmTemplate"]["resources"]
+            if resource.get("type", "").lower() == "microsoft.alertsmanagement/prometheusrulegroups"
+            for rule in resource.get("properties", {}).get("rules", [])
+        )
+    ]
+
+    return filtered_templates
 
 
 # pylint: disable=line-too-long
@@ -39,8 +51,7 @@ def put_rules(cmd, default_rule_group_id, default_rule_group_name, mac_region, a
     for _ in range(3):
         try:
             headers = ['User-Agent=arc-azuremonitormetrics.put_rules.' + default_rule_group_name]
-            send_raw_request(cmd.cli_ctx, "PUT", url,
-                             body=body, headers=headers)
+            send_raw_request(cmd.cli_ctx, "PUT", url, body=body, headers=headers)
             break
         except CLIError as e:
             error = e
@@ -49,30 +60,30 @@ def put_rules(cmd, default_rule_group_id, default_rule_group_name, mac_region, a
 
 
 # pylint: disable=line-too-long
-def create_rules(cmd, cluster_subscription, cluster_resource_group_name, cluster_name, azure_monitor_workspace_resource_id, mac_region):
+def create_rules(cmd, cluster_subscription, cluster_resource_group_name, cluster_name, azure_monitor_workspace_resource_id, mac_region, raw_parameters):
     default_rules_template = get_recording_rules_template(cmd, azure_monitor_workspace_resource_id)
-    default_rule_group_name = "NodeRecordingRulesRuleGroup-{0}".format(cluster_name)
-    default_rule_group_id = "/subscriptions/{0}/resourceGroups/{1}/providers/Microsoft.AlertsManagement/prometheusRuleGroups/{2}".format(
-        cluster_subscription,
-        cluster_resource_group_name,
-        default_rule_group_name
-    )
-    url = "{0}{1}?api-version={2}".format(
-        cmd.cli_ctx.cloud.endpoints.resource_manager,
-        default_rule_group_id,
-        RULES_API
-    )
-    put_rules(cmd, default_rule_group_id, default_rule_group_name, mac_region, azure_monitor_workspace_resource_id, cluster_name, default_rules_template, url, True, 0)
 
-    default_rule_group_name = "KubernetesRecordingRulesRuleGroup-{0}".format(cluster_name)
-    default_rule_group_id = "/subscriptions/{0}/resourceGroups/{1}/providers/Microsoft.AlertsManagement/prometheusRuleGroups/{2}".format(
-        cluster_subscription,
-        cluster_resource_group_name,
-        default_rule_group_name
-    )
-    url = "{0}{1}?api-version={2}".format(
-        cmd.cli_ctx.cloud.endpoints.resource_manager,
-        default_rule_group_id,
-        RULES_API
-    )
-    put_rules(cmd, default_rule_group_id, default_rule_group_name, mac_region, azure_monitor_workspace_resource_id, cluster_name, default_rules_template, url, True, 1)
+    enable_windows_recording_rules = raw_parameters.get("enable_windows_recording_rules", False)
+
+    for index, rule_template in enumerate(default_rules_template):
+        rule_name = rule_template["name"]
+        is_windows_rule = "win" in rule_name.lower()
+
+        enable_rules = not (is_windows_rule and not enable_windows_recording_rules)
+
+        rule_group_name = truncate_rule_group_name(f"{rule_template['name']}-{cluster_name}")
+        rule_group_id = f"/subscriptions/{cluster_subscription}/resourceGroups/{cluster_resource_group_name}/providers/Microsoft.AlertsManagement/prometheusRuleGroups/{rule_group_name}"
+        url = f"{cmd.cli_ctx.cloud.endpoints.resource_manager}{rule_group_id}?api-version={RULES_API}"
+
+        put_rules(
+            cmd,
+            rule_group_id,
+            rule_group_name,
+            mac_region,
+            azure_monitor_workspace_resource_id,
+            cluster_name,
+            default_rules_template,
+            url,
+            enable_rules,
+            index
+        )