From 4b10cd19e5dd0308354a2a7a91b90f6095aa019e Mon Sep 17 00:00:00 2001
From: YASH <139299779+Yash-2707@users.noreply.github.com>
Date: Wed, 2 Oct 2024 22:12:40 +0530
Subject: [PATCH] Update simulate_rules.py

Added a function that checks whether the API key and API url are present or not and give the error according if not present , and added try catch block also

Signed-off-by: YASH <139299779+Yash-2707@users.noreply.github.com>
---
 scripts/simulate_rules.py | 109 +++++++++++++++++++++++---------------
 1 file changed, 67 insertions(+), 42 deletions(-)

diff --git a/scripts/simulate_rules.py b/scripts/simulate_rules.py
index 7a98ba99f..958b9e8ac 100644
--- a/scripts/simulate_rules.py
+++ b/scripts/simulate_rules.py
@@ -1,7 +1,6 @@
-# A script that simulates the creation of rules in the database
-# Written for demonstration purposes only
 import logging
 import os
+import sys
 
 from keep.api.core.db import create_rule, get_api_key
 
@@ -12,54 +11,80 @@
 keep_api_url = os.environ.get("KEEP_API_URL")
 
 
+def validate_env_vars():
+    if not keep_api_key:
+        logger.error("KEEP_API_KEY is missing. Please set it in the environment.")
+        sys.exit(1)
+    if not keep_api_url:
+        logger.error("KEEP_API_URL is missing. Please set it in the environment.")
+        sys.exit(1)
+
+
+def create_rule_in_db(tenant_id, name, timeframe, definition, definition_cel, created_by, grouping_criteria=None, group_description=None):
+    try:
+        create_rule(
+            tenant_id=tenant_id,
+            name=name,
+            timeframe=timeframe,
+            definition=definition,
+            definition_cel=definition_cel,
+            created_by=created_by,
+            grouping_criteria=grouping_criteria,
+            group_description=group_description,
+        )
+        logger.info(f"Rule '{name}' inserted successfully.")
+    except Exception as e:
+        logger.error(f"Failed to create rule '{name}': {e}")
+        sys.exit(1)
+
+
 def simulate_rules():
     # Define the tenant ID
     logger.info("Inserting Rules")
     tenant_id = get_api_key(keep_api_key).tenant_id
     created_by = "keep"
-    definition_template = '{{"sql": "((labels.alertname like :labels.alertname_1))", "params": {{"labels.alertname_1": "%cpu%"}}}}'
-    definition_cel_template = '(labels.alertname.contains("cpu"))'
-    timeframe = 600
-
-    # Rule #1 - CPU (group by labels.instance)
-    create_rule(
-        tenant_id=tenant_id,
-        name="CPU (group by labels.instance)",
-        timeframe=timeframe,
-        definition=definition_template,
-        definition_cel=definition_cel_template,
-        created_by=created_by,
-        grouping_criteria=["labels.instance"],
-        group_description="CPU usage exceeded on {{ group_attributes.num_of_alerts }} pods of {{ labels.instance }} || {{ group_attributes.start_time }} | {{ group_attributes.last_update_time }}",
-    )
+    
+    # Rule configurations
+    rules = [
+        {
+            "name": "CPU (group by labels.instance)",
+            "timeframe": 600,
+            "definition": '{"sql": "((labels.alertname like :labels.alertname_1))", "params": {"labels.alertname_1": "%cpu%"}}',
+            "definition_cel": '(labels.alertname.contains("cpu"))',
+            "grouping_criteria": ["labels.instance"],
+            "group_description": "CPU usage exceeded on {{ group_attributes.num_of_alerts }} pods of {{ labels.instance }} || {{ group_attributes.start_time }} | {{ group_attributes.last_update_time }}"
+        },
+        {
+            "name": "CPU (no grouping)",
+            "timeframe": 600,
+            "definition": '{"sql": "((labels.alertname like :labels.alertname_1))", "params": {"labels.alertname_1": "%cpu%"}}',
+            "definition_cel": '(labels.alertname.contains("cpu"))',
+        },
+        {
+            "name": "MQ (group by labels.queue)",
+            "timeframe": 600,
+            "definition": '{"sql": "((name = :name_1))", "params": {"name_1": "mq_third_full"}}',
+            "definition_cel": '(name == "mq_third_full")',
+            "grouping_criteria": ["labels.queue"],
+            "group_description": "The {{ labels.queue }} is more than third full on {{ group_attributes.num_of_alerts }} queue managers | {{ group_attributes.start_time }} || {{ group_attributes.last_update_time }}"
+        }
+    ]
 
-    # Rule #2 - CPU (no grouping)
-    create_rule(
-        tenant_id=tenant_id,
-        name="CPU (no grouping)",
-        timeframe=timeframe,
-        definition=definition_template,
-        definition_cel=definition_cel_template,
-        created_by=created_by,
-    )
+    for rule in rules:
+        create_rule_in_db(
+            tenant_id=tenant_id,
+            name=rule["name"],
+            timeframe=rule["timeframe"],
+            definition=rule["definition"],
+            definition_cel=rule["definition_cel"],
+            created_by=created_by,
+            grouping_criteria=rule.get("grouping_criteria"),
+            group_description=rule.get("group_description")
+        )
 
-    # Rule #3 - MQ (group by labels.queue)
-    mq_definition_template = (
-        '{{"sql": "((name = :name_1))", "params": {{"name_1": "mq_third_full"}}}}'
-    )
-    mq_definition_cel_template = '(name == "mq_third_full")'
-    create_rule(
-        tenant_id=tenant_id,
-        name="MQ (group by labels.queue)",
-        timeframe=timeframe,
-        definition=mq_definition_template,
-        definition_cel=mq_definition_cel_template,
-        created_by=created_by,
-        grouping_criteria=["labels.queue"],
-        group_description="The {{ labels.queue }} is more than third full on {{ group_attributes.num_of_alerts }} queue managers | {{ group_attributes.start_time }} || {{ group_attributes.last_update_time }}",
-    )
-    logger.info("Rules inserted successfully")
+    logger.info("All rules inserted successfully.")
 
 
 if __name__ == "__main__":
+    validate_env_vars()
     simulate_rules()