[SVLS-8713] feat(logs-agent): optional API key resolver and Datadog intake URL override

crates/datadog-logs-agent/src/flusher.rs

@@ -1,7 +1,10 @@
 // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
+use std::future::Future;
 use std::io::Write as _;
+use std::pin::Pin;
+use std::sync::Arc;
 
 use futures::future::join_all;
 use reqwest::Client;
@@ -15,12 +18,18 @@ use crate::errors::FlushError;
 /// Maximum number of send attempts before giving up on a batch.
 const MAX_FLUSH_ATTEMPTS: u32 = 3;
 
+/// Returns the API key for the primary endpoint, or `None` if it cannot be resolved.
+pub type ApiKeyResolver =
+    Arc<dyn Fn() -> Pin<Box<dyn Future<Output = Option<String>> + Send>> + Send + Sync>;
+
 /// Drains log batches from an [`AggregatorHandle`] and ships them to Datadog.
 #[derive(Clone)]
 pub struct LogFlusher {
     config: LogFlusherConfig,
     client: Client,
     aggregator_handle: AggregatorHandle,
+    api_key_resolver: Option<ApiKeyResolver>,
+    datadog_intake_url: Option<String>,
 }
 
 impl LogFlusher {
@@ -38,9 +47,30 @@ impl LogFlusher {
             config,
             client,
             aggregator_handle,
+            api_key_resolver: None,
+            datadog_intake_url: None,
         }
     }
 
+    /// Get the primary endpoint's API key from `resolver` on every flush instead of using
+    /// `LogFlusherConfig::api_key`. If it returns `None`, the primary endpoint is skipped
+    /// for that flush (its batches are dropped, with one `error!`); additional endpoints
+    /// still use their own keys.
+    #[must_use]
+    pub fn with_api_key_resolver(mut self, resolver: ApiKeyResolver) -> Self {
+        self.api_key_resolver = Some(resolver);
+        self
+    }
+
+    /// In `Destination::Datadog` mode, send to `url` (full URL, including the path) instead
+    /// of `https://http-intake.logs.{site}/api/v2/logs`. Ignored in Observability Pipelines
+    /// Worker mode.
+    #[must_use]
+    pub fn with_datadog_intake_url(mut self, url: impl Into<String>) -> Self {
+        self.datadog_intake_url = Some(url.into());
+        self
+    }
+
     /// Drain the aggregator, ship all pending batches to Datadog, and redrive any
     /// builders that failed transiently in the previous invocation.
     ///
@@ -55,7 +85,8 @@ impl LogFlusher {
     /// A vec of `RequestBuilder`s that still failed after all in-call retries.
     /// The caller should pass these back on the next invocation to re-attempt
     /// delivery. An empty vec means every batch was delivered successfully
-    /// (or encountered a permanent error and was dropped — those are logged).
+    /// (or encountered a permanent error, or the primary API key could not be
+    /// resolved, and was dropped — those are logged).
     ///
     /// Failures on additional endpoints are logged as warnings but their
     /// builders are not included in the returned vec (best-effort delivery).
@@ -88,6 +119,19 @@ impl LogFlusher {
             }
         };
 
+        let primary_api_key = match &self.api_key_resolver {
+            None => Some(self.config.api_key.clone()),
+            Some(resolve) => {
+                let key = resolve().await;
+                if key.is_none() {
+                    error!(
+                        "failed to resolve the logs API key; skipping the primary endpoint for this flush"
+                    );
+                }
+                key
+            }
+        };
+
         if batches.is_empty() {
             debug!("no log batches to flush");
             return failed;
@@ -99,12 +143,17 @@ impl LogFlusher {
 
         let batch_futures = batches.iter().map(|batch| {
             let primary_url = primary_url.clone();
+            let primary_api_key = primary_api_key.as_deref();
             async move {
                 // Primary endpoint — failures are tracked for cross-invocation retry.
                 let is_primary_datadog = matches!(self.config.mode, Destination::Datadog);
-                let primary_result = self
-                    .ship_batch(batch, &primary_url, primary_use_compression, &self.config.api_key, is_primary_datadog)
-                    .await;
+                let primary_result = match primary_api_key {
+                    Some(api_key) => {
+                        self.ship_batch(batch, &primary_url, primary_use_compression, api_key, is_primary_datadog)
+                            .await
+                    }
+                    None => Ok(()),
+                };
 
                 // Additional endpoints — best-effort; failures are only logged.
                 // Additional endpoints are always Datadog intakes regardless of the
@@ -145,7 +194,9 @@ impl LogFlusher {
     fn resolve_endpoint(&self) -> (String, bool) {
         match &self.config.mode {
             Destination::Datadog => {
-                let url = format!("https://http-intake.logs.{}/api/v2/logs", self.config.site);
+                let url = self.datadog_intake_url.clone().unwrap_or_else(|| {
+                    format!("https://http-intake.logs.{}/api/v2/logs", self.config.site)
+                });
                 (url, self.config.use_compression)
             }
             Destination::ObservabilityPipelinesWorker { url } => {

crates/datadog-logs-agent/src/flusher/tests.rs

@@ -1158,3 +1158,209 @@ async fn test_default_logging_warns_once_per_dropped_oversized_entry() {
     assert_eq!(logged.warnings(), 2);
     assert_eq!(logged.errors(), 0, "default logging uses warn!, not error!");
 }
+
+/// A resolver that counts its calls and returns `keys` in order; the last key repeats.
+fn resolver(keys: &[Option<&'static str>], calls: Arc<AtomicUsize>) -> ApiKeyResolver {
+    let keys: Vec<Option<String>> = keys.iter().map(|k| k.map(str::to_string)).collect();
+    Arc::new(
+        move || -> std::pin::Pin<Box<dyn std::future::Future<Output = Option<String>> + Send>> {
+            let n = calls.fetch_add(1, Ordering::SeqCst);
+            let key = keys.get(n).or(keys.last()).cloned().flatten();
+            Box::pin(async move { key })
+        },
+    )
+}
+
+/// In Datadog mode, `with_datadog_intake_url` replaces the public intake URL; headers
+/// and compression are unchanged.
+#[tokio::test]
+async fn test_datadog_intake_url_override_is_used_in_datadog_mode() {
+    let (service, handle) = AggregatorService::new();
+    let _task = tokio::spawn(service.run());
+    let mut server = mockito::Server::new_async().await;
+    let mock = server
+        .mock("POST", "/custom/path")
+        .match_header("DD-API-KEY", "test-api-key")
+        .match_header("DD-PROTOCOL", "agent-json")
+        .match_header("Content-Encoding", "zstd")
+        .with_status(202)
+        .expect(1)
+        .create_async()
+        .await;
+    let config = LogFlusherConfig {
+        mode: Destination::Datadog,
+        use_compression: true,
+        ..config_for_mock("http://unused")
+    };
+    let client = reqwest::Client::builder().build().expect("client");
+    let flusher = LogFlusher::new(config, client, handle.clone())
+        .with_datadog_intake_url(format!("{}/custom/path", server.url()));
+    handle
+        .insert_batch(vec![make_entry("override")])
+        .expect("insert");
+
+    assert!(flusher.flush(vec![]).await.is_empty());
+    mock.assert_async().await;
+}
+
+/// `with_datadog_intake_url` does not affect Observability Pipelines Worker mode.
+#[tokio::test]
+async fn test_datadog_intake_url_override_is_ignored_for_opw() {
+    let (service, handle) = AggregatorService::new();
+    let _task = tokio::spawn(service.run());
+    let mut server = mockito::Server::new_async().await;
+    let mock = server
+        .mock("POST", "/api/v2/logs")
+        .with_status(202)
+        .expect(1)
+        .create_async()
+        .await;
+    let client = reqwest::Client::builder().build().expect("client");
+    let flusher = LogFlusher::new(config_for_mock(&server.url()), client, handle.clone())
+        .with_datadog_intake_url("http://127.0.0.1:1/never");
+    handle
+        .insert_batch(vec![make_entry("opw")])
+        .expect("insert");
+
+    assert!(flusher.flush(vec![]).await.is_empty());
+    mock.assert_async().await;
+}
+
+/// With a resolver, the primary uses the resolved key, not `LogFlusherConfig::api_key`.
+#[tokio::test]
+async fn test_api_key_resolver_key_is_used_for_the_primary() {
+    let (service, handle) = AggregatorService::new();
+    let _task = tokio::spawn(service.run());
+    let mut server = mockito::Server::new_async().await;
+    let mock = server
+        .mock("POST", "/api/v2/logs")
+        .match_header("DD-API-KEY", "resolved-key")
+        .with_status(202)
+        .expect(1)
+        .create_async()
+        .await;
+    let client = reqwest::Client::builder().build().expect("client");
+    let flusher = LogFlusher::new(config_for_mock(&server.url()), client, handle.clone())
+        .with_api_key_resolver(resolver(
+            &[Some("resolved-key")],
+            Arc::new(AtomicUsize::new(0)),
+        ));
+    handle
+        .insert_batch(vec![make_entry("keyed")])
+        .expect("insert");
+
+    assert!(flusher.flush(vec![]).await.is_empty());
+    mock.assert_async().await;
+}
+
+/// The resolver is called on every flush, even when there is nothing to send.
+#[tokio::test]
+async fn test_api_key_resolver_is_called_on_every_flush() {
+    let calls = Arc::new(AtomicUsize::new(0));
+    let (service, handle) = AggregatorService::new();
+    let _task = tokio::spawn(service.run());
+    let client = reqwest::Client::builder().build().expect("client");
+    let flusher = LogFlusher::new(config_for_mock("http://127.0.0.1:1"), client, handle)
+        .with_api_key_resolver(resolver(&[Some("k")], calls.clone()));
+
+    for _ in 0..3 {
+        assert!(flusher.flush(vec![]).await.is_empty());
+    }
+    assert_eq!(calls.load(Ordering::SeqCst), 3);
+}
+
+/// A key that changes between flushes (e.g. a reloaded secret) is used on the next flush.
+#[tokio::test]
+async fn test_reloaded_api_key_is_used_on_the_next_flush() {
+    let (service, handle) = AggregatorService::new();
+    let _task = tokio::spawn(service.run());
+    let mut server = mockito::Server::new_async().await;
+    let first = server
+        .mock("POST", "/api/v2/logs")
+        .match_header("DD-API-KEY", "key-1")
+        .with_status(202)
+        .expect(1)
+        .create_async()
+        .await;
+    let second = server
+        .mock("POST", "/api/v2/logs")
+        .match_header("DD-API-KEY", "key-2")
+        .with_status(202)
+        .expect(1)
+        .create_async()
+        .await;
+    let client = reqwest::Client::builder().build().expect("client");
+    let flusher = LogFlusher::new(config_for_mock(&server.url()), client, handle.clone())
+        .with_api_key_resolver(resolver(
+            &[Some("key-1"), Some("key-2")],
+            Arc::new(AtomicUsize::new(0)),
+        ));
+
+    handle
+        .insert_batch(vec![make_entry("one")])
+        .expect("insert");
+    assert!(flusher.flush(vec![]).await.is_empty());
+    handle
+        .insert_batch(vec![make_entry("two")])
+        .expect("insert");
+    assert!(flusher.flush(vec![]).await.is_empty());
+    first.assert_async().await;
+    second.assert_async().await;
+}
+
+/// If the key cannot be resolved, the primary is skipped with one error line, and
+/// additional endpoints are still sent with their own keys.
+#[tokio::test(flavor = "current_thread")]
+async fn test_unresolved_api_key_skips_the_primary_but_sends_extras() {
+    let (logged, _guard) = LevelCounter::install();
+    let (service, handle) = AggregatorService::new();
+    let _task = tokio::spawn(service.run());
+    let mut primary = mockito::Server::new_async().await;
+    let mut extra = mockito::Server::new_async().await;
+    let primary_mock = primary
+        .mock("POST", "/api/v2/logs")
+        .expect(0)
+        .create_async()
+        .await;
+    let extra_mock = extra
+        .mock("POST", "/extra")
+        .match_header("DD-API-KEY", "extra-key")
+        .with_status(202)
+        .expect(1)
+        .create_async()
+        .await;
+    let config = LogFlusherConfig {
+        additional_endpoints: vec![LogsAdditionalEndpoint {
+            api_key: "extra-key".to_string(),
+            url: format!("{}/extra", extra.url()),
+            is_reliable: true,
+        }],
+        ..config_for_mock(&primary.url())
+    };
+    let client = reqwest::Client::builder().build().expect("client");
+    let flusher = LogFlusher::new(config, client, handle.clone())
+        .with_api_key_resolver(resolver(&[None], Arc::new(AtomicUsize::new(0))));
+    handle
+        .insert_batch(vec![make_entry("log")])
+        .expect("insert");
+
+    assert!(flusher.flush(vec![]).await.is_empty());
+    primary_mock.assert_async().await;
+    extra_mock.assert_async().await;
+    assert_eq!((logged.warnings(), logged.errors()), (0, 1));
+}
+
+/// The unresolved-key error is logged on every flush, even with nothing to send.
+#[tokio::test(flavor = "current_thread")]
+async fn test_unresolved_api_key_logs_an_error_even_without_batches() {
+    let (logged, _guard) = LevelCounter::install();
+    let (service, handle) = AggregatorService::new();
+    let _task = tokio::spawn(service.run());
+    let client = reqwest::Client::builder().build().expect("client");
+    let flusher = LogFlusher::new(config_for_mock("http://127.0.0.1:1"), client, handle)
+        .with_api_key_resolver(resolver(&[None], Arc::new(AtomicUsize::new(0))));
+
+    assert!(flusher.flush(vec![]).await.is_empty());
+    assert!(flusher.flush(vec![]).await.is_empty());
+    assert_eq!(logged.errors(), 2);
+}

crates/datadog-logs-agent/src/lib.rs

@@ -22,7 +22,7 @@ mod test_support;
 // Re-export the most commonly used types at the crate root
 pub use aggregator::{AggregatorHandle, AggregatorService, OversizePolicy};
 pub use config::{Destination, LogFlusherConfig};
-pub use flusher::LogFlusher;
+pub use flusher::{ApiKeyResolver, LogFlusher};
 pub use intake_entry::IntakeEntry;
 pub use logs_additional_endpoint::LogsAdditionalEndpoint;
 pub use server::{LogServer, LogServerConfig};