Merge pull request #65 from DrDroidLab/feature/multi-task-result-vpc

intiialise support for multiple task results

Author: dhruv-droid-r2

agent/settings.py

@@ -11,6 +11,7 @@
 """
 
 import os
+import uuid
 from pathlib import Path
 
 import yaml
@@ -29,7 +30,8 @@ def load_yaml(filepath, native_k8s_connector_mode=False):
     if native_k8s_connector_mode:
         if not loaded_connection:
             loaded_connection = {}
-        loaded_connection.update({'native_k8s_connector': {'type': 'KUBERNETES'}})
+        uuid_hex = uuid.uuid4().hex
+        loaded_connection.update({f'native_k8_connection_{uuid_hex}': {'type': 'KUBERNETES'}})
     return loaded_connection
 
 

integrations/source_api_processors/grafana_api_processor.py

@@ -22,7 +22,7 @@ def __init__(self, grafana_host, grafana_api_key, ssl_verify='true'):
     def test_connection(self):
         try:
             url = '{}/api/datasources'.format(self.__host)
-            response = requests.get(url, headers=self.headers, verify=self.__ssl_verify)
+            response = requests.get(url, headers=self.headers, verify=self.__ssl_verify, timeout=20)
             if response and response.status_code == 200:
                 return True
             else:
@@ -99,17 +99,41 @@ def fetch_promql_metric_timeseries(self, promql_datasource_uid, query, start, en
             logger.error(f"Exception occurred while getting promql metric timeseries with error: {e}")
             raise e
 
-    def panel_query_datasource_api(self, tr: TimeRange, queries):
+    def fetch_alert_rules(self):
+        try:
+            url = '{}/api/v1/provisioning/alert-rules'.format(self.__host)
+            response = requests.get(url, headers=self.headers, verify=self.__ssl_verify)
+            if response and response.status_code == 200:
+                return response.json()
+            else:
+                raise Exception(
+                    f"Failed to fetch alert rules. Status Code: {response.status_code}. Response Text: {response.text}")
+        except Exception as e:
+            logger.error(f"Exception occurred while fetching grafana alert rules with error: {e}")
+            raise e
+
+    def panel_query_datasource_api(self, tr: TimeRange, queries, interval_ms=300000):
         try:
             if not queries or len(queries) == 0:
                 raise ValueError("No queries provided.")
+
             url = f"{self.__host}/api/ds/query"
+
             from_tr = int(tr.time_geq * 1000)
             to_tr = int(tr.time_lt * 1000)
-            payload = {"queries": queries, "from": str(from_tr), "to": str(to_tr)}
+
+            for query in queries:
+                query["intervalMs"] = interval_ms # 5 minutes default, in milliseconds
+
+            payload = {
+                "queries": queries,
+                "from": str(from_tr),
+                "to": str(to_tr)
+            }
+
             response = requests.post(url, headers=self.headers, json=payload)
             if response.status_code == 429:
-                logger.error("Grafana query API responded with 429 (rate limited). Headers: %s", response.headers)
+                logger.info("Grafana query API responded with 429 (rate limited). Headers: %s", response.headers)
                 return None
             elif response.status_code == 200:
                 return response.json()

integrations/source_facade.py

@@ -55,7 +55,7 @@ def get_source_manager(self, source: Source):
             raise ValueError(f'No executor found for source: {source}')
         return self._map.get(source)
 
-    def execute_task(self, time_range, global_variable_set, task: PlaybookTask) -> PlaybookTaskResult:
+    def execute_task(self, time_range, global_variable_set, task: PlaybookTask):
         source = task.source
         if source not in self._map:
             raise ValueError(f'No executor found for source: {source}')

integrations/source_manager.py

@@ -14,7 +14,7 @@
     PlaybookExecutionStatusType
 from protos.playbooks.playbook_pb2 import PlaybookTask
 from utils.proto_utils import proto_to_dict, dict_to_proto
-
+from integrations.utils.executor_utils import check_multiple_task_results
 logger = logging.getLogger(__name__)
 
 
@@ -126,7 +126,7 @@ def get_resolved_task(self, global_variable_set: Struct, input_task: PlaybookTas
 
         return resolved_task, resolved_source_task_proto, task_local_variable_map
 
-    def execute_task(self, time_range: TimeRange, global_variable_set, task: PlaybookTask) -> PlaybookTaskResult:
+    def execute_task(self, time_range: TimeRange, global_variable_set, task: PlaybookTask):
         try:
             source_connector_proto = None
             if task.task_connector_sources and len(task.task_connector_sources) > 0:
@@ -142,18 +142,26 @@ def execute_task(self, time_range: TimeRange, global_variable_set, task: Playboo
             try:
                 # Execute task
                 task_type = resolved_source_task.type
-                playbook_task_result: PlaybookTaskResult = self.task_type_callable_map[task_type]['executor'](
+                playbook_task_result = self.task_type_callable_map[task_type]['executor'](
                     time_range, resolved_source_task, source_connector_proto)
-                # Set task local variables in playbook_task_result to be stored in database
-                task_local_variable_map_proto = dict_to_proto(task_local_variable_map,
-                                                              Struct) if task_local_variable_map else Struct()
-                playbook_task_result.task_local_variable_set.CopyFrom(task_local_variable_map_proto)
-                playbook_task_result.status = PlaybookExecutionStatusType.FINISHED
-                # Apply result transformer
-                playbook_task_result = self.apply_task_result_transformer(resolved_task, playbook_task_result)
-                return playbook_task_result
+                if check_multiple_task_results(playbook_task_result):
+                    task_results = []
+                    for result in playbook_task_result:
+                        task_results.append(self.postprocess_task_result(result, resolved_task, task_local_variable_map))
+                    return task_results
+                return self.postprocess_task_result(playbook_task_result, resolved_task, task_local_variable_map)
             except Exception as e:
                 source_str = Source.Name(resolved_task.source).lower()
                 raise Exception(f"Error while executing task for source: {source_str} with error: {e}")
         except Exception as e:
             raise Exception(f"Error while executing task: {e}")
+
+    def postprocess_task_result(self, playbook_task_result: PlaybookTaskResult, resolved_task: PlaybookTask, task_local_variable_map: dict):
+        task_local_variable_map_proto = dict_to_proto(task_local_variable_map,
+                                                              Struct) if task_local_variable_map else Struct()
+        playbook_task_result.task_local_variable_set.CopyFrom(task_local_variable_map_proto)
+        playbook_task_result.status = PlaybookExecutionStatusType.FINISHED
+
+        # Apply result transformer
+        playbook_task_result = self.apply_task_result_transformer(resolved_task, playbook_task_result)
+        return playbook_task_result

integrations/source_manangers/api_source_manager.py

@@ -80,7 +80,7 @@ def __init__(self):
         }
 
     def execute_http_request(self, time_range: TimeRange, api_task: Api,
-                             api_connector_proto: ConnectorProto) -> PlaybookTaskResult:
+                             api_connector_proto: ConnectorProto):
         try:
             http_request = api_task.http_request
             method = http_request.method

integrations/source_manangers/argocd_source_manager.py

@@ -85,7 +85,7 @@ def get_connector_processor(self, argocd_connector, **kwargs):
         return ArgoCDAPIProcessor(**generated_credentials)
 
     def fetch_deployment_info(self, time_range: TimeRange, argocd_task: ArgoCD,
-                              argocd_connector: ConnectorProto) -> PlaybookTaskResult:
+                              argocd_connector: ConnectorProto):
         # Loop through the commits and get the diff for each one
         try:
             deployment_info = self.get_connector_processor(argocd_connector).get_deployment_info()
@@ -139,7 +139,7 @@ def fetch_deployment_info(self, time_range: TimeRange, argocd_task: ArgoCD,
             raise Exception(f"Error while executing ArgoCD fetch_deployment_info task: {e}")
 
     def rollback_application(self, time_range: TimeRange, argocd_task: ArgoCD,
-                             argocd_connector: ConnectorProto) -> PlaybookTaskResult:
+                             argocd_connector: ConnectorProto):
         try:
             app_name = argocd_task.rollback_application.app_name
             revision = argocd_task.rollback_application.revision

integrations/source_manangers/azure_source_manager.py

@@ -53,7 +53,7 @@ def get_connector_processor(self, azure_connector, **kwargs):
         return AzureApiProcessor(**generated_credentials)
 
     def filter_log_events(self, time_range: TimeRange, azure_task: Azure,
-                          azure_connector: ConnectorProto) -> PlaybookTaskResult:
+                          azure_connector: ConnectorProto):
         try:
             tr_end_time = time_range.time_lt
             end_time = int(tr_end_time * 1000)

integrations/source_manangers/bash_source_manager.py

@@ -53,7 +53,7 @@ def get_connector_processor(self, bash_connector, **kwargs):
         return BashProcessor(**generated_credentials)
 
     def execute_command(self, time_range: TimeRange, bash_task: Bash,
-                        remote_server_connector: ConnectorProto) -> PlaybookTaskResult:
+                        remote_server_connector: ConnectorProto):
         try:
             bash_command: Bash.Command = bash_task.command
             remote_server_str = bash_command.remote_server.value if bash_command.remote_server else None

integrations/source_manangers/big_query_source_manager.py

@@ -42,7 +42,7 @@ def get_connector_processor(self, bq_connector, **kwargs):
         return BigQueryApiProcessor(**generated_credentials)
 
     def execute_query_table(self, time_range: TimeRange, bq_task: BigQuery,
-                            bq_connector: ConnectorProto) -> PlaybookTaskResult:
+                            bq_connector: ConnectorProto):
         try:
             if not bq_connector:
                 raise Exception("Task execution Failed:: No BigQuery source found")

integrations/source_manangers/clickhouse_source_manager.py

@@ -55,7 +55,7 @@ def get_connector_processor(self, clickhouse_connector, **kwargs):
         return ClickhouseDBProcessor(**generated_credentials)
 
     def execute_sql_query(self, time_range: TimeRange, clickhouse_task: SqlDataFetch,
-                          clickhouse_connector: ConnectorProto) -> PlaybookTaskResult:
+                          clickhouse_connector: ConnectorProto):
         try:
             if not clickhouse_connector:
                 raise Exception("Task execution Failed:: No Clickhouse source found")