Skip to main content

omni_dev/datadog/
client.rs

1//! Datadog REST API client.
2//!
3//! Thin `reqwest` wrapper that injects the `DD-API-KEY` and
4//! `DD-APPLICATION-KEY` headers on every request and retries 429 responses via
5//! the shared [`retry_429`](crate::utils::http::retry_429) driver (which honours
6//! `Retry-After` / `X-RateLimit-Reset`).
7
8use anyhow::{Context, Result};
9use reqwest::Client;
10use url::Url;
11
12use crate::datadog::auth::{base_url_for_site, DatadogCredentials};
13use crate::datadog::error::DatadogError;
14use crate::request_log;
15use crate::utils::http::{retry_429, REQUEST_TIMEOUT};
16use crate::utils::secret::Secret;
17
18/// HTTP client for Datadog REST APIs.
19#[derive(Debug)]
20pub struct DatadogClient {
21    client: Client,
22    base_url: String,
23    api_key: Secret,
24    app_key: Secret,
25}
26
27impl DatadogClient {
28    /// Creates a new Datadog API client.
29    ///
30    /// `base_url` should be the full API host, e.g. `https://api.datadoghq.com`.
31    /// For production use, construct via [`Self::from_credentials`]; tests
32    /// pass a wiremock URL directly.
33    pub fn new(base_url: &str, api_key: &str, app_key: &str) -> Result<Self> {
34        let client = Client::builder()
35            .timeout(REQUEST_TIMEOUT)
36            .build()
37            .context("Failed to build HTTP client")?;
38
39        Ok(Self {
40            client,
41            base_url: base_url.trim_end_matches('/').to_string(),
42            api_key: api_key.into(),
43            app_key: app_key.into(),
44        })
45    }
46
47    /// Creates a client from stored credentials.
48    ///
49    /// Respects `DATADOG_API_URL` as an optional override: when set in the
50    /// process environment it replaces the site-derived base URL. Used for
51    /// tests (wiremock) and on-prem Datadog installs.
52    pub fn from_credentials(creds: &DatadogCredentials) -> Result<Self> {
53        Self::from_credentials_with(&crate::utils::env::SystemEnv, creds)
54    }
55
56    /// [`from_credentials`](Self::from_credentials) over an injected
57    /// [`EnvSource`](crate::utils::env::EnvSource).
58    ///
59    /// Tests pass a pure `MapEnv` to exercise the `DATADOG_API_URL` override
60    /// without mutating the process environment (issue #1030).
61    pub(crate) fn from_credentials_with(
62        env: &impl crate::utils::env::EnvSource,
63        creds: &DatadogCredentials,
64    ) -> Result<Self> {
65        let base_url = env
66            .var(crate::datadog::auth::DATADOG_API_URL)
67            .filter(|s| !s.is_empty())
68            .unwrap_or_else(|| base_url_for_site(&creds.site));
69        Self::new(
70            &base_url,
71            creds.api_key.expose_secret(),
72            creds.app_key.expose_secret(),
73        )
74    }
75
76    /// Returns the API base URL (without trailing slash).
77    #[must_use]
78    pub fn base_url(&self) -> &str {
79        &self.base_url
80    }
81
82    /// Builds an absolute API URL by joining `path` onto `base_url`.
83    ///
84    /// `path` is the full path portion including the leading `/api/…` segment
85    /// (the version varies: `/api/v1/…`, `/api/v2/…`). Centralises the
86    /// `Url::parse(…).context("Invalid Datadog base URL")` spelling repeated
87    /// across the `*_api.rs` modules. Takes `base_url` (rather than `&self`) so
88    /// the free `build_*_url` functions — and their unit tests, which pass
89    /// literal base URLs — can call it unchanged.
90    pub(crate) fn api_url(base_url: &str, path: &str) -> Result<Url> {
91        Url::parse(&format!("{base_url}{path}")).context("Invalid Datadog base URL")
92    }
93
94    /// Checks `response` for success and deserialises its JSON body into `T`.
95    ///
96    /// Non-success responses become a [`DatadogError`] via
97    /// [`Self::response_to_error`] (preserving the 429 rate-limit summary); on
98    /// success the body is parsed with `context` attached on failure. Used by
99    /// the paginated and POST call sites that already hold a response;
100    /// single-shot GETs use [`Self::get_parsed`].
101    pub(crate) async fn parse_response<T: serde::de::DeserializeOwned>(
102        &self,
103        response: reqwest::Response,
104        context: &'static str,
105    ) -> Result<T> {
106        if !response.status().is_success() {
107            return Err(Self::response_to_error(response).await.into());
108        }
109        response.json().await.context(context)
110    }
111
112    /// Sends an authenticated GET and deserialises the JSON body into `T`.
113    ///
114    /// Convenience wrapper over [`Self::get_json`] + [`Self::parse_response`]
115    /// for the common single-request GET-then-parse pattern.
116    pub(crate) async fn get_parsed<T: serde::de::DeserializeOwned>(
117        &self,
118        url: &str,
119        context: &'static str,
120    ) -> Result<T> {
121        let response = self.get_json(url).await?;
122        self.parse_response(response, context).await
123    }
124
125    /// Sends an authenticated GET request and returns the raw response.
126    pub async fn get_json(&self, url: &str) -> Result<reqwest::Response> {
127        retry_429(
128            || {
129                self.client
130                    .get(url)
131                    .header("DD-API-KEY", self.api_key.expose_secret())
132                    .header("DD-APPLICATION-KEY", self.app_key.expose_secret())
133                    .header("Accept", "application/json")
134            },
135            |started, result| {
136                request_log::record_http_result("datadog", "GET", url, started, result);
137            },
138        )
139        .await
140        .context("Failed to send GET request to Datadog API")
141    }
142
143    /// Sends an authenticated POST request with a JSON body and returns the raw response.
144    pub async fn post_json<T: serde::Serialize + Sync + ?Sized>(
145        &self,
146        url: &str,
147        body: &T,
148    ) -> Result<reqwest::Response> {
149        retry_429(
150            || {
151                self.client
152                    .post(url)
153                    .header("DD-API-KEY", self.api_key.expose_secret())
154                    .header("DD-APPLICATION-KEY", self.app_key.expose_secret())
155                    .header("Content-Type", "application/json")
156                    .header("Accept", "application/json")
157                    .json(body)
158            },
159            |started, result| {
160                request_log::record_http_result("datadog", "POST", url, started, result);
161            },
162        )
163        .await
164        .context("Failed to send POST request to Datadog API")
165    }
166
167    /// Consumes a non-success response and turns it into a [`DatadogError`].
168    ///
169    /// For 429 responses, appends a human-readable rate-limit summary
170    /// (extracted from `X-RateLimit-*` headers) to the body, so the caller
171    /// sees why the retry loop gave up.
172    pub async fn response_to_error(response: reqwest::Response) -> DatadogError {
173        let status = response.status().as_u16();
174        let headers = response.headers().clone();
175        let body = response.text().await.unwrap_or_default();
176        let body = if status == 429 {
177            match format_rate_limit(&headers) {
178                Some(suffix) => format!("{body} {suffix}").trim().to_string(),
179                None => body,
180            }
181        } else {
182            body
183        };
184        DatadogError::ApiRequestFailed { status, body }
185    }
186}
187
188fn format_rate_limit(headers: &reqwest::header::HeaderMap) -> Option<String> {
189    let remaining = headers
190        .get("X-RateLimit-Remaining")
191        .and_then(|v| v.to_str().ok());
192    let reset = headers
193        .get("X-RateLimit-Reset")
194        .and_then(|v| v.to_str().ok());
195    let limit = headers
196        .get("X-RateLimit-Limit")
197        .and_then(|v| v.to_str().ok());
198
199    if remaining.is_none() && reset.is_none() && limit.is_none() {
200        return None;
201    }
202
203    let mut parts = Vec::new();
204    if let Some(v) = remaining {
205        parts.push(format!("remaining={v}"));
206    }
207    if let Some(v) = limit {
208        parts.push(format!("limit={v}"));
209    }
210    if let Some(v) = reset {
211        parts.push(format!("reset_in={v}s"));
212    }
213    Some(format!("[rate-limit: {}]", parts.join(", ")))
214}
215
216#[cfg(test)]
217#[allow(clippy::unwrap_used, clippy::expect_used)]
218mod tests {
219    use super::*;
220
221    #[test]
222    fn new_client_strips_trailing_slash() {
223        let client = DatadogClient::new("https://api.datadoghq.com/", "api", "app").unwrap();
224        assert_eq!(client.base_url(), "https://api.datadoghq.com");
225    }
226
227    #[test]
228    fn new_client_preserves_clean_url() {
229        let client = DatadogClient::new("https://api.datadoghq.com", "api", "app").unwrap();
230        assert_eq!(client.base_url(), "https://api.datadoghq.com");
231    }
232
233    #[test]
234    fn client_debug_redacts_keys() {
235        let client = DatadogClient::new(
236            "https://api.datadoghq.com",
237            "sekret-api-key",
238            "sekret-app-key",
239        )
240        .unwrap();
241        // Debug must never print the key values (#1131).
242        let debug = format!("{client:?}");
243        assert!(!debug.contains("sekret-api-key"), "leaked api_key: {debug}");
244        assert!(!debug.contains("sekret-app-key"), "leaked app_key: {debug}");
245        assert!(debug.contains("api_key: <redacted>"));
246        assert!(debug.contains("app_key: <redacted>"));
247    }
248
249    #[test]
250    fn from_credentials_builds_base_url_from_site() {
251        let env = crate::test_support::env::MapEnv::new();
252        let creds = DatadogCredentials {
253            api_key: "api".into(),
254            app_key: "app".into(),
255            site: "us5.datadoghq.com".to_string(),
256        };
257        let client = DatadogClient::from_credentials_with(&env, &creds).unwrap();
258        assert_eq!(client.base_url(), "https://api.us5.datadoghq.com");
259    }
260
261    #[test]
262    fn from_credentials_honours_api_url_override() {
263        let env = crate::test_support::env::MapEnv::new().with(
264            crate::datadog::auth::DATADOG_API_URL,
265            "http://proxy.example:8080",
266        );
267        let creds = DatadogCredentials {
268            api_key: "api".into(),
269            app_key: "app".into(),
270            site: "us5.datadoghq.com".to_string(),
271        };
272        let client = DatadogClient::from_credentials_with(&env, &creds).unwrap();
273        assert_eq!(client.base_url(), "http://proxy.example:8080");
274    }
275
276    #[test]
277    fn from_credentials_ignores_empty_api_url_override() {
278        let env =
279            crate::test_support::env::MapEnv::new().with(crate::datadog::auth::DATADOG_API_URL, "");
280        let creds = DatadogCredentials {
281            api_key: "api".into(),
282            app_key: "app".into(),
283            site: "datadoghq.com".to_string(),
284        };
285        let client = DatadogClient::from_credentials_with(&env, &creds).unwrap();
286        assert_eq!(client.base_url(), "https://api.datadoghq.com");
287    }
288
289    #[tokio::test]
290    async fn get_json_sends_auth_headers() {
291        let server = wiremock::MockServer::start().await;
292        wiremock::Mock::given(wiremock::matchers::method("GET"))
293            .and(wiremock::matchers::path("/test"))
294            .and(wiremock::matchers::header("DD-API-KEY", "my-api"))
295            .and(wiremock::matchers::header("DD-APPLICATION-KEY", "my-app"))
296            .and(wiremock::matchers::header("Accept", "application/json"))
297            .respond_with(
298                wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({"ok": true})),
299            )
300            .expect(1)
301            .mount(&server)
302            .await;
303
304        let client = DatadogClient::new(&server.uri(), "my-api", "my-app").unwrap();
305        let resp = client
306            .get_json(&format!("{}/test", server.uri()))
307            .await
308            .unwrap();
309        assert!(resp.status().is_success());
310    }
311
312    #[tokio::test]
313    async fn post_json_sends_body_and_auth() {
314        let server = wiremock::MockServer::start().await;
315        wiremock::Mock::given(wiremock::matchers::method("POST"))
316            .and(wiremock::matchers::path("/test"))
317            .and(wiremock::matchers::header("DD-API-KEY", "my-api"))
318            .and(wiremock::matchers::header("DD-APPLICATION-KEY", "my-app"))
319            .and(wiremock::matchers::header(
320                "Content-Type",
321                "application/json",
322            ))
323            .and(wiremock::matchers::body_json(serde_json::json!({
324                "query": "hello"
325            })))
326            .respond_with(
327                wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({"id": "1"})),
328            )
329            .expect(1)
330            .mount(&server)
331            .await;
332
333        let client = DatadogClient::new(&server.uri(), "my-api", "my-app").unwrap();
334        let body = serde_json::json!({"query": "hello"});
335        let resp = client
336            .post_json(&format!("{}/test", server.uri()), &body)
337            .await
338            .unwrap();
339        assert!(resp.status().is_success());
340    }
341
342    #[tokio::test]
343    async fn get_json_retries_on_429_via_retry_after() {
344        let server = wiremock::MockServer::start().await;
345        wiremock::Mock::given(wiremock::matchers::method("GET"))
346            .and(wiremock::matchers::path("/test"))
347            .respond_with(wiremock::ResponseTemplate::new(429).append_header("Retry-After", "0"))
348            .up_to_n_times(1)
349            .mount(&server)
350            .await;
351        wiremock::Mock::given(wiremock::matchers::method("GET"))
352            .and(wiremock::matchers::path("/test"))
353            .respond_with(
354                wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({"ok": true})),
355            )
356            .up_to_n_times(1)
357            .mount(&server)
358            .await;
359
360        let client = DatadogClient::new(&server.uri(), "api", "app").unwrap();
361        let resp = client
362            .get_json(&format!("{}/test", server.uri()))
363            .await
364            .unwrap();
365        assert!(resp.status().is_success());
366    }
367
368    #[tokio::test]
369    async fn get_json_retries_on_429_via_x_ratelimit_reset() {
370        let server = wiremock::MockServer::start().await;
371        wiremock::Mock::given(wiremock::matchers::method("GET"))
372            .and(wiremock::matchers::path("/test"))
373            .respond_with(
374                wiremock::ResponseTemplate::new(429).append_header("X-RateLimit-Reset", "0"),
375            )
376            .up_to_n_times(1)
377            .mount(&server)
378            .await;
379        wiremock::Mock::given(wiremock::matchers::method("GET"))
380            .and(wiremock::matchers::path("/test"))
381            .respond_with(
382                wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({"ok": true})),
383            )
384            .up_to_n_times(1)
385            .mount(&server)
386            .await;
387
388        let client = DatadogClient::new(&server.uri(), "api", "app").unwrap();
389        let resp = client
390            .get_json(&format!("{}/test", server.uri()))
391            .await
392            .unwrap();
393        assert!(resp.status().is_success());
394    }
395
396    #[tokio::test]
397    async fn post_json_retries_on_429() {
398        let server = wiremock::MockServer::start().await;
399        wiremock::Mock::given(wiremock::matchers::method("POST"))
400            .and(wiremock::matchers::path("/test"))
401            .respond_with(wiremock::ResponseTemplate::new(429).append_header("Retry-After", "0"))
402            .up_to_n_times(1)
403            .mount(&server)
404            .await;
405        wiremock::Mock::given(wiremock::matchers::method("POST"))
406            .and(wiremock::matchers::path("/test"))
407            .respond_with(wiremock::ResponseTemplate::new(201))
408            .up_to_n_times(1)
409            .mount(&server)
410            .await;
411
412        let client = DatadogClient::new(&server.uri(), "api", "app").unwrap();
413        let resp = client
414            .post_json(
415                &format!("{}/test", server.uri()),
416                &serde_json::json!({"k": "v"}),
417            )
418            .await
419            .unwrap();
420        assert_eq!(resp.status().as_u16(), 201);
421    }
422
423    #[tokio::test]
424    async fn get_json_returns_429_after_max_retries() {
425        let server = wiremock::MockServer::start().await;
426        wiremock::Mock::given(wiremock::matchers::method("GET"))
427            .and(wiremock::matchers::path("/test"))
428            .respond_with(wiremock::ResponseTemplate::new(429).append_header("Retry-After", "0"))
429            .mount(&server)
430            .await;
431
432        let client = DatadogClient::new(&server.uri(), "api", "app").unwrap();
433        let resp = client
434            .get_json(&format!("{}/test", server.uri()))
435            .await
436            .unwrap();
437        assert_eq!(resp.status().as_u16(), 429);
438    }
439
440    #[tokio::test]
441    async fn response_to_error_surfaces_rate_limit_headers_on_429() {
442        let server = wiremock::MockServer::start().await;
443        wiremock::Mock::given(wiremock::matchers::method("GET"))
444            .and(wiremock::matchers::path("/test"))
445            .respond_with(
446                wiremock::ResponseTemplate::new(429)
447                    .append_header("Retry-After", "0")
448                    .append_header("X-RateLimit-Remaining", "0")
449                    .append_header("X-RateLimit-Reset", "42")
450                    .append_header("X-RateLimit-Limit", "100")
451                    .set_body_string("too many"),
452            )
453            .mount(&server)
454            .await;
455
456        let client = DatadogClient::new(&server.uri(), "api", "app").unwrap();
457        let resp = client
458            .get_json(&format!("{}/test", server.uri()))
459            .await
460            .unwrap();
461        let err = DatadogClient::response_to_error(resp).await;
462        let msg = err.to_string();
463        assert!(msg.contains("429"));
464        assert!(msg.contains("too many"));
465        assert!(msg.contains("remaining=0"));
466        assert!(msg.contains("limit=100"));
467        assert!(msg.contains("reset_in=42s"));
468    }
469
470    #[tokio::test]
471    async fn response_to_error_does_not_add_rate_limit_suffix_on_non_429() {
472        let server = wiremock::MockServer::start().await;
473        wiremock::Mock::given(wiremock::matchers::method("GET"))
474            .and(wiremock::matchers::path("/test"))
475            .respond_with(wiremock::ResponseTemplate::new(401).set_body_string("Unauthorized"))
476            .mount(&server)
477            .await;
478
479        let client = DatadogClient::new(&server.uri(), "api", "app").unwrap();
480        let resp = client
481            .get_json(&format!("{}/test", server.uri()))
482            .await
483            .unwrap();
484        let err = DatadogClient::response_to_error(resp).await;
485        let msg = err.to_string();
486        assert!(msg.contains("401"));
487        assert!(msg.contains("Unauthorized"));
488        assert!(!msg.contains("rate-limit"));
489    }
490
491    #[tokio::test]
492    async fn response_to_error_omits_suffix_when_no_rate_limit_headers() {
493        let server = wiremock::MockServer::start().await;
494        wiremock::Mock::given(wiremock::matchers::method("GET"))
495            .and(wiremock::matchers::path("/test"))
496            .respond_with(
497                wiremock::ResponseTemplate::new(429)
498                    .append_header("Retry-After", "0")
499                    .set_body_string("slow down"),
500            )
501            .mount(&server)
502            .await;
503
504        let client = DatadogClient::new(&server.uri(), "api", "app").unwrap();
505        let resp = client
506            .get_json(&format!("{}/test", server.uri()))
507            .await
508            .unwrap();
509        let err = DatadogClient::response_to_error(resp).await;
510        let msg = err.to_string();
511        assert!(msg.contains("slow down"));
512        assert!(!msg.contains("rate-limit"));
513    }
514}