Merge pull request #47 from codedwithhs/feature-AIRA-36-Connector---Implement-Datadog

AIRA-36: Connector - Implement Datadog

Author: learn-design-develop

aira/config.py

@@ -51,13 +51,21 @@ class JSMConfig(BaseModel):
     api_token: SecretStr
 
 
+class DatadogConfig(BaseModel):
+    type: Literal["datadog"]
+    api_key: SecretStr
+    app_key: SecretStr
+    # Datadog site URL varies by region (e.g., datadoghq.com, datadoghq.eu)
+    site: Optional[str] = "datadoghq.com"
+
+
 class SlackConfig(BaseModel):
     type: Literal["slack"]
     webhook_url: SecretStr
 
 
 # Discriminated Unions for Connectors and Actions
-AnyConnection = Union[GitHubConfig, PagerDutyConfig, JSMConfig]
+AnyConnection = Union[GitHubConfig, PagerDutyConfig, JSMConfig, DatadogConfig]
 AnyAction = Union[SlackConfig]
 
 

aira/connectors/__init__.py

@@ -5,7 +5,8 @@
 from .source_control import GitHubConnector
 from .alerting import PagerDutyConnector
 from .alerting import JSMConnector
-from .collaboration.slack import SlackConnector
+from .collaboration import SlackConnector
+from .observability import DatadogConnector
 
 
 # The Registry Map: Maps a 'type' string from config to the actual class.
@@ -15,6 +16,7 @@
     "jsm": JSMConnector,
     "github": GitHubConnector,
     "slack": SlackConnector,
+    "datadog": DatadogConnector,
 }
 
 

aira/connectors/observability/__init__.py

@@ -0,0 +1,7 @@
+# aira/connectors/obervability/__init__.py
+
+from .datadog import DatadogConnector
+
+# This line explicitly declares which names are part of this package's
+# public interface, silencing the "unused import" warning.
+__all__ = ["DatadogConnector"]

aira/connectors/observability/datadog.py

@@ -0,0 +1,99 @@
+import requests
+from typing import Dict, Any, Tuple
+from datetime import datetime, timedelta, timezone
+
+from ..base import ObservabilityProvider
+from ...config import DatadogConfig
+
+
+class DatadogConnector(ObservabilityProvider):
+    """
+    Connector for interacting with the Datadog API.
+    """
+
+    def __init__(self, name: str, config: Dict[str, Any]):
+        """
+        Initializes the connector and validates its specific configuration.
+        """
+        super().__init__(name, config)
+        self.validated_config = DatadogConfig(**self.config)
+        self.api_base_url = f"https://api.{self.validated_config.site}"
+        self.headers = {
+            "DD-API-KEY": self.validated_config.api_key.get_secret_value(),
+            "DD-APPLICATION-KEY": self.validated_config.app_key.get_secret_value(),
+            "Content-Type": "application/json",
+        }
+
+    def test_connection(self) -> Tuple[bool, str]:
+        """
+        Validates the Datadog API and App keys by making a lightweight API call.
+        """
+        try:
+            # The validate endpoint is designed for this purpose
+            response = requests.get(
+                f"{self.api_base_url}/api/v1/validate", headers=self.headers, timeout=10
+            )
+            response.raise_for_status()
+            if response.json().get("valid"):
+                return True, "Datadog connection successful and keys are valid."
+            else:
+                return (
+                    False,
+                    "Datadog connection failed: The provided keys are not valid.",
+                )
+        except requests.exceptions.HTTPError as e:
+            if e.response.status_code in [401, 403]:
+                return False, "Connection failed: Invalid Datadog API or App Key."
+            return False, f"Connection failed: HTTP {e.response.status_code} error."
+        except requests.exceptions.RequestException as e:
+            return False, f"Connection failed: Network error - {e}."
+
+    def fetch_logs(self, query: str, time_window_minutes: int = 15) -> str:
+        """
+        Fetches and formats logs from Datadog Logs.
+
+        Args:
+            query (str): The search query to execute (e.g., 'service:api-checkout status:error').
+            time_window_minutes (int): The number of minutes to look back for logs.
+
+        Returns:
+            A formatted string of log lines or an error/empty message.
+        """
+        url = f"{self.api_base_url}/api/v2/logs/events/search"
+        print(f"-> Fetching logs from Datadog with query: '{query}'...")
+
+        now = datetime.now(timezone.utc)
+        from_time = now - timedelta(minutes=time_window_minutes)
+
+        payload = {
+            "filter": {
+                "query": query,
+                "from": from_time.isoformat(),
+                "to": now.isoformat(),
+            },
+            "sort": "-timestamp",
+            "page": {
+                "limit": 25  # Limit to the most recent 25 logs to keep context concise
+            },
+        }
+
+        try:
+            response = requests.post(
+                url, headers=self.headers, json=payload, timeout=15
+            )
+            response.raise_for_status()
+            logs = response.json().get("data", [])
+
+            if not logs:
+                return f"No logs found in Datadog for query '{query}' in the last {time_window_minutes} minutes."
+
+            summaries = [
+                f"- [{log['attributes'].get('status', 'INFO').upper()}] {log['attributes'].get('message', '')}"
+                for log in logs
+            ]
+            print(f"   ...found {len(logs)} log entries.")
+            return "\n".join(summaries)
+        except requests.exceptions.HTTPError as e:
+            return f"Error: Could not fetch logs from Datadog. HTTP {e.response.status_code}."
+        except requests.exceptions.RequestException as e:
+            return f"Error: Network issue while fetching logs from Datadog: {e}"

aira/templates/config.template.yaml

@@ -33,6 +33,12 @@ connections:
     user_email: [email protected]
     api_token: "${JSM_API_TOKEN}"
 
+  datadog_us1:
+    type: datadog
+    api_key: "${DD_API_KEY}"
+    app_key: "${DD_APP_KEY}"
+    site: DD_SITE_URL_PLACEHOLDER # Use "datadoghq.eu" for EU region
+
 # --- Actions Configuration ---
 # Define all the services the agent can take action on.
 actions:

aira/templates/non_secret_prompts.json

@@ -28,5 +28,10 @@
     "prompt": "What is the email address for your JSM user?",
     "placeholder": "[email protected]",
     "default": "[email protected]"
+  },
+  "dd_site_url": {
+    "prompt": "What is the Datadog site URL?",
+    "placeholder": "datadoghq.com or datadoghq.eu",
+    "default": "datadoghq.com"
   }
 }

aira/templates/secret_prompts.json

@@ -18,5 +18,13 @@
   "JSM_API_TOKEN": {
     "prompt": "Enter your JSM API Token",
     "url": "https://support.atlassian.com/atlassian-account/docs/manage-api-tokens-for-your-atlassian-account/"
+  },
+  "DD_API_KEY": {
+    "prompt": "Enter your Datadog API Key",
+    "url": "https://docs.datadoghq.com/account_management/api-app-keys/"
+  },
+  "DD_APP_KEY": {
+    "prompt": "Enter your Datadog Application Key",
+    "url": "https://docs.datadoghq.com/account_management/api-app-keys"
   }
 }

config.example.yaml

@@ -49,12 +49,11 @@ connections:
     user_email: [email protected]
     api_token: "${JSM_API_TOKEN}"
 
-  # --- Example for a future Datadog Connector ---
-  # datadog_us1:
-  #   type: datadog
-  #   api_key: "${DD_API_KEY}"
-  #   app_key: "${DD_APP_KEY}"
-  #   site: "datadoghq.com" # Use "datadoghq.eu" for EU region
+  datadog_us1:
+    type: datadog
+    api_key: "${DD_API_KEY}"
+    app_key: "${DD_APP_KEY}"
+    site: "datadoghq.com" # Use "datadoghq.eu" for EU region
 
 
 # --- Actions Configuration ---

tests/unit/connectors/observability/test_datadog.py

@@ -0,0 +1,81 @@
+import pytest
+from pydantic import ValidationError
+from aira.connectors.observability.datadog import DatadogConnector
+
+
+@pytest.fixture
+def valid_datadog_config() -> dict:
+    """Provides a valid configuration dictionary for the Datadog connector."""
+    return {
+        "type": "datadog",
+        "api_key": "fake_api_key",
+        "app_key": "fake_app_key",
+        "site": "datadoghq.com",
+    }
+
+
+def test_datadog_connector_instantiation_success(valid_datadog_config):
+    """Tests that the connector initializes successfully with valid config."""
+    try:
+        DatadogConnector(name="test_datadog", config=valid_datadog_config)
+    except (ValueError, ValidationError) as e:
+        pytest.fail(f"Connector instantiation failed unexpectedly: {e}")
+
+
+def test_datadog_connector_instantiation_missing_keys():
+    """Tests that instantiation fails if required keys are missing."""
+    with pytest.raises(ValidationError):
+        DatadogConnector(name="test_datadog", config={"type": "datadog"})
+
+
+def test_connection_success(requests_mock, valid_datadog_config):
+    """Tests a successful connection validation."""
+    requests_mock.get(
+        "https://api.datadoghq.com/api/v1/validate",
+        json={"valid": True},
+        status_code=200,
+    )
+    connector = DatadogConnector(name="test_datadog", config=valid_datadog_config)
+    success, message = connector.test_connection()
+    assert success is True
+    assert "connection successful and keys are valid" in message
+
+
+def test_connection_failure_invalid_keys(requests_mock, valid_datadog_config):
+    """Tests a failed connection due to invalid keys (HTTP 403)."""
+    requests_mock.get("https://api.datadoghq.com/api/v1/validate", status_code=403)
+    connector = DatadogConnector(name="test_datadog", config=valid_datadog_config)
+    success, message = connector.test_connection()
+    assert success is False
+    assert "Invalid Datadog API or App Key" in message
+
+
+def test_fetch_logs_success(requests_mock, valid_datadog_config):
+    """Tests successfully fetching and formatting logs."""
+    mock_response = {
+        "data": [
+            {"attributes": {"status": "error", "message": "Service unavailable"}},
+            {"attributes": {"status": "info", "message": "User login successful"}},
+        ]
+    }
+    requests_mock.post(
+        "https://api.datadoghq.com/api/v2/logs/events/search",
+        json=mock_response,
+        status_code=200,
+    )
+    connector = DatadogConnector(name="test_datadog", config=valid_datadog_config)
+    result = connector.fetch_logs(query="service:test", time_window_minutes=10)
+    assert "[ERROR] Service unavailable" in result
+    assert "[INFO] User login successful" in result
+
+
+def test_fetch_logs_no_logs_found(requests_mock, valid_datadog_config):
+    """Tests the response when no logs are found."""
+    requests_mock.post(
+        "https://api.datadoghq.com/api/v2/logs/events/search",
+        json={"data": []},
+        status_code=200,
+    )
+    connector = DatadogConnector(name="test_datadog", config=valid_datadog_config)
+    result = connector.fetch_logs(query="service:test", time_window_minutes=10)
+    assert "No logs found" in result