feat(llm-obs): add dataset batch-update, clone, restore subcommands

docs/COMMANDS.md

@@ -66,7 +66,7 @@ pup <domain> <subgroup> <action> [options] # Nested commands
 | data-deletion | requests (list, create, cancel) | src/commands/data_deletion.rs | ✅ |
 | data-governance | scanner-rules (list) | src/commands/data_governance.rs | ✅ |
 | obs-pipelines | list, get, create, update, delete, validate | src/commands/obs_pipelines.rs | ✅ |
-| llm-obs | projects (create, list), experiments (create, list, update, delete, summary, events (list, get), metric-values, dimension-values), datasets (create, list), spans (search) | src/commands/llm_obs.rs | ✅ |
+| llm-obs | projects (create, list), experiments (create, list, update, delete, summary, events (list, get), metric-values, dimension-values), datasets (create, list, batch-update, clone, restore), spans (search) | src/commands/llm_obs.rs | ✅ |
 | reference-tables | list, get, create, batch-query | src/commands/reference_tables.rs | ✅ |
 | network | flows list, devices (list, get, interfaces, tags), interfaces (list, update) | src/commands/network.rs | ✅ |
 | cloud | aws, gcp, azure, oci | src/commands/cloud.rs | ✅ |
@@ -243,7 +243,7 @@ Available on all commands:
 
 ### v0.28.0 — New Command Groups and Full Pipeline Implementation
 
-- ✅ **llm-obs** (new) — LLM Observability: projects (create, list), experiments (create, list, update, delete, summary, events (list, get), metric-values, dimension-values), datasets (create, list), spans (search)
+- ✅ **llm-obs** (new) — LLM Observability: projects (create, list), experiments (create, list, update, delete, summary, events (list, get), metric-values, dimension-values), datasets (create, list, batch-update, clone, restore), spans (search)
 - ✅ **reference-tables** (new) — Reference table management (list, get, create, batch-query)
 - ✅ **obs-pipelines** (upgraded from placeholder) — Full CRUD: list, get, create, update, delete, validate
 - **costs** — Added cloud cost configs: `aws-config`, `azure-config`, `gcp-config` (list, get, create, delete each)

src/client.rs

@@ -350,7 +350,7 @@ static UNSTABLE_OPS: &[&str] = &[
     "v2.delete_aws_cloud_auth_persona_mapping",
     "v2.get_aws_cloud_auth_persona_mapping",
     "v2.list_aws_cloud_auth_persona_mappings",
-    // LLM Observability (18)
+    // LLM Observability (21)
     "v2.create_llm_obs_project",
     "v2.list_llm_obs_projects",
     "v2.create_llm_obs_experiment",
@@ -359,6 +359,9 @@ static UNSTABLE_OPS: &[&str] = &[
     "v2.delete_llm_obs_experiments",
     "v2.create_llm_obs_dataset",
     "v2.list_llm_obs_datasets",
+    "v2.batch_update_llm_obs_dataset",
+    "v2.clone_llm_obs_dataset",
+    "v2.restore_llm_obs_dataset_version",
     "v2.create_llm_obs_annotation_queue",
     "v2.list_llm_obs_annotation_queues",
     "v2.update_llm_obs_annotation_queue",
@@ -1273,7 +1276,7 @@ mod tests {
 
     #[test]
     fn test_unstable_ops_count() {
-        assert_eq!(UNSTABLE_OPS.len(), 162);
+        assert_eq!(UNSTABLE_OPS.len(), 165);
     }
 
     #[test]

src/commands/llm_obs.rs

@@ -4,9 +4,11 @@ use datadog_api_client::datadogV2::api_llm_observability::{
 };
 use datadog_api_client::datadogV2::model::{
     LLMObsAnnotationQueueInteractionsRequest, LLMObsAnnotationQueueRequest,
-    LLMObsAnnotationQueueUpdateRequest, LLMObsCustomEvalConfigUpdateRequest, LLMObsDatasetRequest,
-    LLMObsDeleteAnnotationQueueInteractionsRequest, LLMObsDeleteExperimentsRequest,
-    LLMObsExperimentRequest, LLMObsExperimentUpdateRequest, LLMObsProjectRequest,
+    LLMObsAnnotationQueueUpdateRequest, LLMObsCustomEvalConfigUpdateRequest,
+    LLMObsDatasetBatchUpdateRequest, LLMObsDatasetCloneRequest, LLMObsDatasetRequest,
+    LLMObsDatasetRestoreVersionRequest, LLMObsDeleteAnnotationQueueInteractionsRequest,
+    LLMObsDeleteExperimentsRequest, LLMObsExperimentRequest, LLMObsExperimentUpdateRequest,
+    LLMObsProjectRequest,
 };
 
 use crate::client;
@@ -108,6 +110,51 @@ pub async fn datasets_list(cfg: &Config, project_id: &str) -> Result<()> {
     formatter::output(cfg, &resp)
 }
 
+pub async fn datasets_batch_update(
+    cfg: &Config,
+    project_id: &str,
+    dataset_id: &str,
+    file: &str,
+) -> Result<()> {
+    let body: LLMObsDatasetBatchUpdateRequest = util::read_json_file(file)?;
+    let api = make_api(cfg);
+    let resp = api
+        .batch_update_llm_obs_dataset(project_id.to_string(), dataset_id.to_string(), body)
+        .await
+        .map_err(|e| anyhow::anyhow!("failed to batch update dataset records: {e:?}"))?;
+    formatter::output(cfg, &resp)
+}
+
+pub async fn datasets_clone(
+    cfg: &Config,
+    project_id: &str,
+    dataset_id: &str,
+    file: &str,
+) -> Result<()> {
+    let body: LLMObsDatasetCloneRequest = util::read_json_file(file)?;
+    let api = make_api(cfg);
+    let resp = api
+        .clone_llm_obs_dataset(project_id.to_string(), dataset_id.to_string(), body)
+        .await
+        .map_err(|e| anyhow::anyhow!("failed to clone dataset: {e:?}"))?;
+    formatter::output(cfg, &resp)
+}
+
+pub async fn datasets_restore(
+    cfg: &Config,
+    project_id: &str,
+    dataset_id: &str,
+    file: &str,
+) -> Result<()> {
+    let body: LLMObsDatasetRestoreVersionRequest = util::read_json_file(file)?;
+    let api = make_api(cfg);
+    api.restore_llm_obs_dataset_version(project_id.to_string(), dataset_id.to_string(), body)
+        .await
+        .map_err(|e| anyhow::anyhow!("failed to restore dataset version: {e:?}"))?;
+    println!("Dataset {dataset_id} restored.");
+    Ok(())
+}
+
 // ---- Experiment analytics (no typed equivalent — unstable MCP endpoints) ----
 
 pub async fn experiments_summary(cfg: &Config, experiment_id: &str) -> Result<()> {
@@ -2637,4 +2684,165 @@ mod tests {
         assert!(result.is_ok(), "spans_search failed: {:?}", result.err());
         cleanup_env();
     }
+
+    // ---- datasets_batch_update ----
+
+    #[tokio::test]
+    async fn test_llm_obs_datasets_batch_update() {
+        let _lock = lock_env().await;
+        std::env::set_var("DD_TOKEN_STORAGE", "file");
+        let mut server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+
+        let tmp = write_temp_json(
+            "pup_test_ds_batch_update.json",
+            r#"{"data":{"id":"ds-1","type":"datasets","attributes":{"insert_records":[],"delete_records":[]}}}"#,
+        );
+        let resp_body = r#"{"data":[]}"#;
+        let _mock = mock_any(&mut server, "POST", resp_body).await;
+
+        let result =
+            super::datasets_batch_update(&cfg, "proj-1", "ds-1", tmp.to_str().unwrap()).await;
+        assert!(
+            result.is_ok(),
+            "datasets_batch_update failed: {:?}",
+            result.err()
+        );
+        let _ = std::fs::remove_file(tmp);
+        cleanup_env();
+        std::env::remove_var("DD_TOKEN_STORAGE");
+    }
+
+    #[tokio::test]
+    async fn test_llm_obs_datasets_batch_update_400() {
+        let _lock = lock_env().await;
+        std::env::set_var("DD_TOKEN_STORAGE", "file");
+        let mut server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+
+        let tmp = write_temp_json(
+            "pup_test_ds_batch_update_400.json",
+            r#"{"data":{"id":"ds-1","type":"datasets","attributes":{"insert_records":[],"delete_records":[]}}}"#,
+        );
+        let _mock = server
+            .mock("POST", mockito::Matcher::Any)
+            .match_query(mockito::Matcher::Any)
+            .with_status(400)
+            .with_header("content-type", "application/json")
+            .with_body(r#"{"errors":["bad request"]}"#)
+            .create_async()
+            .await;
+
+        let result =
+            super::datasets_batch_update(&cfg, "proj-1", "ds-1", tmp.to_str().unwrap()).await;
+        assert!(result.is_err(), "should fail on 400");
+        let _ = std::fs::remove_file(tmp);
+        cleanup_env();
+        std::env::remove_var("DD_TOKEN_STORAGE");
+    }
+
+    // ---- datasets_clone ----
+
+    #[tokio::test]
+    async fn test_llm_obs_datasets_clone() {
+        let _lock = lock_env().await;
+        std::env::set_var("DD_TOKEN_STORAGE", "file");
+        let mut server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+
+        let tmp = write_temp_json(
+            "pup_test_ds_clone.json",
+            r#"{"data":{"id":"ds-1","type":"datasets","attributes":{"name":"cloned-dataset"}}}"#,
+        );
+        let resp_body = r#"{"data":{"id":"ds-2","type":"datasets","attributes":{"name":"cloned-dataset","description":null,"metadata":null,"created_at":"2024-01-01T00:00:00Z","updated_at":"2024-01-01T00:00:00Z","current_version":1}}}"#;
+        let _mock = mock_any(&mut server, "POST", resp_body).await;
+
+        let result = super::datasets_clone(&cfg, "proj-1", "ds-1", tmp.to_str().unwrap()).await;
+        assert!(result.is_ok(), "datasets_clone failed: {:?}", result.err());
+        let _ = std::fs::remove_file(tmp);
+        cleanup_env();
+        std::env::remove_var("DD_TOKEN_STORAGE");
+    }
+
+    #[tokio::test]
+    async fn test_llm_obs_datasets_clone_404() {
+        let _lock = lock_env().await;
+        std::env::set_var("DD_TOKEN_STORAGE", "file");
+        let mut server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+
+        let tmp = write_temp_json(
+            "pup_test_ds_clone_404.json",
+            r#"{"data":{"id":"ds-1","type":"datasets","attributes":{"name":"cloned-dataset"}}}"#,
+        );
+        let _mock = server
+            .mock("POST", mockito::Matcher::Any)
+            .match_query(mockito::Matcher::Any)
+            .with_status(404)
+            .with_header("content-type", "application/json")
+            .with_body(r#"{"errors":["not found"]}"#)
+            .create_async()
+            .await;
+
+        let result =
+            super::datasets_clone(&cfg, "proj-1", "ds-missing", tmp.to_str().unwrap()).await;
+        assert!(result.is_err(), "should fail on 404");
+        let _ = std::fs::remove_file(tmp);
+        cleanup_env();
+        std::env::remove_var("DD_TOKEN_STORAGE");
+    }
+
+    // ---- datasets_restore ----
+
+    #[tokio::test]
+    async fn test_llm_obs_datasets_restore() {
+        let _lock = lock_env().await;
+        std::env::set_var("DD_TOKEN_STORAGE", "file");
+        let mut server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+
+        let tmp = write_temp_json(
+            "pup_test_ds_restore.json",
+            r#"{"data":{"id":"ds-1","type":"datasets","attributes":{"dataset_version":2}}}"#,
+        );
+        let resp_body = r#"{"data":{"id":"ds-1","type":"datasets","attributes":{"name":"my-dataset","description":null,"metadata":null,"created_at":"2024-01-01T00:00:00Z","updated_at":"2024-01-01T00:00:00Z","current_version":2}}}"#;
+        let _mock = mock_any(&mut server, "POST", resp_body).await;
+
+        let result = super::datasets_restore(&cfg, "proj-1", "ds-1", tmp.to_str().unwrap()).await;
+        assert!(
+            result.is_ok(),
+            "datasets_restore failed: {:?}",
+            result.err()
+        );
+        let _ = std::fs::remove_file(tmp);
+        cleanup_env();
+        std::env::remove_var("DD_TOKEN_STORAGE");
+    }
+
+    #[tokio::test]
+    async fn test_llm_obs_datasets_restore_400() {
+        let _lock = lock_env().await;
+        std::env::set_var("DD_TOKEN_STORAGE", "file");
+        let mut server = mockito::Server::new_async().await;
+        let cfg = test_config(&server.url());
+
+        let tmp = write_temp_json(
+            "pup_test_ds_restore_400.json",
+            r#"{"data":{"id":"ds-1","type":"datasets","attributes":{"dataset_version":99}}}"#,
+        );
+        let _mock = server
+            .mock("POST", mockito::Matcher::Any)
+            .match_query(mockito::Matcher::Any)
+            .with_status(400)
+            .with_header("content-type", "application/json")
+            .with_body(r#"{"errors":["invalid version"]}"#)
+            .create_async()
+            .await;
+
+        let result = super::datasets_restore(&cfg, "proj-1", "ds-1", tmp.to_str().unwrap()).await;
+        assert!(result.is_err(), "should fail on 400");
+        let _ = std::fs::remove_file(tmp);
+        cleanup_env();
+        std::env::remove_var("DD_TOKEN_STORAGE");
+    }
 }

src/main.rs

@@ -8682,6 +8682,33 @@ enum LlmObsDatasetsActions {
         #[arg(long, help = "Project ID (required)")]
         project_id: String,
     },
+    /// Batch insert, update, and delete records in a dataset
+    BatchUpdate {
+        #[arg(long, help = "Project ID (required)")]
+        project_id: String,
+        #[arg(long, help = "Dataset ID (required)")]
+        dataset_id: String,
+        #[arg(long, help = "JSON file with batch update body (required)")]
+        file: String,
+    },
+    /// Clone a dataset into a new dataset
+    Clone {
+        #[arg(long, help = "Project ID (required)")]
+        project_id: String,
+        #[arg(long, help = "Dataset ID to clone (required)")]
+        dataset_id: String,
+        #[arg(long, help = "JSON file with clone body (required)")]
+        file: String,
+    },
+    /// Restore a dataset to a previous version
+    Restore {
+        #[arg(long, help = "Project ID (required)")]
+        project_id: String,
+        #[arg(long, help = "Dataset ID (required)")]
+        dataset_id: String,
+        #[arg(long, help = "JSON file with restore version body (required)")]
+        file: String,
+    },
 }
 
 #[derive(Subcommand)]
@@ -14661,6 +14688,35 @@ async fn main_inner() -> anyhow::Result<()> {
                     LlmObsDatasetsActions::List { project_id } => {
                         commands::llm_obs::datasets_list(&cfg, &project_id).await?;
                     }
+                    LlmObsDatasetsActions::BatchUpdate {
+                        project_id,
+                        dataset_id,
+                        file,
+                    } => {
+                        commands::llm_obs::datasets_batch_update(
+                            &cfg,
+                            &project_id,
+                            &dataset_id,
+                            &file,
+                        )
+                        .await?;
+                    }
+                    LlmObsDatasetsActions::Clone {
+                        project_id,
+                        dataset_id,
+                        file,
+                    } => {
+                        commands::llm_obs::datasets_clone(&cfg, &project_id, &dataset_id, &file)
+                            .await?;
+                    }
+                    LlmObsDatasetsActions::Restore {
+                        project_id,
+                        dataset_id,
+                        file,
+                    } => {
+                        commands::llm_obs::datasets_restore(&cfg, &project_id, &dataset_id, &file)
+                            .await?;
+                    }
                 },
                 LlmObsActions::Spans { action } => match action {
                     LlmObsSpansActions::Search {