Merge pull request #431 from piochelepiotr/piotr.wolski/data-streams-experimental

feat(data-streams): experimental DSM Kafka commands

Author: platinummonkey

agents/kafka.md

@@ -0,0 +1,145 @@
+---
+description: Inspect Kafka topics/brokers, schema registry, client configs, and on-demand read of Kafka messages. Auto-discovers cluster ID, bootstrap servers, latest produced offset, and consumer-group offsets/lag from Datadog metrics before invoking reads.
+---
+
+# Kafka Agent
+
+You are a specialized agent for inspecting Kafka clusters through Datadog via the `pup` CLI. Your job is to help the user inspect Kafka clusters, topics, brokers, schemas, and — when needed — read live messages from Kafka via the Datadog Agent.
+
+## Important Context
+
+**CLI Tool**: This agent uses the `pup` CLI to execute Datadog API commands.
+
+**Environment Variables**:
+- `DD_API_KEY` / `DD_APP_KEY`: Required if you are not using OAuth2 (`pup auth login`). Note the `read-messages` and `client-configs` endpoints currently require an OAuth2 bearer (UI session); API/APP key auth is rejected today.
+- `DD_SITE`: Datadog site. Default `datadoghq.com`. Use `datad0g.com` for staging.
+
+**API surface**: these commands hit experimental Datadog routes that are **not** part of the public API contract and may change.
+
+## Permission Model
+
+`read-messages` requires the `data_streams_capture_messages` permission and is rate-limited to 10 calls/minute per user.
+
+## Available Commands
+
+```bash
+# Topic config history
+pup kafka topic-configs \
+  --kafka-cluster-id <id> --topic <topic>
+
+# Broker config history
+pup kafka broker-configs \
+  --kafka-cluster-id <id> --broker-id <broker>
+
+# Producer/consumer client configs (one or more service:type pairs)
+pup kafka client-configs \
+  --kafka-cluster-id <id> \
+  --service <svc>:producer \
+  --service <svc>:consumer
+
+# Schema registry — full version history of a subject on a cluster
+pup kafka subject-schemas \
+  --kafka-cluster-id <id> --subject <subject>
+
+# Read live messages (rate-limited, agent-mediated)
+pup kafka read-messages \
+  --cluster <id> --topic <topic> \
+  --bootstrap-servers <host:port,...> \
+  [--partition N] [--start-offset N] [--start-timestamp ms] \
+  [--n-messages-retrieved N] [--max-scanned-messages N] \
+  [--filter expr] [--consumer-group-id <id>]
+```
+
+### `--filter` expressions
+
+`--filter` is a jq-style expression evaluated agent-side against each deserialized message. The message context exposes top-level fields `.key`, `.value`, `.headers`, `.topic`, `.partition`, `.offset`, and `.timestamp`; navigate nested fields with dotted paths (e.g. `.value.user.country`).
+
+- Operators: `==`, `!=`, `>`, `<`, `>=`, `<=`, `contains`.
+- Combine with ` and ` / ` or ` (note: `or` has higher precedence — it is split first).
+- String literals must be quoted with `"` or `'`. Numeric literals are parsed as int/float.
+- A bare path (no operator) is an existence check — true when the field resolves to a non-null value.
+
+Examples:
+
+```bash
+--filter='.value.status == "failed"'
+--filter='.value.amount > 100'
+--filter='.headers.tenant == "acme" and .value.priority >= 5'
+--filter='.value.tags contains "urgent"'
+--filter='.value.error'   # existence
+```
+
+## Auto-discovering arguments via Datadog metrics
+
+Before calling `read-messages`, you almost never have the `kafka_cluster_id` / `bootstrap_servers` / partition / offset on hand. Resolve them by querying Datadog metrics with `pup metrics query`. **These tools are usable only when `kafka.broker.count` is reported for the cluster** — if that metric is empty, do not call `read-messages`.
+
+The relevant metrics (all share the same tag set: `kafka_cluster_id`, `bootstrap_servers`, `topic`, `partition`, and for consumer-group metrics `consumer_group`):
+
+| Metric | What it tells you |
+|---|---|
+| `kafka.broker.count` | Known clusters. Tags: `kafka_cluster_id`, `bootstrap_servers`. |
+| `kafka.broker_offset` | Latest produced offset per partition. Tags: `kafka_cluster_id`, `topic`, `partition`. |
+| `kafka.consumer_offset` | Last committed offset of a consumer group. Tags: `kafka_cluster_id`, `topic`, `partition`, `consumer_group`. |
+| `kafka.consumer_lag` | Consumer lag in offsets. Same tags. |
+| `kafka.estimated_consumer_lag` | Consumer lag in seconds. Same tags. |
+
+### Resolution recipes
+
+**1. Resolve `kafka_cluster_id` + `bootstrap_servers` from a topic name:**
+```bash
+pup metrics query \
+  --query='max:kafka.broker.count{topic:<TOPIC>} by {kafka_cluster_id,bootstrap_servers}' \
+  --from='now-15m'
+```
+The single (or top) returned series's tag values are your `--cluster` and `--bootstrap-servers`.
+
+**2. Find the partition with the most data, and the latest produced offset:**
+```bash
+pup metrics query \
+  --query='max:kafka.broker_offset{topic:<TOPIC>} by {partition}' \
+  --from='now-15m'
+```
+Pick the partition with the highest value. For a tail read use `--start-offset = max - n_messages_retrieved`.
+
+**3. Tail what a consumer hasn't yet processed:**
+```bash
+pup metrics query \
+  --query='max:kafka.consumer_offset{topic:<TOPIC>,consumer_group:<GROUP>} by {partition}' \
+  --from='now-15m'
+```
+Use that value as `--start-offset` and pass `--consumer-group-id <GROUP>`. Pair with `kafka.consumer_lag` (offsets) or `kafka.estimated_consumer_lag` (seconds) to report how far behind the consumer is.
+
+If a query returns no series, surface that to the user instead of guessing.
+
+## Worked example
+
+> "Get the last 5 messages on topic `orders-events`."
+
+1. Resolve cluster + bootstrap:
+   ```bash
+   pup metrics query \
+     --query='max:kafka.broker.count{topic:orders-events} by {kafka_cluster_id,bootstrap_servers}' \
+     --from='now-15m'
+   ```
+2. Find the busiest partition and its latest offset:
+   ```bash
+   pup metrics query \
+     --query='max:kafka.broker_offset{topic:orders-events} by {partition}' \
+     --from='now-15m'
+   ```
+3. Confirm with the user, then read:
+   ```bash
+   pup kafka read-messages \
+     --cluster <id-from-step-1> \
+     --topic orders-events \
+     --bootstrap-servers <bootstrap-from-step-1> \
+     --partition <p-from-step-2> \
+     --start-offset <max-5> \
+     --n-messages-retrieved 5
+   ```
+
+## Failure modes
+
+- **`kafka.broker.count` returns no data** — the cluster is not reporting Kafka telemetry to Datadog; `read-messages` will hang or fail. Stop and tell the user.
+- **`HTTP 403 / data_streams_capture_messages`** — the caller lacks the permission. Ask the user to request it; do not retry.
+- **`HTTP 504` / no response** — no Datadog Agent reachable by Remote Config can connect to the cluster. The cluster may be air-gapped or the Agent isn't deployed where it can see the brokers.

src/client.rs

@@ -862,6 +862,37 @@ fn apply_auth(
     anyhow::bail!("no authentication configured")
 }
 
+/// POST a JSON:API document. Wraps `attributes` in `{data:{type,attributes}}`
+/// and sends with `Content-Type: application/vnd.api+json`. Use for routes
+/// whose decoder is configured for JSON:API.
+pub async fn raw_post_jsonapi(
+    cfg: &Config,
+    path: &str,
+    resource_type: &str,
+    attributes: serde_json::Value,
+) -> anyhow::Result<serde_json::Value> {
+    let url = format!("{}{}", cfg.api_base_url(), path);
+    let envelope = serde_json::json!({
+        "data": { "type": resource_type, "attributes": attributes },
+    });
+    let client = reqwest::Client::new();
+    let mut req = client.post(&url);
+    req = apply_auth(req, cfg, "POST", path)?;
+    let resp = req
+        .header("Content-Type", "application/vnd.api+json")
+        .header("Accept", "application/vnd.api+json")
+        .header("User-Agent", useragent::get())
+        .json(&envelope)
+        .send()
+        .await?;
+    if !resp.status().is_success() {
+        let status = resp.status();
+        let body = resp.text().await.unwrap_or_default();
+        anyhow::bail!("POST {url} failed (HTTP {status}): {body}");
+    }
+    Ok(resp.json().await?)
+}
+
 /// Like `raw_post`, but returns the parsed JSON body even on non-2xx responses.
 /// Callers are responsible for inspecting the body for errors.
 pub async fn raw_post_lenient(

src/commands/kafka.rs

@@ -0,0 +1,117 @@
+//! [Experimental] Kafka inspection commands.
+//!
+//! These endpoints are experimental — the API surface is not covered by
+//! Datadog's public API compatibility guarantees and may change without
+//! notice.
+//!
+//! Authentication: OAuth2 bearer token (e.g. `pup auth login`). DD_API_KEY +
+//! DD_APP_KEY is not accepted by these UI routes.
+
+use anyhow::Result;
+use serde_json::{json, Value};
+
+use crate::client;
+use crate::config::Config;
+use crate::formatter;
+
+const TOPIC_CONFIGS_PATH: &str = "/api/ui/data_streams/kafka_topic_configs";
+const BROKER_CONFIGS_PATH: &str = "/api/ui/data_streams/kafka_broker_configs";
+const CLIENT_CONFIGS_PATH: &str = "/api/ui/data_streams/kafka_client_configs";
+const READ_MESSAGES_PATH: &str = "/api/ui/data_streams/kafka_actions/read_messages";
+const SUBJECT_SCHEMAS_PATH: &str = "/api/ui/data_streams/subject_kafka_schemas";
+
+pub async fn topic_configs(cfg: &Config, kafka_cluster_id: &str, topic: &str) -> Result<()> {
+    let query = [("kafka_cluster_id", kafka_cluster_id), ("topic", topic)];
+    let resp = client::raw_get(cfg, TOPIC_CONFIGS_PATH, &query)
+        .await
+        .map_err(|e| anyhow::anyhow!("failed to get kafka topic configs: {e:?}"))?;
+    formatter::output(cfg, &resp)
+}
+
+pub async fn broker_configs(cfg: &Config, kafka_cluster_id: &str, broker_id: &str) -> Result<()> {
+    let query = [
+        ("kafka_cluster_id", kafka_cluster_id),
+        ("broker_id", broker_id),
+    ];
+    let resp = client::raw_get(cfg, BROKER_CONFIGS_PATH, &query)
+        .await
+        .map_err(|e| anyhow::anyhow!("failed to get kafka broker configs: {e:?}"))?;
+    formatter::output(cfg, &resp)
+}
+
+/// `services` is a list of `service:config_type` pairs where `config_type` is
+/// either `producer` or `consumer`.
+pub async fn client_configs(
+    cfg: &Config,
+    kafka_cluster_id: &str,
+    services: Vec<(String, String)>,
+) -> Result<()> {
+    if services.is_empty() {
+        anyhow::bail!("at least one --service SERVICE:producer|consumer is required");
+    }
+    let services_json: Vec<Value> = services
+        .into_iter()
+        .map(|(service, config_type)| json!({ "service": service, "config_type": config_type }))
+        .collect();
+    let body = json!({
+        "kafka_cluster_id": kafka_cluster_id,
+        "services": services_json,
+    });
+    let resp = client::raw_post(cfg, CLIENT_CONFIGS_PATH, body)
+        .await
+        .map_err(|e| anyhow::anyhow!("failed to get kafka client configs: {e:?}"))?;
+    formatter::output(cfg, &resp)
+}
+
+/// Dispatches a Kafka read-messages action to the agent via Remote Config and
+/// polls until the agent responds.
+#[allow(clippy::too_many_arguments)]
+pub async fn read_messages(
+    cfg: &Config,
+    cluster: &str,
+    topic: &str,
+    bootstrap_servers: &str,
+    partition: Option<i32>,
+    start_offset: i64,
+    start_timestamp: Option<i64>,
+    n_messages_retrieved: u32,
+    max_scanned_messages: u32,
+    filter: Option<String>,
+    consumer_group_id: Option<String>,
+) -> Result<()> {
+    let mut attrs = json!({
+        "cluster": cluster,
+        "topic": topic,
+        "bootstrap_servers": bootstrap_servers,
+        "start_offset": start_offset,
+        "n_messages_retrieved": n_messages_retrieved,
+        "max_scanned_messages": max_scanned_messages,
+    });
+    if let Some(p) = partition {
+        attrs["partition"] = json!(p);
+    }
+    if let Some(ts) = start_timestamp {
+        attrs["start_timestamp"] = json!(ts);
+    }
+    if let Some(f) = filter {
+        attrs["filter"] = json!(f);
+    }
+    if let Some(cg) = consumer_group_id {
+        attrs["consumer_group_id"] = json!(cg);
+    }
+
+    let resp =
+        client::raw_post_jsonapi(cfg, READ_MESSAGES_PATH, "kafka_action_read_messages", attrs)
+            .await
+            .map_err(|e| anyhow::anyhow!("failed to read kafka messages: {e:?}"))?;
+    formatter::output(cfg, &resp)
+}
+
+/// All version history of a single Schema Registry subject on a Kafka cluster.
+pub async fn subject_schemas(cfg: &Config, kafka_cluster_id: &str, subject: &str) -> Result<()> {
+    let query = [("kafka_cluster_id", kafka_cluster_id), ("subject", subject)];
+    let resp = client::raw_get(cfg, SUBJECT_SCHEMAS_PATH, &query)
+        .await
+        .map_err(|e| anyhow::anyhow!("failed to get subject kafka schemas: {e:?}"))?;
+    formatter::output(cfg, &resp)
+}

src/commands/mod.rs

@@ -45,6 +45,7 @@ pub mod incidents;
 pub mod infrastructure;
 pub mod integrations;
 pub mod investigations;
+pub mod kafka;
 pub mod llm_obs;
 pub mod logs;
 pub mod logs_restriction;