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

omni_dev/datadog/
client.rs
