From 3a9a234ef2a220f4f0974cc02059387076a26b87 Mon Sep 17 00:00:00 2001 From: Rachel Yang Date: Tue, 2 Jun 2026 13:59:53 -0400 Subject: [PATCH 1/2] feat(apm): add sampling-rules CRUD commands MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds `pup apm sampling-rules {list,get,create,update,delete}` for managing customer per-(service,env,resource) head-based sampling rules. Backed by Remote Config product APM_TRACING with provenance=customer. Rules surface on traces with `_dd.p.dm:-11` and `ingestion_reason:remote_rule`. ## Commands pup apm sampling-rules list [--service S --env E] pup apm sampling-rules get pup apm sampling-rules create --service S --env E --resource R --sample-rate F pup apm sampling-rules update --service S --env E --resource R --sample-rate F pup apm sampling-rules delete When `--service` and `--env` are both given, list narrows via the backend's `/configs/by_target` endpoint. ## Scopes - apm_remote_configuration_read (added to default + read_only) - apm_remote_configuration_write (added to default) Both scopes are on pup's OUISvc registrable client allowlist as of the recent OAuth rollout. Backend support is in place: - dd-source PR `rachel.yang/apm-trace-configurations-oauth-auth` (Rapid service-default OAuth on trace-configurations) - dd-go PR `rachel.yang/rc-api-proxy-apm-tracing-oauth` (rc-api-proxy authn_methods on apm_tracing routes) ## Tests 8 tests added in `commands::apm::tests`, matching the service_remapping_* pattern: test_sampling_rules_list test_sampling_rules_list_by_target test_sampling_rules_get test_sampling_rules_get_not_found test_sampling_rules_create test_sampling_rules_create_api_error test_sampling_rules_update test_sampling_rules_delete All pass locally. The 5 unrelated test failures in `cases`, `dbm`, `logs`, `metrics`, `traces` are pre-existing DNS-resolution flakes against `unused.local` and not touched by this change. ## Follow-ups - `pup apm adaptive-sampling` commands (allotment + onboarding) — separate PR. Requires the `apm_service_ingest_{read,write}` scopes which are also already on the OUISvc allowlist. - Flip the agent-skills `dd-apm/sampling` skill PR from draft to ready once these commands ship in a release. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/auth/types.rs | 3 + src/commands/apm.rs | 333 ++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 100 +++++++++++++ 3 files changed, 436 insertions(+) diff --git a/src/auth/types.rs b/src/auth/types.rs index 168e7799..60814f5c 100644 --- a/src/auth/types.rs +++ b/src/auth/types.rs @@ -49,6 +49,7 @@ pub fn all_known_scopes() -> Vec<&'static str> { pub fn read_only_scopes() -> Vec<&'static str> { vec![ "apm_read", + "apm_remote_configuration_read", "apm_service_catalog_read", "audit_logs_read", "aws_configuration_read", @@ -106,6 +107,8 @@ pub fn default_scopes() -> Vec<&'static str> { vec![ // APM "apm_read", + "apm_remote_configuration_read", + "apm_remote_configuration_write", "apm_service_catalog_read", "apm_service_renaming_write", // Audit diff --git a/src/commands/apm.rs b/src/commands/apm.rs index 4fe973cc..67c2d7dc 100644 --- a/src/commands/apm.rs +++ b/src/commands/apm.rs @@ -161,6 +161,110 @@ pub async fn service_remapping_delete(cfg: &Config, id: String, version: i64) -> client::raw_delete(cfg, &format!("/api/v2/service-naming-rules/{id}/{version}")).await } +// ============================================================================= +// APM sampling rules — customer per-(service, env) resource sampling rules. +// Backed by RC product APM_TRACING (provenance:customer). These rules surface +// on traces with `_dd.p.dm:-11` and `ingestion_reason:remote_rule`. +// ============================================================================= + +const SAMPLING_RULES_BASE: &str = "/api/unstable/remote_config/products/apm_tracing/configs"; + +pub async fn sampling_rules_list( + cfg: &Config, + service: Option, + env: Option, +) -> Result<()> { + // If service + env are both given, prefer the narrowed by_target endpoint. + if let (Some(svc), Some(e)) = (service.as_deref(), env.as_deref()) { + let path = format!("{SAMPLING_RULES_BASE}/by_target"); + let data = client::raw_get(cfg, &path, &[("service", svc), ("env", e)]).await?; + return formatter::output(cfg, &data); + } + let data = client::raw_get(cfg, SAMPLING_RULES_BASE, &[]).await?; + formatter::output(cfg, &data) +} + +pub async fn sampling_rules_get(cfg: &Config, id: String) -> Result<()> { + let data = client::raw_get(cfg, &format!("{SAMPLING_RULES_BASE}/{id}"), &[]).await?; + formatter::output(cfg, &data) +} + +pub async fn sampling_rules_create( + cfg: &Config, + service: String, + env: String, + resource: String, + sample_rate: f64, +) -> Result<()> { + let body = serde_json::json!({ + "data": { + "type": "apm_tracing_config", + "attributes": { + "action": "enable", + "lib_config": { + "library_language": "all", + "library_version": "latest", + "service_name": service, + "env": env, + "tracing_sampling_rules": [{ + "service": service, + "provenance": "customer", + "resource": resource, + "sample_rate": sample_rate, + }], + }, + "service_target": { + "service": service, + "env": env, + }, + } + } + }); + let data = client::raw_post(cfg, SAMPLING_RULES_BASE, body).await?; + formatter::output(cfg, &data) +} + +pub async fn sampling_rules_update( + cfg: &Config, + id: String, + service: String, + env: String, + resource: String, + sample_rate: f64, +) -> Result<()> { + let body = serde_json::json!({ + "data": { + "id": id, + "type": "apm_tracing_config", + "attributes": { + "action": "enable", + "lib_config": { + "library_language": "all", + "library_version": "latest", + "service_name": service, + "env": env, + "tracing_sampling_rules": [{ + "service": service, + "provenance": "customer", + "resource": resource, + "sample_rate": sample_rate, + }], + }, + "service_target": { + "service": service, + "env": env, + }, + } + } + }); + let data = client::raw_put(cfg, &format!("{SAMPLING_RULES_BASE}/{id}"), body).await?; + formatter::output(cfg, &data) +} + +pub async fn sampling_rules_delete(cfg: &Config, id: String) -> Result<()> { + client::raw_delete(cfg, &format!("{SAMPLING_RULES_BASE}/{id}")).await +} + pub async fn service_config_get( cfg: &Config, service_name: String, @@ -721,4 +825,233 @@ mod tests { mock.assert_async().await; cleanup_env(); } + + // ===== sampling rules ===== + + #[tokio::test] + async fn test_sampling_rules_list() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let mock = server + .mock( + "GET", + "/api/unstable/remote_config/products/apm_tracing/configs", + ) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(r#"{"data": []}"#) + .create_async() + .await; + + let result = super::sampling_rules_list(&cfg, None, None).await; + assert!( + result.is_ok(), + "sampling_rules_list failed: {:?}", + result.err() + ); + mock.assert_async().await; + cleanup_env(); + } + + #[tokio::test] + async fn test_sampling_rules_list_by_target() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let mock = server + .mock( + "GET", + "/api/unstable/remote_config/products/apm_tracing/configs/by_target?service=api&env=prod", + ) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(r#"{"data": []}"#) + .create_async() + .await; + + let result = + super::sampling_rules_list(&cfg, Some("api".into()), Some("prod".into())).await; + assert!( + result.is_ok(), + "sampling_rules_list by_target failed: {:?}", + result.err() + ); + mock.assert_async().await; + cleanup_env(); + } + + #[tokio::test] + async fn test_sampling_rules_get() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let mock = server + .mock( + "GET", + "/api/unstable/remote_config/products/apm_tracing/configs/abc123", + ) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(r#"{"data": {"id": "abc123"}}"#) + .create_async() + .await; + + let result = super::sampling_rules_get(&cfg, "abc123".into()).await; + assert!( + result.is_ok(), + "sampling_rules_get failed: {:?}", + result.err() + ); + mock.assert_async().await; + cleanup_env(); + } + + #[tokio::test] + async fn test_sampling_rules_get_not_found() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + server + .mock( + "GET", + "/api/unstable/remote_config/products/apm_tracing/configs/missing", + ) + .with_status(404) + .with_header("content-type", "application/json") + .with_body(r#"{"errors": ["Not Found"]}"#) + .create_async() + .await; + + let result = super::sampling_rules_get(&cfg, "missing".into()).await; + assert!(result.is_err(), "expected error on 404"); + cleanup_env(); + } + + #[tokio::test] + async fn test_sampling_rules_create() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let mock = server + .mock( + "POST", + "/api/unstable/remote_config/products/apm_tracing/configs", + ) + .with_status(201) + .with_header("content-type", "application/json") + .with_body(r#"{"data": {"id": "new-config-id"}}"#) + .create_async() + .await; + + let result = super::sampling_rules_create( + &cfg, + "api".into(), + "prod".into(), + "*".into(), + 0.1, + ) + .await; + assert!( + result.is_ok(), + "sampling_rules_create failed: {:?}", + result.err() + ); + mock.assert_async().await; + cleanup_env(); + } + + #[tokio::test] + async fn test_sampling_rules_create_api_error() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + server + .mock( + "POST", + "/api/unstable/remote_config/products/apm_tracing/configs", + ) + .with_status(422) + .with_header("content-type", "application/json") + .with_body(r#"{"errors": ["Invalid sample_rate"]}"#) + .create_async() + .await; + + let result = super::sampling_rules_create( + &cfg, + "api".into(), + "prod".into(), + "*".into(), + -1.0, + ) + .await; + assert!(result.is_err(), "expected error on 422"); + cleanup_env(); + } + + #[tokio::test] + async fn test_sampling_rules_update() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let mock = server + .mock( + "PUT", + "/api/unstable/remote_config/products/apm_tracing/configs/abc123", + ) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(r#"{"data": {"id": "abc123"}}"#) + .create_async() + .await; + + let result = super::sampling_rules_update( + &cfg, + "abc123".into(), + "api".into(), + "prod".into(), + "*".into(), + 0.5, + ) + .await; + assert!( + result.is_ok(), + "sampling_rules_update failed: {:?}", + result.err() + ); + mock.assert_async().await; + cleanup_env(); + } + + #[tokio::test] + async fn test_sampling_rules_delete() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let mock = server + .mock( + "DELETE", + "/api/unstable/remote_config/products/apm_tracing/configs/abc123", + ) + .with_status(204) + .create_async() + .await; + + let result = super::sampling_rules_delete(&cfg, "abc123".into()).await; + assert!( + result.is_ok(), + "sampling_rules_delete failed: {:?}", + result.err() + ); + mock.assert_async().await; + cleanup_env(); + } } diff --git a/src/main.rs b/src/main.rs index 0e2d7dbf..7be14426 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8088,6 +8088,14 @@ enum ApmActions { #[command(subcommand)] action: ApmServiceRemappingActions, }, + /// Manage APM customer sampling rules (per-service per-resource head-based sampling rates). + /// Backed by Remote Config product APM_TRACING with provenance=customer. + /// Rules show on traces with `_dd.p.dm:-11` and `ingestion_reason:remote_rule`. + #[command(name = "sampling-rules")] + SamplingRules { + #[command(subcommand)] + action: ApmSamplingRulesActions, + }, /// View APM service instance configuration #[command(name = "service-config")] ServiceConfig { @@ -8273,6 +8281,55 @@ enum ApmServiceRemappingActions { }, } +#[derive(Subcommand)] +enum ApmSamplingRulesActions { + /// List sampling rules. With `--service` + `--env`, narrows to that target. + List { + #[arg(long, help = "Filter by service name (must be combined with --env)")] + service: Option, + #[arg(long, help = "Filter by environment (must be combined with --service)")] + env: Option, + }, + /// Get a sampling rule config by ID + Get { + #[arg(help = "Config ID")] + id: String, + }, + /// Create a customer sampling rule for (service, env, resource). + /// Rate is between 0.0 and 1.0. Anything > 1e-6 is honored. + Create { + #[arg(long, help = "Service name (required)")] + service: String, + #[arg(long, help = "Environment (required, must match DD_ENV on the service)")] + env: String, + #[arg( + long, + help = "Resource glob — `*` matches all resources for the service, or e.g. 'GET /api/users'" + )] + resource: String, + #[arg(long, help = "Sample rate between 0.0 and 1.0")] + sample_rate: f64, + }, + /// Update an existing sampling rule by ID (replaces all attributes) + Update { + #[arg(help = "Config ID")] + id: String, + #[arg(long, help = "Service name")] + service: String, + #[arg(long, help = "Environment")] + env: String, + #[arg(long, help = "Resource glob")] + resource: String, + #[arg(long, help = "Sample rate between 0.0 and 1.0")] + sample_rate: f64, + }, + /// Delete a sampling rule by ID + Delete { + #[arg(help = "Config ID")] + id: String, + }, +} + #[derive(Subcommand)] enum ApmServiceConfigActions { /// Get service instance configuration. @@ -14259,6 +14316,49 @@ async fn main_inner() -> anyhow::Result<()> { commands::apm::service_remapping_delete(&cfg, id, version).await?; } }, + ApmActions::SamplingRules { action } => match action { + ApmSamplingRulesActions::List { service, env } => { + commands::apm::sampling_rules_list(&cfg, service, env).await?; + } + ApmSamplingRulesActions::Get { id } => { + commands::apm::sampling_rules_get(&cfg, id).await?; + } + ApmSamplingRulesActions::Create { + service, + env, + resource, + sample_rate, + } => { + commands::apm::sampling_rules_create( + &cfg, + service, + env, + resource, + sample_rate, + ) + .await?; + } + ApmSamplingRulesActions::Update { + id, + service, + env, + resource, + sample_rate, + } => { + commands::apm::sampling_rules_update( + &cfg, + id, + service, + env, + resource, + sample_rate, + ) + .await?; + } + ApmSamplingRulesActions::Delete { id } => { + commands::apm::sampling_rules_delete(&cfg, id).await?; + } + }, ApmActions::ServiceConfig { action } => match action { ApmServiceConfigActions::Get { service_name, From 8cfbaf5333ce7acd4643d8cb3883a8a7e0aa62de Mon Sep 17 00:00:00 2001 From: Rachel Yang Date: Tue, 2 Jun 2026 14:13:43 -0400 Subject: [PATCH 2/2] feat(apm): add adaptive-sampling commands MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds `pup apm adaptive-sampling` subcommands for managing Datadog's adaptive sampling (server-driven dynamic sampling) — onboarding services and configuring the monthly byte/percent allotment. Datadog auto-tunes per-(service, env, resource) sampling rates to fit the budget. Generated rules surface on traces with `_dd.p.dm:-12` and `ingestion_reason:adaptive_rule`. Builds on the sampling-rules PR (#PR_B1) and the same OAuth rollout. ## Commands pup apm adaptive-sampling onboarding-status [--service S --env E] pup apm adaptive-sampling onboard --service S --env E pup apm adaptive-sampling offboard --service S --env E pup apm adaptive-sampling get-allotment pup apm adaptive-sampling set-allotment (--bytes N | --percent F) pup apm adaptive-sampling check pup apm adaptive-sampling preview (--bytes N | --percent F) `set-allotment` and `preview` take exactly one of `--bytes` or `--percent` (enforced via clap `conflicts_with`). Strategy is inferred: `fixed_target` for bytes, `percent_total` for percent. ## Scopes - apm_service_ingest_read (added to default + read_only) - apm_service_ingest_write (added to default) Both already on pup's OUISvc registrable client allowlist (added in the same rollout as `apm_remote_configuration_{read,write}` for the sampling-rules commands). The adaptive-sampling endpoints require BOTH `apm_service_ingest_*` AND `apm_remote_configuration_*` (RoutePermissions with operator AND) — the latter is added by the sampling-rules PR this builds on. ## Endpoints GET/POST /api/ui/adaptive_sampling/onboarding_status GET/POST /api/ui/adaptive_sampling/allotment_config POST /api/ui/adaptive_sampling/preview GET /api/ui/adaptive_sampling/allotment_check All in dd-source `trace-configurations` service, which got `ValidOAuthAccessToken` in the dd-source-side OAuth PR. ## Tests 9 tests added in `commands::apm::tests`: test_adaptive_sampling_onboarding_status_no_filter test_adaptive_sampling_onboarding_status_with_filter test_adaptive_sampling_onboard test_adaptive_sampling_offboard test_adaptive_sampling_get_allotment test_adaptive_sampling_set_allotment_bytes test_adaptive_sampling_set_allotment_percent test_adaptive_sampling_check test_adaptive_sampling_preview All pass locally. The same 5 pre-existing DNS-flakes in unrelated modules (cases/dbm/logs/metrics/traces) remain — not touched here. ## Follow-up Once both B1 (sampling-rules) and this PR (B2) ship in a pup release, flip the agent-skills `dd-apm/sampling` skill PR from draft to ready. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/auth/types.rs | 3 + src/commands/apm.rs | 374 ++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 102 ++++++++++++ 3 files changed, 479 insertions(+) diff --git a/src/auth/types.rs b/src/auth/types.rs index 60814f5c..698a361a 100644 --- a/src/auth/types.rs +++ b/src/auth/types.rs @@ -51,6 +51,7 @@ pub fn read_only_scopes() -> Vec<&'static str> { "apm_read", "apm_remote_configuration_read", "apm_service_catalog_read", + "apm_service_ingest_read", "audit_logs_read", "aws_configuration_read", "azure_configuration_read", @@ -110,6 +111,8 @@ pub fn default_scopes() -> Vec<&'static str> { "apm_remote_configuration_read", "apm_remote_configuration_write", "apm_service_catalog_read", + "apm_service_ingest_read", + "apm_service_ingest_write", "apm_service_renaming_write", // Audit "audit_logs_read", diff --git a/src/commands/apm.rs b/src/commands/apm.rs index 67c2d7dc..57cad2dc 100644 --- a/src/commands/apm.rs +++ b/src/commands/apm.rs @@ -265,6 +265,149 @@ pub async fn sampling_rules_delete(cfg: &Config, id: String) -> Result<()> { client::raw_delete(cfg, &format!("{SAMPLING_RULES_BASE}/{id}")).await } +// ============================================================================= +// APM adaptive sampling — let Datadog auto-tune per-resource sampling rates to +// fit a monthly byte/percent allotment. Generated rules surface on traces with +// `_dd.p.dm:-12` and `ingestion_reason:adaptive_rule`. +// +// Strategy values: +// - "fixed_target" — set a hard byte target (use with --bytes) +// - "percent_total" — set a percent of allotment cap (use with --percent) +// ============================================================================= + +const ADAPTIVE_SAMPLING_BASE: &str = "/api/ui/adaptive_sampling"; + +fn allotment_attributes(bytes: Option, percent: Option) -> serde_json::Value { + let strategy = if bytes.is_some() { + "fixed_target" + } else { + "percent_total" + }; + let mut attrs = serde_json::json!({ "strategy": strategy }); + if let Some(b) = bytes { + attrs["allotment_bytes"] = serde_json::json!(b); + } + if let Some(p) = percent { + attrs["allotment_percent"] = serde_json::json!(p); + } + attrs +} + +pub async fn adaptive_sampling_onboarding_status( + cfg: &Config, + service: Option, + env: Option, +) -> Result<()> { + let path = format!("{ADAPTIVE_SAMPLING_BASE}/onboarding_status"); + let mut params: Vec<(&str, &str)> = Vec::new(); + if let Some(s) = service.as_deref() { + params.push(("service", s)); + } + if let Some(e) = env.as_deref() { + params.push(("env", e)); + } + let data = client::raw_get(cfg, &path, ¶ms).await?; + formatter::output(cfg, &data) +} + +async fn post_onboarding( + cfg: &Config, + service: String, + env: String, + onboarded: bool, +) -> Result<()> { + let body = serde_json::json!({ + "data": { + "id": "1", + "type": "apm_adaptive_sampling_onboarding_status", + "attributes": { + "service": service, + "env": env, + "onboarded": onboarded, + } + } + }); + let data = client::raw_post( + cfg, + &format!("{ADAPTIVE_SAMPLING_BASE}/onboarding_status"), + body, + ) + .await?; + formatter::output(cfg, &data) +} + +pub async fn adaptive_sampling_onboard( + cfg: &Config, + service: String, + env: String, +) -> Result<()> { + post_onboarding(cfg, service, env, true).await +} + +pub async fn adaptive_sampling_offboard( + cfg: &Config, + service: String, + env: String, +) -> Result<()> { + post_onboarding(cfg, service, env, false).await +} + +pub async fn adaptive_sampling_get_allotment(cfg: &Config) -> Result<()> { + let path = format!("{ADAPTIVE_SAMPLING_BASE}/allotment_config"); + let data = client::raw_get(cfg, &path, &[]).await?; + formatter::output(cfg, &data) +} + +pub async fn adaptive_sampling_set_allotment( + cfg: &Config, + bytes: Option, + percent: Option, +) -> Result<()> { + let attrs = allotment_attributes(bytes, percent); + let body = serde_json::json!({ + "data": { + "id": "1", + "type": "apm_adaptive_sampling_allotment_config", + "attributes": attrs, + } + }); + let data = client::raw_post( + cfg, + &format!("{ADAPTIVE_SAMPLING_BASE}/allotment_config"), + body, + ) + .await?; + formatter::output(cfg, &data) +} + +pub async fn adaptive_sampling_check(cfg: &Config) -> Result<()> { + let path = format!("{ADAPTIVE_SAMPLING_BASE}/allotment_check"); + let data = client::raw_get(cfg, &path, &[]).await?; + formatter::output(cfg, &data) +} + +pub async fn adaptive_sampling_preview( + cfg: &Config, + bytes: Option, + percent: Option, +) -> Result<()> { + let attrs = allotment_attributes(bytes, percent); + let body = serde_json::json!({ + "data": { + "id": "1", + "type": "apm_adaptive_sampling_allotment_preview", + "attributes": attrs, + } + }); + let data = client::raw_post( + cfg, + &format!("{ADAPTIVE_SAMPLING_BASE}/preview"), + body, + ) + .await?; + formatter::output(cfg, &data) +} + pub async fn service_config_get( cfg: &Config, service_name: String, @@ -1054,4 +1197,235 @@ mod tests { mock.assert_async().await; cleanup_env(); } + + // ===== adaptive sampling ===== + + #[tokio::test] + async fn test_adaptive_sampling_onboarding_status_no_filter() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let mock = server + .mock("GET", "/api/ui/adaptive_sampling/onboarding_status") + .with_status(200) + .with_header("content-type", "application/json") + .with_body(r#"{"data": []}"#) + .create_async() + .await; + + let result = super::adaptive_sampling_onboarding_status(&cfg, None, None).await; + assert!( + result.is_ok(), + "adaptive_sampling_onboarding_status failed: {:?}", + result.err() + ); + mock.assert_async().await; + cleanup_env(); + } + + #[tokio::test] + async fn test_adaptive_sampling_onboarding_status_with_filter() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let mock = server + .mock( + "GET", + "/api/ui/adaptive_sampling/onboarding_status?service=api&env=prod", + ) + .with_status(200) + .with_header("content-type", "application/json") + .with_body(r#"{"data": {"onboarded": true}}"#) + .create_async() + .await; + + let result = super::adaptive_sampling_onboarding_status( + &cfg, + Some("api".into()), + Some("prod".into()), + ) + .await; + assert!( + result.is_ok(), + "adaptive_sampling_onboarding_status with filter failed: {:?}", + result.err() + ); + mock.assert_async().await; + cleanup_env(); + } + + #[tokio::test] + async fn test_adaptive_sampling_onboard() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let mock = server + .mock("POST", "/api/ui/adaptive_sampling/onboarding_status") + .with_status(200) + .with_header("content-type", "application/json") + .with_body(r#"{"data": {"onboarded": true}}"#) + .create_async() + .await; + + let result = + super::adaptive_sampling_onboard(&cfg, "api".into(), "prod".into()).await; + assert!( + result.is_ok(), + "adaptive_sampling_onboard failed: {:?}", + result.err() + ); + mock.assert_async().await; + cleanup_env(); + } + + #[tokio::test] + async fn test_adaptive_sampling_offboard() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let mock = server + .mock("POST", "/api/ui/adaptive_sampling/onboarding_status") + .with_status(200) + .with_header("content-type", "application/json") + .with_body(r#"{"data": {"onboarded": false}}"#) + .create_async() + .await; + + let result = + super::adaptive_sampling_offboard(&cfg, "api".into(), "prod".into()).await; + assert!( + result.is_ok(), + "adaptive_sampling_offboard failed: {:?}", + result.err() + ); + mock.assert_async().await; + cleanup_env(); + } + + #[tokio::test] + async fn test_adaptive_sampling_get_allotment() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let mock = server + .mock("GET", "/api/ui/adaptive_sampling/allotment_config") + .with_status(200) + .with_header("content-type", "application/json") + .with_body(r#"{"data": {"strategy": "fixed_target", "allotment_bytes": 100000}}"#) + .create_async() + .await; + + let result = super::adaptive_sampling_get_allotment(&cfg).await; + assert!( + result.is_ok(), + "adaptive_sampling_get_allotment failed: {:?}", + result.err() + ); + mock.assert_async().await; + cleanup_env(); + } + + #[tokio::test] + async fn test_adaptive_sampling_set_allotment_bytes() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let mock = server + .mock("POST", "/api/ui/adaptive_sampling/allotment_config") + .with_status(200) + .with_header("content-type", "application/json") + .with_body(r#"{"data": {}}"#) + .create_async() + .await; + + let result = + super::adaptive_sampling_set_allotment(&cfg, Some(100_000), None).await; + assert!( + result.is_ok(), + "adaptive_sampling_set_allotment with bytes failed: {:?}", + result.err() + ); + mock.assert_async().await; + cleanup_env(); + } + + #[tokio::test] + async fn test_adaptive_sampling_set_allotment_percent() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let mock = server + .mock("POST", "/api/ui/adaptive_sampling/allotment_config") + .with_status(200) + .with_header("content-type", "application/json") + .with_body(r#"{"data": {}}"#) + .create_async() + .await; + + let result = super::adaptive_sampling_set_allotment(&cfg, None, Some(50.0)).await; + assert!( + result.is_ok(), + "adaptive_sampling_set_allotment with percent failed: {:?}", + result.err() + ); + mock.assert_async().await; + cleanup_env(); + } + + #[tokio::test] + async fn test_adaptive_sampling_check() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let mock = server + .mock("GET", "/api/ui/adaptive_sampling/allotment_check") + .with_status(200) + .with_header("content-type", "application/json") + .with_body( + r#"{"data": {"allotment_bytes": 100000, "ingested_bytes": 50000, "projected_monthly_ingested_bytes": 150000}}"#, + ) + .create_async() + .await; + + let result = super::adaptive_sampling_check(&cfg).await; + assert!( + result.is_ok(), + "adaptive_sampling_check failed: {:?}", + result.err() + ); + mock.assert_async().await; + cleanup_env(); + } + + #[tokio::test] + async fn test_adaptive_sampling_preview() { + let _lock = lock_env().await; + let mut server = mockito::Server::new_async().await; + let cfg = test_config(&server.url()); + + let mock = server + .mock("POST", "/api/ui/adaptive_sampling/preview") + .with_status(200) + .with_header("content-type", "application/json") + .with_body(r#"{"data": {"monthly_quota": 100000, "monthly_target": 50000}}"#) + .create_async() + .await; + + let result = super::adaptive_sampling_preview(&cfg, Some(50_000), None).await; + assert!( + result.is_ok(), + "adaptive_sampling_preview failed: {:?}", + result.err() + ); + mock.assert_async().await; + cleanup_env(); + } } diff --git a/src/main.rs b/src/main.rs index 7be14426..4947dc5e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8096,6 +8096,14 @@ enum ApmActions { #[command(subcommand)] action: ApmSamplingRulesActions, }, + /// Manage APM adaptive sampling — onboard services and configure the monthly allotment. + /// Datadog auto-tunes per-resource sampling rates to fit the configured byte/percent budget. + /// Generated rules show on traces with `_dd.p.dm:-12` and `ingestion_reason:adaptive_rule`. + #[command(name = "adaptive-sampling")] + AdaptiveSampling { + #[command(subcommand)] + action: ApmAdaptiveSamplingActions, + }, /// View APM service instance configuration #[command(name = "service-config")] ServiceConfig { @@ -8330,6 +8338,69 @@ enum ApmSamplingRulesActions { }, } +#[derive(Subcommand)] +enum ApmAdaptiveSamplingActions { + /// Get the onboarding status. With `--service` and `--env`, returns one entry; otherwise lists all. + #[command(name = "onboarding-status")] + OnboardingStatus { + #[arg(long, help = "Filter by service name (optional)")] + service: Option, + #[arg(long, help = "Filter by environment (optional)")] + env: Option, + }, + /// Onboard a (service, env) pair to adaptive sampling + Onboard { + #[arg(long, help = "Service name (required)")] + service: String, + #[arg(long, help = "Environment (required)")] + env: String, + }, + /// Offboard a (service, env) pair from adaptive sampling + Offboard { + #[arg(long, help = "Service name (required)")] + service: String, + #[arg(long, help = "Environment (required)")] + env: String, + }, + /// Read the org's adaptive sampling allotment configuration + #[command(name = "get-allotment")] + GetAllotment, + /// Set the org's adaptive sampling allotment. Provide exactly one of --bytes or --percent. + #[command(name = "set-allotment")] + SetAllotment { + #[arg( + long, + conflicts_with = "percent", + help = "Monthly target in bytes (strategy=fixed_target)" + )] + bytes: Option, + #[arg( + long, + conflicts_with = "bytes", + help = "Percent of total monthly allotment (strategy=percent_total)" + )] + percent: Option, + }, + /// Check whether the configured allotment is sufficient for current ingestion + Check, + /// Preview the allotment Datadog would compute for a strategy without applying it. + /// Provide exactly one of --bytes or --percent. + Preview { + #[arg( + long, + conflicts_with = "percent", + help = "Monthly target in bytes (strategy=fixed_target)" + )] + bytes: Option, + #[arg( + long, + conflicts_with = "bytes", + help = "Percent of total monthly allotment (strategy=percent_total)" + )] + percent: Option, + }, +} + #[derive(Subcommand)] enum ApmServiceConfigActions { /// Get service instance configuration. @@ -14359,6 +14430,37 @@ async fn main_inner() -> anyhow::Result<()> { commands::apm::sampling_rules_delete(&cfg, id).await?; } }, + ApmActions::AdaptiveSampling { action } => match action { + ApmAdaptiveSamplingActions::OnboardingStatus { service, env } => { + commands::apm::adaptive_sampling_onboarding_status(&cfg, service, env) + .await?; + } + ApmAdaptiveSamplingActions::Onboard { service, env } => { + commands::apm::adaptive_sampling_onboard(&cfg, service, env).await?; + } + ApmAdaptiveSamplingActions::Offboard { service, env } => { + commands::apm::adaptive_sampling_offboard(&cfg, service, env).await?; + } + ApmAdaptiveSamplingActions::GetAllotment => { + commands::apm::adaptive_sampling_get_allotment(&cfg).await?; + } + ApmAdaptiveSamplingActions::SetAllotment { bytes, percent } => { + if bytes.is_none() && percent.is_none() { + anyhow::bail!("must provide --bytes or --percent"); + } + commands::apm::adaptive_sampling_set_allotment(&cfg, bytes, percent) + .await?; + } + ApmAdaptiveSamplingActions::Check => { + commands::apm::adaptive_sampling_check(&cfg).await?; + } + ApmAdaptiveSamplingActions::Preview { bytes, percent } => { + if bytes.is_none() && percent.is_none() { + anyhow::bail!("must provide --bytes or --percent"); + } + commands::apm::adaptive_sampling_preview(&cfg, bytes, percent).await?; + } + }, ApmActions::ServiceConfig { action } => match action { ApmServiceConfigActions::Get { service_name,