fix new relic tool

robusta-dev · Feb 4, 2025 · 12f27cc · 12f27cc
1 parent 93af740
commit 12f27cc
Showing 1 changed file with 13 additions and 13 deletions.
diff --git a/holmes/plugins/toolsets/newrelic.py b/holmes/plugins/toolsets/newrelic.py
@@ -2,9 +2,9 @@
 import requests
 import json
 import logging
-from typing import Any
+from typing import Any, Optional
 from holmes.core.tools import CallablePrerequisite, Tool, ToolParameter, Toolset, ToolsetTag
-from pydantic import BaseModel, Optional
+from pydantic import BaseModel
 
 class BaseNewRelicTool(Tool):
     toolset: "NewRelicToolset"
@@ -37,7 +37,7 @@ def invoke(self, params: Any) -> str:
             "query": f"""
             {{
                 actor {{
-                    account(id: {self.toolset.account_id}) {{
+                    account(id: {self.toolset.nr_account_id}) {{
                         nrql(query: \"SELECT * FROM Log WHERE app = '{app}' SINCE {since}\") {{
                             results
                         }}
@@ -50,7 +50,7 @@ def invoke(self, params: Any) -> str:
         url = "https://api.newrelic.com/graphql"
         headers = {
             "Content-Type": "application/json",
-            "Api-Key": self.toolset.api_key,
+            "Api-Key": self.toolset.nr_api_key,
         }
 
         response = requests.post(url, headers=headers, json=query)
@@ -99,7 +99,7 @@ def invoke(self, params: Any) -> str:
             "query": f"""
             {{
                 actor {{
-                    account(id: {self.toolset.account_id}) {{
+                    account(id: {self.toolset.nr_account_id}) {{
                         nrql(query: \"{query_string}\") {{
                             results
                         }}
@@ -112,7 +112,7 @@ def invoke(self, params: Any) -> str:
         url = "https://api.newrelic.com/graphql"
         headers = {
             "Content-Type": "application/json",
-            "Api-Key": self.toolset.api_key,
+            "Api-Key": self.toolset.nr_api_key,
         }
 
         response = requests.post(url, headers=headers, json=query)
@@ -130,12 +130,12 @@ def get_parameterized_one_liner(self, params) -> str:
 
 
 class NewrelicConfig(BaseModel):
-    api_key: Optional[str] = None
-    account_id: Optional[str] = None
+    nr_api_key: Optional[str] = None
+    nr_account_id: Optional[str] = None
 
 class NewRelicToolset(Toolset):
-    api_key: Optional[str] = None
-    account_id: Optional[str] = None
+    nr_api_key: Optional[str] = None
+    nr_account_id: Optional[str] = None
 
     def __init__(self):
         super().__init__(
@@ -157,9 +157,9 @@ def prerequisites_callable(self, config: dict[str, Any]) -> bool:
 
         try:
             nr_config = NewrelicConfig(**config)
-            self.account_id = nr_config.account_id
-            self.api_key = nr_config.api_key
-            return self.account_id and self.api_key
+            self.nr_account_id = nr_config.nr_account_id
+            self.nr_api_key = nr_config.nr_api_key
+            return self.nr_account_id and self.nr_api_key
         except Exception:
             logging.exception("Failed to set up new relic toolset")
             return False