From c34e7fa2383e1efe98da9c569c98cef6a440171d Mon Sep 17 00:00:00 2001 From: mdecalf Date: Mon, 18 May 2026 20:23:59 +0200 Subject: [PATCH 1/2] feat(datadog): multi-account support for all four Datadog toolsets Many organizations run separate Datadog orgs for staging and production. Until now a single Holmes instance could only query one Datadog org per toolset, which meant maintaining two Holmes deployments or writing brittle custom toolset YAML. This change adds first-class multi-account support to all four Datadog toolsets (logs, metrics, traces, general) while keeping full backward compatibility with the existing single-account configuration. ## Config schema (datadog_api.py) - New DatadogAccount Pydantic model holds per-account credentials (name, api_key, app_key, api_url, default). api_url defaults to https://api.datadoghq.com (US1), preserving the upstream default. api_key and app_key have min_length=1 to fail fast on blank strings. - DatadogBaseConfig gains accounts: List[DatadogAccount]. A model_validator normalises the legacy single-account shorthand (api_key/app_key/api_url at top level) into a one-element list named 'default'. api_url is omitted from the DatadogAccount constructor when None so the field default applies correctly. - _ui_required_fields = ['api_key', 'app_key'] keeps the UI schema requiring credentials; _hidden_fields = ['accounts'] hides the complex nested field from form UIs. - get_account(name) resolves the per-call account name or returns the default; raises KeyError listing valid names for LLM self-correction. - get_headers(account) takes a DatadogAccount instead of the full config. ## Per-toolset changes (logs, metrics, traces, general) - Every tool gains an optional account string parameter. - _invoke resolves the account via get_account() and uses it for get_headers() and all URL constructors. - _perform_healthcheck takes a DatadogAccount; prerequisites_callable iterates all accounts and fails fast on the first unhealthy one. - _reload_instructions() passes dd_config to the Jinja context so the system prompt lists available accounts when more than one is present. ## Jinja instruction templates (all four) Guard block listing accounts only when more than one is configured. Single-account users see no change in their system prompt. ## datadog_url_utils.py - URL generators accept DatadogAccount instead of the full typed config. - generate_datadog_logs_url: use '*' not in indexes instead of indexes \!= ['*'] to treat any list containing '*' as unfiltered. ## Tests (67 pass, 16 skipped/slow) - New test_datadog_accounts.py: config parsing, validation edge cases, get_account() resolution. - Account-routing tests for all four toolsets. - Updated existing assertions for new account-scoped error messages. ## Docs Added 'Multiple Datadog accounts' section to the Datadog toolset docs. Signed-off-by: mdecalf --- docs/data-sources/builtin-toolsets/datadog.md | 33 +++ holmes/core/tools.py | 22 +- .../plugins/toolsets/datadog/datadog_api.py | 227 ++++++++++++++++-- .../datadog_general_instructions.jinja2 | 12 + .../datadog/datadog_logs_instructions.jinja2 | 12 + .../datadog_metrics_instructions.jinja2 | 12 + .../toolsets/datadog/datadog_url_utils.py | 48 ++-- .../instructions_datadog_traces.jinja2 | 12 + .../datadog/toolset_datadog_general.py | 107 ++++++--- .../toolsets/datadog/toolset_datadog_logs.py | 76 ++++-- .../datadog/toolset_datadog_metrics.py | 148 +++++++++--- .../datadog/toolset_datadog_traces.py | 98 ++++++-- .../datadog/logs/test_check_prerequisites.py | 74 +++++- .../datadog/metrics/test_datadog_metrics.py | 79 +++++- .../toolsets/datadog/test_datadog_accounts.py | 174 ++++++++++++++ .../datadog/test_toolset_datadog_general.py | 76 ++++++ .../datadog/traces/test_datadog_traces.py | 65 +++++ 17 files changed, 1117 insertions(+), 158 deletions(-) create mode 100644 tests/plugins/toolsets/datadog/test_datadog_accounts.py diff --git a/docs/data-sources/builtin-toolsets/datadog.md b/docs/data-sources/builtin-toolsets/datadog.md index 3d62e24a3c..396d08939e 100644 --- a/docs/data-sources/builtin-toolsets/datadog.md +++ b/docs/data-sources/builtin-toolsets/datadog.md @@ -186,6 +186,39 @@ holmes ask "list Datadog monitors" That's it! You're now connected to Datadog with all toolsets enabled. +## Multiple Datadog accounts + +Each Datadog toolset can be configured against several Datadog accounts +(organizations) at once. This is useful when, for example, staging and +production live in separate Datadog orgs. The LLM picks the target account +on a per-tool-call basis via an `account` parameter. + +Replace the single-account `api_key`/`app_key`/`api_url` block with an +`accounts:` list. The two forms are mutually exclusive — supply one or the +other. + + toolsets: + datadog/metrics: + enabled: true + config: + accounts: + - name: staging + api_key: "{{ env.DATADOG_API_KEY_STG }}" + app_key: "{{ env.DATADOG_APP_KEY_STG }}" + api_url: https://api.datadoghq.eu + default: true # optional — first account wins otherwise + - name: production + api_key: "{{ env.DATADOG_API_KEY_PRD }}" + app_key: "{{ env.DATADOG_APP_KEY_PRD }}" + api_url: https://api.datadoghq.eu + +When two or more accounts are configured, every tool exposed by that +toolset accepts an optional `account` parameter whose value matches one of +the configured account names. Omitting it routes the call to the account +marked `default: true` (or to the first one if no default is set). The list +of available accounts is injected into the toolset's system prompt so the +LLM can pick the right one for each call. + ## Available Toolsets HolmesGPT provides four specialized Datadog toolsets: diff --git a/holmes/core/tools.py b/holmes/core/tools.py index 57feb1788c..8c1a0ab79c 100644 --- a/holmes/core/tools.py +++ b/holmes/core/tools.py @@ -1098,22 +1098,36 @@ def get_config_schema(self) -> Optional[Dict[str, Any]]: return None return {cls.__name__: cls.build_schema_entry() for cls in self.config_classes} - def _load_llm_instructions(self, jinja_template: str): + def _load_llm_instructions( + self, jinja_template: str, extra_context: Optional[Dict[str, Any]] = None + ): tool_names = [t.name for t in self.tools] + context: Dict[str, Any] = {"tool_names": tool_names, "config": self.config} + if extra_context: + context.update(extra_context) self.llm_instructions = load_and_render_prompt( prompt=jinja_template, - context={"tool_names": tool_names, "config": self.config}, + context=context, ) - def _load_llm_instructions_from_file(self, file_dir: str, filename: str) -> None: + def _load_llm_instructions_from_file( + self, + file_dir: str, + filename: str, + extra_context: Optional[Dict[str, Any]] = None, + ) -> None: """Helper method to load LLM instructions from a jinja2 template file. Args: file_dir: Directory where the template file is located (typically os.path.dirname(__file__)) filename: Name of the jinja2 template file (e.g., "toolset_grafana_dashboard.jinja2") + extra_context: Optional additional variables merged into the Jinja context. """ template_file_path = os.path.abspath(os.path.join(file_dir, filename)) - self._load_llm_instructions(jinja_template=f"file://{template_file_path}") + self._load_llm_instructions( + jinja_template=f"file://{template_file_path}", + extra_context=extra_context, + ) class YAMLToolset(Toolset): diff --git a/holmes/plugins/toolsets/datadog/datadog_api.py b/holmes/plugins/toolsets/datadog/datadog_api.py index bd78a51c14..ee1dbba71b 100644 --- a/holmes/plugins/toolsets/datadog/datadog_api.py +++ b/holmes/plugins/toolsets/datadog/datadog_api.py @@ -3,11 +3,11 @@ import re import threading from datetime import datetime, timedelta, timezone -from typing import Any, ClassVar, Dict, Optional, Tuple, Union +from typing import Any, ClassVar, Dict, List, Optional, Tuple, Union from urllib.parse import urlparse, urlunparse import requests # type: ignore -from pydantic import AnyUrl, Field +from pydantic import AnyUrl, BaseModel, ConfigDict, Field, model_validator from requests.structures import CaseInsensitiveDict # type: ignore from tenacity import retry, retry_if_exception, stop_after_attempt, wait_incrementing from tenacity.wait import wait_base @@ -95,8 +95,78 @@ def convert_api_url_to_app_url(api_url: Union[str, AnyUrl]) -> str: return app_url +class DatadogAccount(BaseModel): + """A single Datadog organization (account) the toolset can query. + + Each account carries its own credentials and site URL so a single + Holmes deployment can target several Datadog organizations (e.g. + staging + production) from the same toolset. + """ + + model_config = ConfigDict(extra="forbid") + + name: str = Field( + min_length=1, + title="Account name", + description=( + "Identifier used by the LLM to target this account from a tool " + "call (for example `staging` or `production`). Must be unique " + "within a toolset." + ), + examples=["default", "staging", "production"], + ) + api_key: str = Field( + min_length=1, + title="API Key", + description="Datadog API key for authentication", + json_schema_extra={"format": "password"}, + ) + app_key: str = Field( + min_length=1, + title="Application Key", + description="Datadog application key for authentication", + json_schema_extra={"format": "password"}, + ) + api_url: AnyUrl = Field( + default="https://api.datadoghq.com", + title="API URL", + description=( + "Datadog site API base URL for this account. Defaults to the US1 " + "site (https://api.datadoghq.com). See " + "https://docs.datadoghq.com/getting_started/site/ for the full " + "list." + ), + examples=[ + "https://api.datadoghq.com", + "https://api.datadoghq.eu", + ], + ) + default: bool = Field( + default=False, + title="Default", + description=( + "When true, this account is used by tool calls that don't set " + "an explicit `account` parameter. At most one account may be " + "flagged default; if none is, the first account in the list " + "wins." + ), + ) + + class DatadogBaseConfig(ToolsetConfig): - """Base configuration for all Datadog toolsets""" + """Base configuration for all Datadog toolsets. + + Supports two interchangeable schemas: + + - **Single-account shorthand** (legacy / common case): set `api_key`, + `app_key` and `api_url` directly on the config. Internally normalised + to a single :class:`DatadogAccount` named ``default``. + - **Multi-account**: populate the `accounts:` list to query several + Datadog organizations from the same toolset. Each tool call accepts + an optional `account` parameter to pick the target. + + The two forms are mutually exclusive within one config block. + """ _deprecated_mappings: ClassVar[Dict[str, Optional[str]]] = { "dd_api_key": "api_key", @@ -104,24 +174,45 @@ class DatadogBaseConfig(ToolsetConfig): "site_api_url": "api_url", "request_timeout": "timeout_seconds", } - - api_key: str = Field( + # The UI form guides users through the common single-account path, so mark + # the shorthand credential fields as required in the exported schema even + # though they are Optional at the Pydantic level (to allow the + # `accounts:` alternative). The `accounts` list is hidden because nested + # complex objects don't render well in form UIs — advanced users supply it + # directly in their YAML/Helm values. + _ui_required_fields: ClassVar[List[str]] = ["api_key", "app_key"] + _hidden_fields: ClassVar[List[str]] = ["accounts"] + + # Single-account shorthand. When `accounts:` below is populated, these + # fields must be omitted. + api_key: Optional[str] = Field( + default=None, title="API Key", - description="Datadog API key for authentication", + description=( + "Datadog API key. Single-account shorthand; for multi-account " + "setups, populate the `accounts:` list instead." + ), examples=["{{ env.DATADOG_API_KEY }}"], json_schema_extra={"format": "password"}, ) - app_key: str = Field( + app_key: Optional[str] = Field( + default=None, title="Application Key", - description="Datadog application key for authentication", + description=( + "Datadog application key. Single-account shorthand; for " + "multi-account setups, populate `accounts:` instead." + ), examples=["{{ env.DATADOG_APP_KEY }}"], json_schema_extra={"format": "password"}, ) - api_url: AnyUrl = Field( + api_url: Optional[AnyUrl] = Field( + default=None, title="API URL", description=( - "Datadog site API base URL. Pick the URL matching your Datadog region — " - "see https://docs.datadoghq.com/getting_started/site/ for the full list." + "Datadog site API base URL. Single-account shorthand; for " + "multi-account setups, populate `accounts:` instead. See " + "https://docs.datadoghq.com/getting_started/site/ for the " + "full list." ), examples=[ "https://api.datadoghq.com", @@ -132,12 +223,104 @@ class DatadogBaseConfig(ToolsetConfig): "https://api.ddog-gov.com", ], ) + accounts: List[DatadogAccount] = Field( + default_factory=list, + title="Datadog accounts", + description=( + "Datadog accounts (organizations) the toolset can query. Leave " + "empty to use the single-account shorthand above; populate with " + "two or more entries to enable multi-account routing via the " + "`account` tool parameter." + ), + ) timeout_seconds: int = Field( default=60, title="Timeout", description="HTTP request timeout in seconds", ) + @model_validator(mode="after") + def _normalize_accounts(self) -> "DatadogBaseConfig": + """Coalesce the legacy single-account fields into ``self.accounts``.""" + legacy_set = any( + v is not None for v in (self.api_key, self.app_key, self.api_url) + ) + if self.accounts and legacy_set: + raise ValueError( + "Datadog config: provide either the top-level " + "api_key/app_key/api_url shorthand OR the `accounts:` list, " + "not both." + ) + if not self.accounts: + if not legacy_set: + raise ValueError( + "Datadog config: missing credentials. Provide either " + "top-level api_key/app_key/api_url or an `accounts:` " + "list." + ) + missing = [ + name + for name, value in ( + ("api_key", self.api_key), + ("app_key", self.app_key), + ) + if value is None + ] + if missing: + raise ValueError( + f"Datadog config: missing required field(s) {missing}. " + "api_key and app_key must be set when using the " + "single-account shorthand." + ) + account_kwargs: Dict[str, Any] = { + "name": "default", + "api_key": self.api_key, + "app_key": self.app_key, + "default": True, + } + if self.api_url is not None: + account_kwargs["api_url"] = self.api_url + self.accounts = [DatadogAccount(**account_kwargs)] + # `self.accounts` is now the source of truth. We keep the legacy + # top-level fields populated (when the shorthand was used) so that + # downstream code reading them stays backward-compatible — but new + # call sites should always read from `accounts` via `get_account()`. + names = [account.name for account in self.accounts] + duplicates = sorted({n for n in names if names.count(n) > 1}) + if duplicates: + raise ValueError( + f"Datadog config: duplicate account name(s): {duplicates}." + ) + defaults = [account for account in self.accounts if account.default] + if len(defaults) > 1: + raise ValueError( + "Datadog config: at most one account may be flagged " + f"`default: true` (got {[a.name for a in defaults]})." + ) + if not defaults: + self.accounts[0].default = True + return self + + def get_account(self, name: Optional[str] = None) -> DatadogAccount: + """Resolve the named Datadog account, or the default when omitted. + + Raises :class:`KeyError` with the list of valid account names when + ``name`` is unknown so the LLM can self-correct on retry. + """ + if name is None: + for account in self.accounts: + if account.default: + return account + return self.accounts[0] + for account in self.accounts: + if account.name == name: + return account + available = ", ".join(repr(account.name) for account in self.accounts) + raise KeyError( + f"unknown Datadog account {name!r}; available accounts: " + f"[{available}]" + ) + class DataDogRequestError(Exception): payload: dict @@ -159,19 +342,29 @@ def __init__( self.response_headers = response_headers -def get_headers(dd_config: DatadogBaseConfig) -> Dict[str, str]: - """Get standard headers for Datadog API requests. +#: Shared description for the per-tool ``account`` parameter exposed by every +#: Datadog tool. The exact list of available accounts is injected into the +#: toolset's system instructions via Jinja so the LLM knows what's valid. +ACCOUNT_TOOL_PARAM_DESCRIPTION = ( + "Datadog account to query. Optional. If omitted, the configured default " + "account is used. The list of available accounts is in the toolset's " + "system instructions." +) + + +def get_headers(account: DatadogAccount) -> Dict[str, str]: + """Get standard headers for Datadog API requests for a given account. Args: - dd_config: Datadog configuration object + account: The Datadog account whose credentials should be used. Returns: - Dictionary of headers for Datadog API requests + Dictionary of headers for Datadog API requests. """ return { "Content-Type": "application/json", - "DD-API-KEY": dd_config.api_key, - "DD-APPLICATION-KEY": dd_config.app_key, + "DD-API-KEY": account.api_key, + "DD-APPLICATION-KEY": account.app_key, } diff --git a/holmes/plugins/toolsets/datadog/datadog_general_instructions.jinja2 b/holmes/plugins/toolsets/datadog/datadog_general_instructions.jinja2 index 62cfe33bde..82debd965c 100644 --- a/holmes/plugins/toolsets/datadog/datadog_general_instructions.jinja2 +++ b/holmes/plugins/toolsets/datadog/datadog_general_instructions.jinja2 @@ -1,3 +1,15 @@ +{%- if dd_config and dd_config.accounts and dd_config.accounts | length > 1 %} +## Datadog accounts + +This toolset is configured against multiple Datadog accounts: +{%- for account in dd_config.accounts %} +- `{{ account.name }}`{% if account.default %} (default){% endif %} — {{ account.api_url }} +{%- endfor %} + +Pass the `account` parameter on each tool call to target a specific account. +Omit it to use the default ({{ (dd_config.accounts | selectattr('default') | first).name }}). + +{% endif -%} ## Datadog General API Tools Usage Guide ### When to Use This Toolset diff --git a/holmes/plugins/toolsets/datadog/datadog_logs_instructions.jinja2 b/holmes/plugins/toolsets/datadog/datadog_logs_instructions.jinja2 index 861bf90c4f..adc58d1286 100644 --- a/holmes/plugins/toolsets/datadog/datadog_logs_instructions.jinja2 +++ b/holmes/plugins/toolsets/datadog/datadog_logs_instructions.jinja2 @@ -1,3 +1,15 @@ +{%- if dd_config and dd_config.accounts and dd_config.accounts | length > 1 %} +## Datadog accounts + +This toolset is configured against multiple Datadog accounts: +{%- for account in dd_config.accounts %} +- `{{ account.name }}`{% if account.default %} (default){% endif %} — {{ account.api_url }} +{%- endfor %} + +Pass the `account` parameter on each tool call to target a specific account. +Omit it to use the default ({{ (dd_config.accounts | selectattr('default') | first).name }}). + +{% endif -%} ## Datadog Logs Tools Usage Guide Before running logs queries: diff --git a/holmes/plugins/toolsets/datadog/datadog_metrics_instructions.jinja2 b/holmes/plugins/toolsets/datadog/datadog_metrics_instructions.jinja2 index 2a96ba56f9..b5fc5a7b02 100644 --- a/holmes/plugins/toolsets/datadog/datadog_metrics_instructions.jinja2 +++ b/holmes/plugins/toolsets/datadog/datadog_metrics_instructions.jinja2 @@ -1,3 +1,15 @@ +{%- if dd_config and dd_config.accounts and dd_config.accounts | length > 1 %} +## Datadog accounts + +This toolset is configured against multiple Datadog accounts: +{%- for account in dd_config.accounts %} +- `{{ account.name }}`{% if account.default %} (default){% endif %} — {{ account.api_url }} +{%- endfor %} + +Pass the `account` parameter on each tool call to target a specific account. +Omit it to use the default ({{ (dd_config.accounts | selectattr('default') | first).name }}). + +{% endif -%} ## Datadog Metrics Tools Usage Guide Before running metrics queries: diff --git a/holmes/plugins/toolsets/datadog/datadog_url_utils.py b/holmes/plugins/toolsets/datadog/datadog_url_utils.py index e1a5b03013..c0fc65edaa 100644 --- a/holmes/plugins/toolsets/datadog/datadog_url_utils.py +++ b/holmes/plugins/toolsets/datadog/datadog_url_utils.py @@ -1,23 +1,20 @@ import re -from typing import Any, Dict, Optional +from typing import Any, Dict, List, Optional from urllib.parse import urlencode, urlparse -from holmes.plugins.toolsets.datadog.datadog_api import convert_api_url_to_app_url -from holmes.plugins.toolsets.datadog.datadog_models import ( - DatadogGeneralConfig, - DatadogLogsConfig, - DatadogMetricsConfig, - DatadogTracesConfig, +from holmes.plugins.toolsets.datadog.datadog_api import ( + DatadogAccount, + convert_api_url_to_app_url, ) def generate_datadog_metrics_explorer_url( - dd_config: DatadogMetricsConfig, + account: DatadogAccount, query: str, from_time: int, to_time: int, ) -> str: - base_url = convert_api_url_to_app_url(dd_config.api_url) + base_url = convert_api_url_to_app_url(account.api_url) params = { "query": query, @@ -30,13 +27,13 @@ def generate_datadog_metrics_explorer_url( def generate_datadog_metrics_list_url( - dd_config: DatadogMetricsConfig, + account: DatadogAccount, from_time: int, host: Optional[str] = None, tag_filter: Optional[str] = None, metric_filter: Optional[str] = None, ) -> str: - base_url = convert_api_url_to_app_url(dd_config.api_url) + base_url = convert_api_url_to_app_url(account.api_url) params = {} if metric_filter: @@ -52,30 +49,30 @@ def generate_datadog_metrics_list_url( def generate_datadog_metric_metadata_url( - dd_config: DatadogMetricsConfig, + account: DatadogAccount, metric_name: str, ) -> str: - base_url = convert_api_url_to_app_url(dd_config.api_url) + base_url = convert_api_url_to_app_url(account.api_url) params = {"metric": metric_name} return f"{base_url}/metric/summary?{urlencode(params)}" def generate_datadog_metric_tags_url( - dd_config: DatadogMetricsConfig, + account: DatadogAccount, metric_name: str, ) -> str: - base_url = convert_api_url_to_app_url(dd_config.api_url) + base_url = convert_api_url_to_app_url(account.api_url) params = {"metric": metric_name} return f"{base_url}/metric/summary?{urlencode(params)}" def generate_datadog_spans_url( - dd_config: DatadogTracesConfig, + account: DatadogAccount, query: str, from_time_ms: int, to_time_ms: int, ) -> str: - base_url = convert_api_url_to_app_url(dd_config.api_url) + base_url = convert_api_url_to_app_url(account.api_url) url_params = { "query": query, @@ -88,12 +85,12 @@ def generate_datadog_spans_url( def generate_datadog_spans_analytics_url( - dd_config: DatadogTracesConfig, + account: DatadogAccount, query: str, from_time_ms: int, to_time_ms: int, ) -> str: - base_url = convert_api_url_to_app_url(dd_config.api_url) + base_url = convert_api_url_to_app_url(account.api_url) url_params = { "query": query, @@ -106,10 +103,11 @@ def generate_datadog_spans_analytics_url( def generate_datadog_logs_url( - dd_config: DatadogLogsConfig, + account: DatadogAccount, + indexes: List[str], params: dict, ) -> str: - base_url = convert_api_url_to_app_url(dd_config.api_url) + base_url = convert_api_url_to_app_url(account.api_url) url_params = { "query": params["filter"]["query"], "from_ts": params["filter"]["from"], @@ -118,8 +116,8 @@ def generate_datadog_logs_url( "storage": params["filter"]["storage_tier"], } - if dd_config.indexes != ["*"]: - url_params["index"] = ",".join(dd_config.indexes) + if indexes and "*" not in indexes: + url_params["index"] = ",".join(indexes) # Construct the full URL return f"{base_url}/logs?{urlencode(url_params)}" @@ -157,11 +155,11 @@ def _build_qs( def generate_datadog_general_url( - dd_config: DatadogGeneralConfig, + account: DatadogAccount, endpoint: str, query_params: Optional[Dict[str, Any]] = None, ) -> Optional[str]: - base_url = convert_api_url_to_app_url(dd_config.api_url) + base_url = convert_api_url_to_app_url(account.api_url) path = urlparse(endpoint).path if "/logs" in path: diff --git a/holmes/plugins/toolsets/datadog/instructions_datadog_traces.jinja2 b/holmes/plugins/toolsets/datadog/instructions_datadog_traces.jinja2 index b80df8fe7a..a51850c291 100644 --- a/holmes/plugins/toolsets/datadog/instructions_datadog_traces.jinja2 +++ b/holmes/plugins/toolsets/datadog/instructions_datadog_traces.jinja2 @@ -1,3 +1,15 @@ +{%- if dd_config and dd_config.accounts and dd_config.accounts | length > 1 %} +## Datadog accounts + +This toolset is configured against multiple Datadog accounts: +{%- for account in dd_config.accounts %} +- `{{ account.name }}`{% if account.default %} (default){% endif %} — {{ account.api_url }} +{%- endfor %} + +Pass the `account` parameter on each tool call to target a specific account. +Omit it to use the default ({{ (dd_config.accounts | selectattr('default') | first).name }}). + +{% endif -%} ## Datadog Traces Toolset Tools to search and analyze distributed traces from Datadog APM. diff --git a/holmes/plugins/toolsets/datadog/toolset_datadog_general.py b/holmes/plugins/toolsets/datadog/toolset_datadog_general.py index 1785895b79..eb9483bd4d 100644 --- a/holmes/plugins/toolsets/datadog/toolset_datadog_general.py +++ b/holmes/plugins/toolsets/datadog/toolset_datadog_general.py @@ -21,7 +21,9 @@ ) from holmes.plugins.toolsets.consts import TOOLSET_CONFIG_MISSING_ERROR from holmes.plugins.toolsets.datadog.datadog_api import ( + ACCOUNT_TOOL_PARAM_DESCRIPTION, MAX_RETRY_COUNT_ON_RATE_LIMIT, + DatadogAccount, DataDogRequestError, enhance_error_message, execute_datadog_http_request, @@ -225,12 +227,7 @@ def __init__(self): ], tags=[ToolsetTag.CORE], ) - template_file_path = os.path.abspath( - os.path.join( - os.path.dirname(__file__), "datadog_general_instructions.jinja2" - ) - ) - self._load_llm_instructions(jinja_template=f"file://{template_file_path}") + self._reload_instructions() def prerequisites_callable(self, config: dict[str, Any]) -> Tuple[bool, str]: """Check prerequisites with configuration.""" @@ -242,7 +239,6 @@ def prerequisites_callable(self, config: dict[str, Any]) -> Tuple[bool, str]: try: dd_config = DatadogGeneralConfig(**config) - self.dd_config = dd_config # Fetch OpenAPI spec on startup for better error messages and documentation logging.debug("Fetching Datadog OpenAPI specification...") @@ -256,19 +252,29 @@ def prerequisites_callable(self, config: dict[str, Any]) -> Tuple[bool, str]: "Could not fetch OpenAPI spec; enhanced error messages will be limited" ) - success, error_msg = self._perform_healthcheck(dd_config) - return success, error_msg + for account in dd_config.accounts: + success, error_msg = self._perform_healthcheck(account, dd_config) + if not success: + return False, error_msg + self.dd_config = dd_config + self._reload_instructions() + return True, "" except Exception as e: logging.exception("Failed to set up Datadog general toolset") return False, f"Invalid Datadog General configuration: {e}" - def _perform_healthcheck(self, dd_config: DatadogGeneralConfig) -> Tuple[bool, str]: - """Perform health check on Datadog API.""" + def _perform_healthcheck( + self, account: DatadogAccount, dd_config: DatadogGeneralConfig + ) -> Tuple[bool, str]: + """Perform health check on Datadog API for one account.""" try: - logging.info("Performing Datadog general API configuration healthcheck...") - base_url = str(dd_config.api_url).rstrip("/") + logging.info( + "Performing Datadog general API healthcheck for account %r...", + account.name, + ) + base_url = str(account.api_url).rstrip("/") url = f"{base_url}/api/v1/validate" - headers = get_headers(dd_config) + headers = get_headers(account) data = execute_datadog_http_request( url=url, @@ -279,16 +285,33 @@ def _perform_healthcheck(self, dd_config: DatadogGeneralConfig) -> Tuple[bool, s ) if data.get("valid", False): - logging.debug("Datadog general API health check completed successfully") + logging.debug( + "Datadog general API healthcheck succeeded for account %r", + account.name, + ) return True, "" - else: - error_msg = "Datadog API key validation failed" - logging.error(f"Datadog General health check failed: {error_msg}") - return False, f"Datadog General health check failed: {error_msg}" + error_msg = ( + f"Datadog General healthcheck failed for account {account.name!r}: " + "API key validation returned valid=false." + ) + logging.error(error_msg) + return False, error_msg except Exception as e: - logging.exception("Failed during Datadog general API health check") - return False, f"Datadog General health check failed: {e}" + logging.exception( + "Datadog general API healthcheck failed for account %r", account.name + ) + return False, ( + f"Datadog General healthcheck failed for account " + f"{account.name!r}: {e}" + ) + + def _reload_instructions(self): + self._load_llm_instructions_from_file( + os.path.dirname(__file__), + "datadog_general_instructions.jinja2", + extra_context={"dd_config": self.dd_config}, + ) def is_endpoint_allowed( @@ -385,6 +408,11 @@ def __init__(self, toolset: "DatadogGeneralToolset"): type="string", required=True, ), + "account": ToolParameter( + description=ACCOUNT_TOOL_PARAM_DESCRIPTION, + type="string", + required=False, + ), }, toolset=toolset, ) @@ -413,6 +441,14 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes error=TOOLSET_CONFIG_MISSING_ERROR, params=params_return, ) + try: + account = self.toolset.dd_config.get_account(params.get("account")) + except KeyError as exc: + return StructuredToolResult( + status=StructuredToolResultStatus.ERROR, + error=str(exc), + params=params_return, + ) endpoint = params.get("endpoint", "") query_params = params.get("query_params", {}) @@ -434,10 +470,10 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes url = None try: # Build full URL (ensure no double slashes) - base_url = str(self.toolset.dd_config.api_url).rstrip("/") + base_url = str(account.api_url).rstrip("/") endpoint = endpoint.lstrip("/") url = f"{base_url}/{endpoint}" - headers = get_headers(self.toolset.dd_config) + headers = get_headers(account) logging.info(f"Full API URL: {url}") @@ -466,7 +502,7 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes ) web_url = generate_datadog_general_url( - self.toolset.dd_config, + account, endpoint, query_params, ) @@ -492,7 +528,7 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes elif e.status_code == 400: # Use enhanced error message for 400 errors error_msg = enhance_error_message( - e, endpoint, "GET", str(self.toolset.dd_config.api_url) + e, endpoint, "GET", str(account.api_url) ) else: error_msg = f"API error {e.status_code}: {str(e)}" @@ -560,6 +596,11 @@ def __init__(self, toolset: "DatadogGeneralToolset"): type="string", required=True, ), + "account": ToolParameter( + description=ACCOUNT_TOOL_PARAM_DESCRIPTION, + type="string", + required=False, + ), }, toolset=toolset, ) @@ -609,6 +650,14 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes error=TOOLSET_CONFIG_MISSING_ERROR, params=params, ) + try: + account = self.toolset.dd_config.get_account(params.get("account")) + except KeyError as exc: + return StructuredToolResult( + status=StructuredToolResultStatus.ERROR, + error=str(exc), + params=params, + ) endpoint = params.get("endpoint", "") body = params.get("body", {}) @@ -630,10 +679,10 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes url = None try: # Build full URL (ensure no double slashes) - base_url = str(self.toolset.dd_config.api_url).rstrip("/") + base_url = str(account.api_url).rstrip("/") endpoint = endpoint.lstrip("/") url = f"{base_url}/{endpoint}" - headers = get_headers(self.toolset.dd_config) + headers = get_headers(account) logging.info(f"Full API URL: {url}") @@ -663,7 +712,7 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes body_query_params = self._body_to_query_params(body) web_url = generate_datadog_general_url( - self.toolset.dd_config, + account, endpoint, body_query_params, ) @@ -689,7 +738,7 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes elif e.status_code == 400: # Use enhanced error message for 400 errors error_msg = enhance_error_message( - e, endpoint, "POST", str(self.toolset.dd_config.api_url) + e, endpoint, "POST", str(account.api_url) ) else: error_msg = f"API error {e.status_code}: {str(e)}" diff --git a/holmes/plugins/toolsets/datadog/toolset_datadog_logs.py b/holmes/plugins/toolsets/datadog/toolset_datadog_logs.py index 724fd1c358..441e093b34 100644 --- a/holmes/plugins/toolsets/datadog/toolset_datadog_logs.py +++ b/holmes/plugins/toolsets/datadog/toolset_datadog_logs.py @@ -16,7 +16,9 @@ ) from holmes.plugins.toolsets.consts import STANDARD_END_DATETIME_TOOL_PARAM_DESCRIPTION from holmes.plugins.toolsets.datadog.datadog_api import ( + ACCOUNT_TOOL_PARAM_DESCRIPTION, MAX_RETRY_COUNT_ON_RATE_LIMIT, + DatadogAccount, DataDogRequestError, execute_datadog_http_request, get_headers, @@ -79,13 +81,15 @@ def __init__(self): self.tools = [GetLogs(toolset=self)] self._reload_instructions() - def _perform_healthcheck(self) -> Tuple[bool, str]: - """Perform health check on Datadog logs API.""" + def _perform_healthcheck(self, account: DatadogAccount) -> Tuple[bool, str]: + """Perform health check on Datadog logs API for one account.""" if not self.dd_config: return False, "Internal error: Datadog configuration not initialized" try: - logging.info("Performing Datadog logs configuration healthcheck...") - headers = get_headers(self.dd_config) + logging.info( + "Performing Datadog logs healthcheck for account %r...", account.name + ) + headers = get_headers(account) payload = { "filter": { "from": "now-1m", @@ -96,7 +100,7 @@ def _perform_healthcheck(self) -> Tuple[bool, str]: "page": {"limit": 1}, } - search_url = f"{self.dd_config.api_url}/api/v2/logs/events/search" + search_url = f"{account.api_url}/api/v2/logs/events/search" execute_datadog_http_request( url=search_url, headers=headers, @@ -109,18 +113,30 @@ def _perform_healthcheck(self) -> Tuple[bool, str]: except DataDogRequestError as e: logging.error( - f"Datadog API error during healthcheck: {e.status_code} - {e.response_text}" + "Datadog logs healthcheck failed for account %r: %s - %s", + account.name, + e.status_code, + e.response_text, ) if e.status_code == 403: return ( False, - "API key lacks required permissions. Make sure your API key has 'apm_read' scope.", + f"Datadog Logs healthcheck failed for account " + f"{account.name!r}: API key lacks required permissions. " + "Make sure your API key has 'apm_read' scope.", ) - else: - return False, f"Datadog API error: {e.status_code} - {e.response_text}" + return ( + False, + f"Datadog Logs healthcheck failed for account " + f"{account.name!r}: API error {e.status_code} - {e.response_text}", + ) except Exception as e: - logging.exception("Failed during Datadog logs health check") - return False, f"Datadog Logs health check failed: {e}" + logging.exception( + "Datadog logs healthcheck failed for account %r", account.name + ) + return False, ( + f"Datadog Logs healthcheck failed for account {account.name!r}: {e}" + ) def prerequisites_callable(self, config: dict[str, Any]) -> Tuple[bool, str]: if not config: @@ -133,19 +149,23 @@ def prerequisites_callable(self, config: dict[str, Any]) -> Tuple[bool, str]: dd_config = DatadogLogsConfig(**config) self.dd_config = dd_config - success, error_msg = self._perform_healthcheck() - return success, error_msg + for account in dd_config.accounts: + success, error_msg = self._perform_healthcheck(account) + if not success: + return False, error_msg + self._reload_instructions() + return True, "" except Exception as e: logging.exception("Failed to set up Datadog Logs toolset") return (False, f"Invalid Datadog Logs configuration: {e}") def _reload_instructions(self): - """Load Datadog logs specific troubleshooting instructions.""" - template_file_path = os.path.abspath( - os.path.join(os.path.dirname(__file__), "datadog_logs_instructions.jinja2") + self._load_llm_instructions_from_file( + os.path.dirname(__file__), + "datadog_logs_instructions.jinja2", + extra_context={"dd_config": self.dd_config}, ) - self._load_llm_instructions(jinja_template=f"file://{template_file_path}") class GetLogs(Tool): @@ -188,6 +208,11 @@ class GetLogs(Tool): type="boolean", required=False, ), + "account": ToolParameter( + description=ACCOUNT_TOOL_PARAM_DESCRIPTION, + type="string", + required=False, + ), } def get_parameterized_one_liner(self, params: dict) -> str: @@ -202,6 +227,15 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes error="Datadog configuration not initialized", params=params, ) + try: + account = self.toolset.dd_config.get_account(params.get("account")) + except KeyError as exc: + return StructuredToolResult( + status=StructuredToolResultStatus.ERROR, + error=str(exc), + params=params, + ) + url = None payload: Optional[Dict[str, Any]] = None try: @@ -221,8 +255,8 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes params["limit"] = limit sort = "timestamp" if params.get("sort_desc", False) else "-timestamp" - url = f"{self.toolset.dd_config.api_url}/api/v2/logs/events/search" - headers = get_headers(self.toolset.dd_config) + url = f"{account.api_url}/api/v2/logs/events/search" + headers = get_headers(account) storage = self.toolset.dd_config.storage_tier payload = { @@ -257,7 +291,9 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes status=StructuredToolResultStatus.SUCCESS, data=response, params=params, - url=generate_datadog_logs_url(self.toolset.dd_config, payload), + url=generate_datadog_logs_url( + account, self.toolset.dd_config.indexes, payload + ), ) except DataDogRequestError as e: diff --git a/holmes/plugins/toolsets/datadog/toolset_datadog_metrics.py b/holmes/plugins/toolsets/datadog/toolset_datadog_metrics.py index 059d7a2270..58bc16c712 100644 --- a/holmes/plugins/toolsets/datadog/toolset_datadog_metrics.py +++ b/holmes/plugins/toolsets/datadog/toolset_datadog_metrics.py @@ -23,7 +23,9 @@ TOOLSET_CONFIG_MISSING_ERROR, ) from holmes.plugins.toolsets.datadog.datadog_api import ( + ACCOUNT_TOOL_PARAM_DESCRIPTION, MAX_RETRY_COUNT_ON_RATE_LIMIT, + DatadogAccount, DataDogRequestError, execute_datadog_http_request, get_headers, @@ -49,6 +51,24 @@ class BaseDatadogMetricsTool(Tool): toolset: "DatadogMetricsToolset" + def _resolve_account( + self, params: dict + ) -> Tuple[Optional[DatadogAccount], Optional[StructuredToolResult]]: + """Resolve the target Datadog account from the ``account`` parameter. + + Returns either (account, None) on success or (None, error_result) + when the account is unknown so the caller can short-circuit. + """ + assert self.toolset.dd_config is not None + try: + return self.toolset.dd_config.get_account(params.get("account")), None + except KeyError as exc: + return None, StructuredToolResult( + status=StructuredToolResultStatus.ERROR, + error=str(exc), + params=params, + ) + ACTIVE_METRICS_DEFAULT_LOOK_BACK_HOURS = 24 ACTIVE_METRICS_DEFAULT_TIME_SPAN_SECONDS = 24 * 60 * 60 @@ -86,6 +106,11 @@ def __init__(self, toolset: "DatadogMetricsToolset"): type="integer", required=False, ), + "account": ToolParameter( + description=ACCOUNT_TOOL_PARAM_DESCRIPTION, + type="string", + required=False, + ), }, toolset=toolset, ) @@ -98,6 +123,11 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes params=params, ) + account, error = self._resolve_account(params) + if error is not None: + return error + assert account is not None + url = None query_params = None @@ -110,8 +140,8 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes default_time_span_seconds=ACTIVE_METRICS_DEFAULT_TIME_SPAN_SECONDS, ) - url = f"{self.toolset.dd_config.api_url}/api/v1/metrics" - headers = get_headers(self.toolset.dd_config) + url = f"{account.api_url}/api/v1/metrics" + headers = get_headers(account) query_params = { "from": from_time, @@ -182,7 +212,7 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes ) url = generate_datadog_metrics_list_url( - self.toolset.dd_config, + account, from_time, params.get("host"), params.get("tag_filter"), @@ -293,6 +323,11 @@ def __init__(self, toolset: "DatadogMetricsToolset"): type="string", required=False, ), + "account": ToolParameter( + description=ACCOUNT_TOOL_PARAM_DESCRIPTION, + type="string", + required=False, + ), }, toolset=toolset, ) @@ -305,6 +340,11 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes params=params, ) + account, error = self._resolve_account(params) + if error is not None: + return error + assert account is not None + url = None query_params = None @@ -317,8 +357,8 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes default_time_span_seconds=DEFAULT_TIME_SPAN_SECONDS, ) - url = f"{self.toolset.dd_config.api_url}/api/v1/query" - headers = get_headers(self.toolset.dd_config) + url = f"{account.api_url}/api/v1/query" + headers = get_headers(account) query_params = { "query": query, @@ -413,7 +453,7 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes } url = generate_datadog_metrics_explorer_url( - self.toolset.dd_config, + account, query, from_time, to_time, @@ -493,6 +533,11 @@ def __init__(self, toolset: "DatadogMetricsToolset"): type="string", required=True, ), + "account": ToolParameter( + description=ACCOUNT_TOOL_PARAM_DESCRIPTION, + type="string", + required=False, + ), }, toolset=toolset, ) @@ -505,6 +550,11 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes params=params, ) + account, error = self._resolve_account(params) + if error is not None: + return error + assert account is not None + try: metric_names_str = get_param_or_raise(params, "metric_names") @@ -521,14 +571,14 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes params=params, ) - headers = get_headers(self.toolset.dd_config) + headers = get_headers(account) results = {} errors = {} for metric_name in metric_names: try: - api_url = f"{self.toolset.dd_config.api_url}/api/v1/metrics/{metric_name}" + api_url = f"{account.api_url}/api/v1/metrics/{metric_name}" data = execute_datadog_http_request( url=api_url, @@ -563,7 +613,7 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes # Generate URL for the first metric (or a general metrics page if multiple) if metric_names: url = generate_datadog_metric_metadata_url( - self.toolset.dd_config, + account, metric_names[0], ) else: @@ -618,6 +668,11 @@ def __init__(self, toolset: "DatadogMetricsToolset"): type="string", required=True, ), + "account": ToolParameter( + description=ACCOUNT_TOOL_PARAM_DESCRIPTION, + type="string", + required=False, + ), }, toolset=toolset, ) @@ -630,25 +685,30 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes params=params, ) + account, error = self._resolve_account(params) + if error is not None: + return error + assert account is not None + api_url = None - query_params = None + query_params: dict = {} try: metric_name = get_param_or_raise(params, "metric_name") - api_url = f"{self.toolset.dd_config.api_url}/api/v2/metrics/{metric_name}/active-configurations" - headers = get_headers(self.toolset.dd_config) + api_url = f"{account.api_url}/api/v2/metrics/{metric_name}/active-configurations" + headers = get_headers(account) data = execute_datadog_http_request( url=api_url, headers=headers, timeout=self.toolset.dd_config.timeout_seconds, method="GET", - payload_or_params={}, + payload_or_params=query_params, ) web_url = generate_datadog_metric_tags_url( - self.toolset.dd_config, + account, metric_name, ) @@ -686,7 +746,7 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes error=error_msg, params=params, invocation=json.dumps({"url": api_url, "params": query_params}) - if api_url and query_params + if api_url is not None else None, ) @@ -728,12 +788,17 @@ def __init__(self): ) self._reload_instructions() - def _perform_healthcheck(self, dd_config: DatadogMetricsConfig) -> Tuple[bool, str]: + def _perform_healthcheck( + self, account: DatadogAccount, dd_config: DatadogMetricsConfig + ) -> Tuple[bool, str]: try: - logging.debug("Performing Datadog metrics configuration healthcheck...") + logging.debug( + "Performing Datadog metrics healthcheck for account %r...", + account.name, + ) - url = f"{dd_config.api_url}/api/v1/validate" - headers = get_headers(dd_config) + url = f"{account.api_url}/api/v1/validate" + headers = get_headers(account) data = execute_datadog_http_request( url=url, @@ -744,16 +809,26 @@ def _perform_healthcheck(self, dd_config: DatadogMetricsConfig) -> Tuple[bool, s ) if data.get("valid", False): - logging.info("Datadog metrics health check completed successfully") + logging.info( + "Datadog metrics healthcheck succeeded for account %r", + account.name, + ) return True, "" - else: - error_msg = "Datadog API key validation failed" - logging.error(f"Datadog Metrics health check failed: {error_msg}") - return False, f"Datadog Metrics health check failed: {error_msg}" + error_msg = ( + f"Datadog Metrics healthcheck failed for account {account.name!r}: " + "API key validation returned valid=false." + ) + logging.error(error_msg) + return False, error_msg except Exception as e: - logging.exception("Failed during Datadog metrics health check") - return False, f"Datadog Metrics health check failed: {e}" + logging.exception( + "Datadog metrics healthcheck failed for account %r", account.name + ) + return False, ( + f"Datadog Metrics healthcheck failed for account " + f"{account.name!r}: {e}" + ) def prerequisites_callable(self, config: dict[str, Any]) -> Tuple[bool, str]: if not config: @@ -764,20 +839,23 @@ def prerequisites_callable(self, config: dict[str, Any]) -> Tuple[bool, str]: try: dd_config = DatadogMetricsConfig(**config) - self.dd_config = dd_config - success, error_msg = self._perform_healthcheck(dd_config) - return success, error_msg + for account in dd_config.accounts: + success, error_msg = self._perform_healthcheck(account, dd_config) + if not success: + return False, error_msg + + self.dd_config = dd_config + self._reload_instructions() + return True, "" except Exception as e: logging.exception("Failed to set up Datadog Metrics toolset") return (False, f"Invalid Datadog Metrics configuration: {e}") def _reload_instructions(self): - """Load Datadog metrics specific troubleshooting instructions.""" - template_file_path = os.path.abspath( - os.path.join( - os.path.dirname(__file__), "datadog_metrics_instructions.jinja2" - ) + self._load_llm_instructions_from_file( + os.path.dirname(__file__), + "datadog_metrics_instructions.jinja2", + extra_context={"dd_config": self.dd_config}, ) - self._load_llm_instructions(jinja_template=f"file://{template_file_path}") diff --git a/holmes/plugins/toolsets/datadog/toolset_datadog_traces.py b/holmes/plugins/toolsets/datadog/toolset_datadog_traces.py index 9a37cfe64a..f503676c8e 100644 --- a/holmes/plugins/toolsets/datadog/toolset_datadog_traces.py +++ b/holmes/plugins/toolsets/datadog/toolset_datadog_traces.py @@ -22,7 +22,9 @@ ) from holmes.plugins.toolsets.consts import STANDARD_END_DATETIME_TOOL_PARAM_DESCRIPTION from holmes.plugins.toolsets.datadog.datadog_api import ( + ACCOUNT_TOOL_PARAM_DESCRIPTION, MAX_RETRY_COUNT_ON_RATE_LIMIT, + DatadogAccount, DataDogRequestError, execute_datadog_http_request, get_headers, @@ -65,9 +67,7 @@ def __init__(self): ], tags=[ToolsetTag.CORE], ) - self._load_llm_instructions_from_file( - os.path.dirname(__file__), "instructions_datadog_traces.jinja2" - ) + self._reload_instructions() def prerequisites_callable(self, config: dict[str, Any]) -> Tuple[bool, str]: """Check prerequisites with configuration.""" @@ -76,18 +76,26 @@ def prerequisites_callable(self, config: dict[str, Any]) -> Tuple[bool, str]: try: dd_config = DatadogTracesConfig(**config) + for account in dd_config.accounts: + success, error_msg = self._perform_healthcheck(account, dd_config) + if not success: + return False, error_msg self.dd_config = dd_config - success, error_msg = self._perform_healthcheck(dd_config) - return success, error_msg + self._reload_instructions() + return True, "" except Exception as e: logging.exception("Failed to set up Datadog traces toolset") return False, f"Invalid Datadog Traces configuration: {e}" - def _perform_healthcheck(self, dd_config: DatadogTracesConfig) -> Tuple[bool, str]: - """Perform health check on Datadog traces API.""" + def _perform_healthcheck( + self, account: DatadogAccount, dd_config: DatadogTracesConfig + ) -> Tuple[bool, str]: + """Perform health check on Datadog traces API for one account.""" try: - logging.info("Performing Datadog traces configuration healthcheck...") - headers = get_headers(dd_config) + logging.info( + "Performing Datadog traces healthcheck for account %r...", account.name + ) + headers = get_headers(account) # The spans API uses POST, not GET payload = { @@ -105,8 +113,7 @@ def _perform_healthcheck(self, dd_config: DatadogTracesConfig) -> Tuple[bool, st } } - # Use search endpoint instead - search_url = f"{dd_config.api_url}/api/v2/spans/events/search" + search_url = f"{account.api_url}/api/v2/spans/events/search" execute_datadog_http_request( url=search_url, @@ -120,18 +127,37 @@ def _perform_healthcheck(self, dd_config: DatadogTracesConfig) -> Tuple[bool, st except DataDogRequestError as e: logging.error( - f"Datadog API error during healthcheck: {e.status_code} - {e.response_text}" + "Datadog traces healthcheck failed for account %r: %s - %s", + account.name, + e.status_code, + e.response_text, ) if e.status_code == 403: return ( False, - "API key lacks required permissions. Make sure your API key has 'apm_read' scope.", + f"Datadog Traces healthcheck failed for account " + f"{account.name!r}: API key lacks required permissions. " + "Make sure your API key has 'apm_read' scope.", ) - else: - return False, f"Datadog API error: {e.status_code} - {e.response_text}" + return ( + False, + f"Datadog Traces healthcheck failed for account " + f"{account.name!r}: API error {e.status_code} - {e.response_text}", + ) except Exception as e: - logging.exception("Failed during Datadog traces health check") - return False, f"Datadog Traces health check failed: {e}" + logging.exception( + "Datadog traces healthcheck failed for account %r", account.name + ) + return False, ( + f"Datadog Traces healthcheck failed for account {account.name!r}: {e}" + ) + + def _reload_instructions(self): + self._load_llm_instructions_from_file( + os.path.dirname(__file__), + "instructions_datadog_traces.jinja2", + extra_context={"dd_config": self.dd_config}, + ) class BaseDatadogTracesTool(Tool): @@ -212,6 +238,11 @@ def __init__(self, toolset: "DatadogTracesToolset"): type="boolean", required=True, ), + "account": ToolParameter( + description=ACCOUNT_TOOL_PARAM_DESCRIPTION, + type="string", + required=False, + ), }, toolset=toolset, ) @@ -228,6 +259,14 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes error="Datadog configuration not initialized", params=params, ) + try: + account = self.toolset.dd_config.get_account(params.get("account")) + except KeyError as exc: + return StructuredToolResult( + status=StructuredToolResultStatus.ERROR, + error=str(exc), + params=params, + ) url = None payload: Optional[Dict[str, Any]] = None @@ -252,8 +291,8 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes sort = "-timestamp" # Use POST endpoint for more complex searches - url = f"{self.toolset.dd_config.api_url}/api/v2/spans/events/search" - headers = get_headers(self.toolset.dd_config) + url = f"{account.api_url}/api/v2/spans/events/search" + headers = get_headers(account) payload = { "data": { @@ -291,7 +330,7 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes ] web_url = generate_datadog_spans_url( - self.toolset.dd_config, + account, query, from_time_ms, to_time_ms, @@ -548,6 +587,11 @@ def __init__(self, toolset: "DatadogTracesToolset"): type="string", required=False, ), + "account": ToolParameter( + description=ACCOUNT_TOOL_PARAM_DESCRIPTION, + type="string", + required=False, + ), }, toolset=toolset, ) @@ -594,6 +638,14 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes error="Datadog configuration not initialized", params=params, ) + try: + account = self.toolset.dd_config.get_account(params.get("account")) + except KeyError as exc: + return StructuredToolResult( + status=StructuredToolResultStatus.ERROR, + error=str(exc), + params=params, + ) url = None payload = None @@ -613,8 +665,8 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes query = params.get("query", "*") # Build the request payload - url = f"{self.toolset.dd_config.api_url}/api/v2/spans/analytics/aggregate" - headers = get_headers(self.toolset.dd_config) + url = f"{account.api_url}/api/v2/spans/analytics/aggregate" + headers = get_headers(account) # Build payload attributes first # Process compute parameter to fix common p95->pc95 style mistakes @@ -658,7 +710,7 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes ) web_url = generate_datadog_spans_analytics_url( - self.toolset.dd_config, + account, query, from_time_ms, to_time_ms, diff --git a/tests/plugins/toolsets/datadog/logs/test_check_prerequisites.py b/tests/plugins/toolsets/datadog/logs/test_check_prerequisites.py index 71f094ad0d..2ee592d191 100644 --- a/tests/plugins/toolsets/datadog/logs/test_check_prerequisites.py +++ b/tests/plugins/toolsets/datadog/logs/test_check_prerequisites.py @@ -4,10 +4,12 @@ from holmes.plugins.toolsets.datadog.datadog_models import ( DEFAULT_STORAGE_TIER, DataDogStorageTier, + DatadogLogsConfig, ) from holmes.plugins.toolsets.datadog.toolset_datadog_logs import ( DatadogLogsToolset, ) +from tests.conftest import create_mock_tool_invoke_context class TestDatadogToolsetCheckPrerequisites: @@ -111,7 +113,10 @@ def test_check_prerequisites_healthcheck_error(self): toolset.check_prerequisites() assert toolset.status == ToolsetStatusEnum.FAILED - assert toolset.error == 'Datadog API error: 401 - {"errors":["Unauthorized"]}' + assert ( + toolset.error + == 'Datadog Logs healthcheck failed for account \'default\': API error 401 - {"errors":["Unauthorized"]}' + ) @patch( "holmes.plugins.toolsets.datadog.toolset_datadog_logs.execute_datadog_http_request" @@ -130,7 +135,10 @@ def test_check_prerequisites_healthcheck_exception(self, mock_execute_request): toolset.check_prerequisites() assert toolset.status == ToolsetStatusEnum.FAILED - assert "Datadog Logs health check failed: Network error" in toolset.error + assert ( + "Datadog Logs healthcheck failed for account 'default': Network error" + in toolset.error + ) @patch( "holmes.plugins.toolsets.datadog.toolset_datadog_logs.execute_datadog_http_request" @@ -217,3 +225,65 @@ def test_check_prerequisites_integration(self): prerequisite = toolset.prerequisites[0] assert hasattr(prerequisite, "callable") assert prerequisite.callable == toolset.prerequisites_callable + + +class TestDatadogLogsAccountRouting: + """Verify the `account` parameter routes log queries to the right account.""" + + def setup_method(self): + self.toolset = DatadogLogsToolset() + self.toolset.dd_config = DatadogLogsConfig( + accounts=[ + { + "name": "staging", + "api_key": "stg-key", + "app_key": "stg-app", + "api_url": "https://api.stg.datadoghq.eu", + }, + { + "name": "production", + "api_key": "prd-key", + "app_key": "prd-app", + "api_url": "https://api.datadoghq.eu", + "default": True, + }, + ], + timeout_seconds=60, + ) + + @patch( + "holmes.plugins.toolsets.datadog.toolset_datadog_logs.execute_datadog_http_request" + ) + def test_account_param_routes_to_target(self, mock_execute): + mock_execute.return_value = {"data": [], "meta": {}} + tool = self.toolset.tools[0] + result = tool._invoke( + {"account": "staging"}, context=create_mock_tool_invoke_context() + ) + + assert result.status == StructuredToolResultStatus.SUCCESS + call_kwargs = mock_execute.call_args[1] + assert call_kwargs["url"].startswith("https://api.stg.datadoghq.eu") + assert call_kwargs["headers"]["DD-API-KEY"] == "stg-key" + + @patch( + "holmes.plugins.toolsets.datadog.toolset_datadog_logs.execute_datadog_http_request" + ) + def test_omitted_account_uses_default(self, mock_execute): + mock_execute.return_value = {"data": [], "meta": {}} + tool = self.toolset.tools[0] + tool._invoke({}, context=create_mock_tool_invoke_context()) + + call_kwargs = mock_execute.call_args[1] + assert call_kwargs["url"].startswith("https://api.datadoghq.eu") + assert call_kwargs["headers"]["DD-API-KEY"] == "prd-key" + + def test_unknown_account_returns_structured_error(self): + tool = self.toolset.tools[0] + result = tool._invoke( + {"account": "nope"}, context=create_mock_tool_invoke_context() + ) + + assert result.status == StructuredToolResultStatus.ERROR + assert "'staging'" in result.error + assert "'production'" in result.error diff --git a/tests/plugins/toolsets/datadog/metrics/test_datadog_metrics.py b/tests/plugins/toolsets/datadog/metrics/test_datadog_metrics.py index 30995d5523..a7549d364c 100644 --- a/tests/plugins/toolsets/datadog/metrics/test_datadog_metrics.py +++ b/tests/plugins/toolsets/datadog/metrics/test_datadog_metrics.py @@ -262,7 +262,9 @@ def test_healthcheck_success(self, mock_get): response.json.return_value = {"valid": True} mock_get.return_value = response - success, error_msg = self.toolset._perform_healthcheck(self.config) + success, error_msg = self.toolset._perform_healthcheck( + self.config.accounts[0], self.config + ) assert success is True assert error_msg == "" @@ -277,10 +279,12 @@ def test_healthcheck_failure(self, mock_get): response.json.return_value = {"valid": False} mock_get.return_value = response - success, error_msg = self.toolset._perform_healthcheck(self.config) + success, error_msg = self.toolset._perform_healthcheck( + self.config.accounts[0], self.config + ) assert success is False - assert "validation failed" in error_msg.lower() + assert "valid=false" in error_msg.lower() @patch("holmes.plugins.toolsets.datadog.datadog_api.requests.get") def test_prerequisites_callable_success(self, mock_get): @@ -319,3 +323,72 @@ def test_prerequisites_callable_invalid_config(self): assert success is False assert "Invalid Datadog Metrics configuration" in error_msg + + +class TestDatadogMetricsAccountRouting: + """Verify the LLM-provided ``account`` parameter routes calls correctly.""" + + def setup_method(self): + self.toolset = DatadogMetricsToolset() + self.toolset.dd_config = DatadogMetricsConfig( + accounts=[ + { + "name": "staging", + "api_key": "stg-key", + "app_key": "stg-app", + "api_url": "https://api.stg.datadoghq.eu", + }, + { + "name": "production", + "api_key": "prd-key", + "app_key": "prd-app", + "api_url": "https://api.datadoghq.eu", + "default": True, + }, + ], + default_limit=1000, + timeout_seconds=60, + ) + + @patch("holmes.plugins.toolsets.datadog.datadog_api.requests.get") + def test_account_param_routes_to_target_url_and_keys(self, mock_get): + response = Mock() + response.status_code = 200 + response.json.return_value = {"metrics": ["system.cpu.user"]} + mock_get.return_value = response + + # tools[0] is ListActiveMetrics + tool = self.toolset.tools[0] + result = tool._invoke( + {"account": "staging"}, context=create_mock_tool_invoke_context() + ) + + assert result.status == StructuredToolResultStatus.SUCCESS + call_args = mock_get.call_args + assert call_args[0][0].startswith("https://api.stg.datadoghq.eu") + assert call_args[1]["headers"]["DD-API-KEY"] == "stg-key" + assert call_args[1]["headers"]["DD-APPLICATION-KEY"] == "stg-app" + + @patch("holmes.plugins.toolsets.datadog.datadog_api.requests.get") + def test_omitted_account_uses_default(self, mock_get): + response = Mock() + response.status_code = 200 + response.json.return_value = {"metrics": []} + mock_get.return_value = response + + tool = self.toolset.tools[0] + tool._invoke({}, context=create_mock_tool_invoke_context()) + + call_args = mock_get.call_args + assert call_args[0][0].startswith("https://api.datadoghq.eu") + assert call_args[1]["headers"]["DD-API-KEY"] == "prd-key" + + def test_unknown_account_returns_structured_error(self): + tool = self.toolset.tools[0] + result = tool._invoke( + {"account": "nope"}, context=create_mock_tool_invoke_context() + ) + + assert result.status == StructuredToolResultStatus.ERROR + assert "'staging'" in result.error + assert "'production'" in result.error diff --git a/tests/plugins/toolsets/datadog/test_datadog_accounts.py b/tests/plugins/toolsets/datadog/test_datadog_accounts.py new file mode 100644 index 0000000000..c1a407f7f6 --- /dev/null +++ b/tests/plugins/toolsets/datadog/test_datadog_accounts.py @@ -0,0 +1,174 @@ +"""Tests for the multi-account Datadog config plumbing. + +Covers the backward-compatible single-account shorthand, the new +``accounts:`` list form, and the ``get_account()`` resolver used by every +Datadog tool to route a call to the right account. +""" + +import pytest + +from holmes.plugins.toolsets.datadog.datadog_api import ( + DatadogAccount, + DatadogBaseConfig, +) + + +def _base_config(**kwargs): + return DatadogBaseConfig(**kwargs) + + +class TestLegacyShorthand: + """The single-account shorthand must keep working unchanged.""" + + def test_shorthand_synthesizes_default_account(self): + cfg = _base_config( + api_key="k", app_key="a", api_url="https://api.datadoghq.com" + ) + assert len(cfg.accounts) == 1 + account = cfg.accounts[0] + assert account.name == "default" + assert account.default is True + assert account.api_key == "k" + assert account.app_key == "a" + assert str(account.api_url).rstrip("/") == "https://api.datadoghq.com" + + def test_shorthand_missing_field_raises(self): + with pytest.raises(ValueError, match="missing required field"): + _base_config(api_key="k", api_url="https://api.datadoghq.com") + + def test_no_credentials_raises(self): + with pytest.raises(ValueError, match="missing credentials"): + _base_config() + + +class TestAccountsList: + """The new ``accounts:`` form supports N accounts.""" + + def test_accounts_list_preserves_names(self): + cfg = _base_config( + accounts=[ + DatadogAccount( + name="staging", + api_key="ks", + app_key="as", + api_url="https://api.datadoghq.eu", + ), + DatadogAccount( + name="production", + api_key="kp", + app_key="ap", + api_url="https://api.datadoghq.eu", + default=True, + ), + ] + ) + assert [a.name for a in cfg.accounts] == ["staging", "production"] + assert cfg.accounts[1].default is True + + def test_first_account_becomes_default_when_unset(self): + cfg = _base_config( + accounts=[ + DatadogAccount( + name="a", api_key="k", app_key="a", api_url="https://x.example" + ), + DatadogAccount( + name="b", api_key="k", app_key="a", api_url="https://x.example" + ), + ] + ) + assert cfg.accounts[0].default is True + assert cfg.accounts[1].default is False + + def test_both_forms_together_is_rejected(self): + with pytest.raises(ValueError, match="either the top-level"): + _base_config( + api_key="k", + app_key="a", + api_url="https://api.datadoghq.com", + accounts=[ + DatadogAccount( + name="x", + api_key="k", + app_key="a", + api_url="https://api.datadoghq.com", + ) + ], + ) + + def test_duplicate_names_rejected(self): + with pytest.raises(ValueError, match="duplicate account name"): + _base_config( + accounts=[ + DatadogAccount( + name="dup", + api_key="k", + app_key="a", + api_url="https://x.example", + ), + DatadogAccount( + name="dup", + api_key="k", + app_key="a", + api_url="https://x.example", + ), + ] + ) + + def test_multiple_defaults_rejected(self): + with pytest.raises(ValueError, match="at most one account"): + _base_config( + accounts=[ + DatadogAccount( + name="a", + api_key="k", + app_key="a", + api_url="https://x.example", + default=True, + ), + DatadogAccount( + name="b", + api_key="k", + app_key="a", + api_url="https://x.example", + default=True, + ), + ] + ) + + +class TestGetAccount: + """`get_account` is what every tool calls to resolve the target account.""" + + @pytest.fixture + def cfg(self): + return _base_config( + accounts=[ + DatadogAccount( + name="staging", + api_key="ks", + app_key="as", + api_url="https://api.datadoghq.eu", + ), + DatadogAccount( + name="production", + api_key="kp", + app_key="ap", + api_url="https://api.datadoghq.eu", + default=True, + ), + ] + ) + + def test_none_returns_default(self, cfg): + assert cfg.get_account(None).name == "production" + + def test_by_name(self, cfg): + assert cfg.get_account("staging").name == "staging" + + def test_unknown_lists_available(self, cfg): + with pytest.raises(KeyError) as exc: + cfg.get_account("nope") + msg = str(exc.value) + assert "'staging'" in msg + assert "'production'" in msg + assert "'nope'" in msg diff --git a/tests/plugins/toolsets/datadog/test_toolset_datadog_general.py b/tests/plugins/toolsets/datadog/test_toolset_datadog_general.py index 3c5ed2d660..0a0f39cbff 100644 --- a/tests/plugins/toolsets/datadog/test_toolset_datadog_general.py +++ b/tests/plugins/toolsets/datadog/test_toolset_datadog_general.py @@ -2,6 +2,8 @@ from unittest.mock import Mock, patch +import responses + from holmes.core.tools import StructuredToolResultStatus from holmes.plugins.toolsets.datadog.toolset_datadog_general import ( DatadogGeneralToolset, @@ -173,3 +175,77 @@ def test_api_get_tool(self, mock_headers, mock_execute): assert result.status == StructuredToolResultStatus.ERROR assert "blacklisted operation" in result.error + + +class TestDatadogGeneralAccountRouting: + """Verify the `account` parameter routes API calls to the right account.""" + + def setup_method(self): + from holmes.plugins.toolsets.datadog.datadog_models import DatadogGeneralConfig + from holmes.plugins.toolsets.datadog.toolset_datadog_general import ( + DatadogGeneralToolset, + ) + + self.toolset = DatadogGeneralToolset() + self.toolset.dd_config = DatadogGeneralConfig( + accounts=[ + { + "name": "staging", + "api_key": "stg-key", + "app_key": "stg-app", + "api_url": "https://api.stg.datadoghq.eu", + }, + { + "name": "production", + "api_key": "prd-key", + "app_key": "prd-app", + "api_url": "https://api.datadoghq.eu", + "default": True, + }, + ], + timeout_seconds=60, + ) + + def test_account_param_routes_to_target(self): + get_tool = self.toolset.tools[0] # DatadogAPIGet + with responses.RequestsMock() as rsps: + rsps.add( + responses.GET, + "https://api.stg.datadoghq.eu/api/v1/monitor", + json={"data": []}, + status=200, + ) + result = get_tool._invoke( + {"endpoint": "/api/v1/monitor", "account": "staging"}, + context=create_mock_tool_invoke_context(), + ) + assert result.status == StructuredToolResultStatus.SUCCESS + assert rsps.calls[0].request.url.startswith("https://api.stg.datadoghq.eu") + assert rsps.calls[0].request.headers["DD-API-KEY"] == "stg-key" + + def test_omitted_account_uses_default(self): + get_tool = self.toolset.tools[0] + with responses.RequestsMock() as rsps: + rsps.add( + responses.GET, + "https://api.datadoghq.eu/api/v1/monitor", + json={"data": []}, + status=200, + ) + get_tool._invoke( + {"endpoint": "/api/v1/monitor"}, + context=create_mock_tool_invoke_context(), + ) + assert rsps.calls[0].request.url.startswith("https://api.datadoghq.eu") + assert rsps.calls[0].request.headers["DD-API-KEY"] == "prd-key" + + def test_unknown_account_returns_structured_error(self): + get_tool = self.toolset.tools[0] + result = get_tool._invoke( + {"endpoint": "/api/v1/monitor", "account": "nope"}, + context=create_mock_tool_invoke_context(), + ) + + assert result.status == StructuredToolResultStatus.ERROR + assert "'staging'" in result.error + assert "'production'" in result.error diff --git a/tests/plugins/toolsets/datadog/traces/test_datadog_traces.py b/tests/plugins/toolsets/datadog/traces/test_datadog_traces.py index 000e4c7bf7..1de871e5f8 100644 --- a/tests/plugins/toolsets/datadog/traces/test_datadog_traces.py +++ b/tests/plugins/toolsets/datadog/traces/test_datadog_traces.py @@ -158,3 +158,68 @@ def test_invoke_no_spans_found(self, mock_execute): # When no data is found, the tool still returns success with empty data assert result.data == mock_execute.return_value assert len(result.data["data"]) == 0 + + +class TestDatadogTracesAccountRouting: + """Verify the `account` parameter routes span queries to the right account.""" + + def setup_method(self): + from holmes.plugins.toolsets.datadog.datadog_api import DatadogAccount + from holmes.plugins.toolsets.datadog.datadog_models import DatadogTracesConfig + + self.toolset = DatadogTracesToolset() + self.toolset.dd_config = DatadogTracesConfig( + accounts=[ + { + "name": "staging", + "api_key": "stg-key", + "app_key": "stg-app", + "api_url": "https://api.stg.datadoghq.eu", + }, + { + "name": "production", + "api_key": "prd-key", + "app_key": "prd-app", + "api_url": "https://api.datadoghq.eu", + "default": True, + }, + ], + timeout_seconds=60, + ) + + @patch( + "holmes.plugins.toolsets.datadog.toolset_datadog_traces.execute_datadog_http_request" + ) + def test_account_param_routes_to_target(self, mock_execute): + mock_execute.return_value = {"data": [], "meta": {"page": {}}} + tool = self.toolset.tools[0] # GetSpans + result = tool._invoke( + {"account": "staging"}, context=create_mock_tool_invoke_context() + ) + + assert result.status == StructuredToolResultStatus.SUCCESS + call_kwargs = mock_execute.call_args[1] + assert call_kwargs["url"].startswith("https://api.stg.datadoghq.eu") + assert call_kwargs["headers"]["DD-API-KEY"] == "stg-key" + + @patch( + "holmes.plugins.toolsets.datadog.toolset_datadog_traces.execute_datadog_http_request" + ) + def test_omitted_account_uses_default(self, mock_execute): + mock_execute.return_value = {"data": [], "meta": {"page": {}}} + tool = self.toolset.tools[0] + tool._invoke({}, context=create_mock_tool_invoke_context()) + + call_kwargs = mock_execute.call_args[1] + assert call_kwargs["url"].startswith("https://api.datadoghq.eu") + assert call_kwargs["headers"]["DD-API-KEY"] == "prd-key" + + def test_unknown_account_returns_structured_error(self): + tool = self.toolset.tools[0] + result = tool._invoke( + {"account": "nope"}, context=create_mock_tool_invoke_context() + ) + + assert result.status == StructuredToolResultStatus.ERROR + assert "'staging'" in result.error + assert "'production'" in result.error From 55be47420a383f1977dad63d033a04b475513d42 Mon Sep 17 00:00:00 2001 From: mdecalf Date: Mon, 8 Jun 2026 14:41:25 +0200 Subject: [PATCH 2/2] fix(toolsets/datadog): address review feedback on multi-account PR - Extend _hidden_fields in DatadogTracesConfig and DatadogLogsConfig to include "accounts" alongside "indexes", preventing the parent ClassVar from being shadowed and leaking `accounts` into the form-UI JSON schema. - Align DatadogLogsToolset._perform_healthcheck signature with the metrics and traces pattern: accept dd_config as an explicit parameter and move self.dd_config assignment to after the per-account healthcheck loop. - Add str(...).rstrip("/") to all 10 URL construction sites in logs, metrics and traces toolsets to prevent Pydantic v2 AnyUrl trailing-slash from producing double-slash paths (e.g. https://api.datadoghq.eu//api/v1). Consistent with the existing pattern already used in the general toolset. --- .../toolsets/datadog/datadog_models.py | 6 ++++-- .../toolsets/datadog/toolset_datadog_logs.py | 19 ++++++++++--------- .../datadog/toolset_datadog_metrics.py | 10 +++++----- .../datadog/toolset_datadog_traces.py | 6 +++--- 4 files changed, 22 insertions(+), 19 deletions(-) diff --git a/holmes/plugins/toolsets/datadog/datadog_models.py b/holmes/plugins/toolsets/datadog/datadog_models.py index d7398ab924..f226625163 100644 --- a/holmes/plugins/toolsets/datadog/datadog_models.py +++ b/holmes/plugins/toolsets/datadog/datadog_models.py @@ -47,7 +47,8 @@ class DatadogTracesConfig(DatadogBaseConfig): # Hide list-typed advanced fields from the frontend form and example YAML. # The runtime still accepts them via raw YAML for users who need to override. - _hidden_fields: ClassVar[List[str]] = ["indexes"] + # Must include "accounts" from the parent so ClassVar shadowing doesn't leak it. + _hidden_fields: ClassVar[List[str]] = ["accounts", "indexes"] indexes: list[str] = Field( default_factory=lambda: ["*"], @@ -62,7 +63,8 @@ class DatadogLogsConfig(DatadogBaseConfig): # Hide the `indexes` list from the frontend form and example YAML # because complex list types don't render as form inputs. Runtime still # accepts it via raw YAML for advanced users. - _hidden_fields: ClassVar[List[str]] = ["indexes"] + # Must include "accounts" from the parent so ClassVar shadowing doesn't leak it. + _hidden_fields: ClassVar[List[str]] = ["accounts", "indexes"] indexes: list[str] = Field( default_factory=lambda: ["*"], diff --git a/holmes/plugins/toolsets/datadog/toolset_datadog_logs.py b/holmes/plugins/toolsets/datadog/toolset_datadog_logs.py index 441e093b34..e689cae31b 100644 --- a/holmes/plugins/toolsets/datadog/toolset_datadog_logs.py +++ b/holmes/plugins/toolsets/datadog/toolset_datadog_logs.py @@ -81,10 +81,10 @@ def __init__(self): self.tools = [GetLogs(toolset=self)] self._reload_instructions() - def _perform_healthcheck(self, account: DatadogAccount) -> Tuple[bool, str]: + def _perform_healthcheck( + self, account: DatadogAccount, dd_config: DatadogLogsConfig + ) -> Tuple[bool, str]: """Perform health check on Datadog logs API for one account.""" - if not self.dd_config: - return False, "Internal error: Datadog configuration not initialized" try: logging.info( "Performing Datadog logs healthcheck for account %r...", account.name @@ -95,17 +95,17 @@ def _perform_healthcheck(self, account: DatadogAccount) -> Tuple[bool, str]: "from": "now-1m", "to": "now", "query": "*", - "indexes": self.dd_config.indexes, + "indexes": dd_config.indexes, }, "page": {"limit": 1}, } - search_url = f"{account.api_url}/api/v2/logs/events/search" + search_url = f"{str(account.api_url).rstrip('/')}/api/v2/logs/events/search" execute_datadog_http_request( url=search_url, headers=headers, payload_or_params=payload, - timeout=self.dd_config.timeout_seconds, + timeout=dd_config.timeout_seconds, method="POST", ) @@ -147,12 +147,13 @@ def prerequisites_callable(self, config: dict[str, Any]) -> Tuple[bool, str]: try: dd_config = DatadogLogsConfig(**config) - self.dd_config = dd_config for account in dd_config.accounts: - success, error_msg = self._perform_healthcheck(account) + success, error_msg = self._perform_healthcheck(account, dd_config) if not success: return False, error_msg + + self.dd_config = dd_config self._reload_instructions() return True, "" @@ -255,7 +256,7 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes params["limit"] = limit sort = "timestamp" if params.get("sort_desc", False) else "-timestamp" - url = f"{account.api_url}/api/v2/logs/events/search" + url = f"{str(account.api_url).rstrip('/')}/api/v2/logs/events/search" headers = get_headers(account) storage = self.toolset.dd_config.storage_tier diff --git a/holmes/plugins/toolsets/datadog/toolset_datadog_metrics.py b/holmes/plugins/toolsets/datadog/toolset_datadog_metrics.py index 58bc16c712..89e141c9f3 100644 --- a/holmes/plugins/toolsets/datadog/toolset_datadog_metrics.py +++ b/holmes/plugins/toolsets/datadog/toolset_datadog_metrics.py @@ -140,7 +140,7 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes default_time_span_seconds=ACTIVE_METRICS_DEFAULT_TIME_SPAN_SECONDS, ) - url = f"{account.api_url}/api/v1/metrics" + url = f"{str(account.api_url).rstrip('/')}/api/v1/metrics" headers = get_headers(account) query_params = { @@ -357,7 +357,7 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes default_time_span_seconds=DEFAULT_TIME_SPAN_SECONDS, ) - url = f"{account.api_url}/api/v1/query" + url = f"{str(account.api_url).rstrip('/')}/api/v1/query" headers = get_headers(account) query_params = { @@ -578,7 +578,7 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes for metric_name in metric_names: try: - api_url = f"{account.api_url}/api/v1/metrics/{metric_name}" + api_url = f"{str(account.api_url).rstrip('/')}/api/v1/metrics/{metric_name}" data = execute_datadog_http_request( url=api_url, @@ -696,7 +696,7 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes try: metric_name = get_param_or_raise(params, "metric_name") - api_url = f"{account.api_url}/api/v2/metrics/{metric_name}/active-configurations" + api_url = f"{str(account.api_url).rstrip('/')}/api/v2/metrics/{metric_name}/active-configurations" headers = get_headers(account) data = execute_datadog_http_request( @@ -797,7 +797,7 @@ def _perform_healthcheck( account.name, ) - url = f"{account.api_url}/api/v1/validate" + url = f"{str(account.api_url).rstrip('/')}/api/v1/validate" headers = get_headers(account) data = execute_datadog_http_request( diff --git a/holmes/plugins/toolsets/datadog/toolset_datadog_traces.py b/holmes/plugins/toolsets/datadog/toolset_datadog_traces.py index f503676c8e..79d5dc36e4 100644 --- a/holmes/plugins/toolsets/datadog/toolset_datadog_traces.py +++ b/holmes/plugins/toolsets/datadog/toolset_datadog_traces.py @@ -113,7 +113,7 @@ def _perform_healthcheck( } } - search_url = f"{account.api_url}/api/v2/spans/events/search" + search_url = f"{str(account.api_url).rstrip('/')}/api/v2/spans/events/search" execute_datadog_http_request( url=search_url, @@ -291,7 +291,7 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes sort = "-timestamp" # Use POST endpoint for more complex searches - url = f"{account.api_url}/api/v2/spans/events/search" + url = f"{str(account.api_url).rstrip('/')}/api/v2/spans/events/search" headers = get_headers(account) payload = { @@ -665,7 +665,7 @@ def _invoke(self, params: dict, context: ToolInvokeContext) -> StructuredToolRes query = params.get("query", "*") # Build the request payload - url = f"{account.api_url}/api/v2/spans/analytics/aggregate" + url = f"{str(account.api_url).rstrip('/')}/api/v2/spans/analytics/aggregate" headers = get_headers(account) # Build payload attributes first