refactor(datasets): unify create paths with shared create_dataset helper

Author: zfarrell

src/datasets.rs

@@ -220,13 +220,12 @@ fn upload_from_stdin(
     (id, ft.format)
 }
 
-pub fn create(
+fn create_dataset(
     workspace_id: &str,
-    label: Option<&str>,
+    label: &str,
     table_name: Option<&str>,
-    file: Option<&str>,
-    upload_id: Option<&str>,
-    source_format: &str,
+    source: serde_json::Value,
+    on_failure: Option<Box<dyn FnOnce()>>,
 ) {
     let profile_config = match config::load("default") {
         Ok(c) => c,
@@ -244,62 +243,13 @@ pub fn create(
         }
     };
 
-    let label_derived;
-    let label: &str = match label {
-        Some(l) => l,
-        None => match file {
-            Some(path) => {
-                label_derived = Path::new(path)
-                    .file_stem()
-                    .and_then(|s| s.to_str())
-                    .unwrap_or("dataset")
-                    .to_string();
-                &label_derived
-            }
-            None => {
-                if upload_id.is_some() {
-                    eprintln!("error: no label provided. Use --label to name the dataset.");
-                    std::process::exit(1);
-                }
-                match stdin_redirect_filename() {
-                    Some(name) => {
-                        label_derived = name;
-                        &label_derived
-                    }
-                    None => {
-                        eprintln!("error: no label provided. Use --label to name the dataset.");
-                        std::process::exit(1);
-                    }
-                }
-            }
-        },
-    };
-
-    let client = reqwest::blocking::Client::new();
-
-    let (upload_id, format, upload_id_was_uploaded): (String, &str, bool) = if let Some(id) = upload_id {
-        (id.to_string(), source_format, false)
-    } else {
-        let (id, fmt) = match file {
-            Some(path) => upload_from_file(&client, &api_key, workspace_id, &profile_config.api_url, path),
-            None => {
-                use std::io::IsTerminal;
-                if std::io::stdin().is_terminal() {
-                    eprintln!("error: no input data. Use --file <path>, --upload-id <id>, or pipe data via stdin.");
-                    std::process::exit(1);
-                }
-                upload_from_stdin(&client, &api_key, workspace_id, &profile_config.api_url)
-            }
-        };
-        (id, fmt, true)
-    };
-
-    let source = json!({ "upload_id": upload_id, "format": format });
     let mut body = json!({ "label": label, "source": source });
     if let Some(tn) = table_name {
         body["table_name"] = json!(tn);
     }
+
     let url = format!("{}/datasets", profile_config.api_url);
+    let client = reqwest::blocking::Client::new();
 
     let resp = match client
         .post(&url)
@@ -318,15 +268,8 @@ pub fn create(
     if !resp.status().is_success() {
         use crossterm::style::Stylize;
         eprintln!("{}", crate::util::api_error(resp.text().unwrap_or_default()).red());
-        // Only show the resume hint when the upload_id came from a fresh upload
-        if upload_id_was_uploaded {
-            eprintln!(
-                "{}",
-                format!(
-                    "Resume dataset creation without re-uploading by passing --upload-id {upload_id}"
-                )
-                .yellow()
-            );
+        if let Some(f) = on_failure {
+            f();
         }
         std::process::exit(1);
     }
@@ -346,12 +289,13 @@ pub fn create(
     println!("full_name:  datasets.main.{}", dataset.table_name);
 }
 
-fn create_from_source(
+pub fn create_from_upload(
     workspace_id: &str,
-    source: serde_json::Value,
     label: Option<&str>,
     table_name: Option<&str>,
-    label_required_hint: &str,
+    file: Option<&str>,
+    upload_id: Option<&str>,
+    source_format: &str,
 ) {
     let profile_config = match config::load("default") {
         Ok(c) => c,
@@ -369,55 +313,72 @@ fn create_from_source(
         }
     };
 
-    let label = match label {
+    let label_derived;
+    let label: &str = match label {
         Some(l) => l,
-        None => {
-            eprintln!("error: --label is required when using {label_required_hint}");
-            std::process::exit(1);
-        }
+        None => match file {
+            Some(path) => {
+                label_derived = Path::new(path)
+                    .file_stem()
+                    .and_then(|s| s.to_str())
+                    .unwrap_or("dataset")
+                    .to_string();
+                &label_derived
+            }
+            None => {
+                if upload_id.is_some() {
+                    eprintln!("error: no label provided. Use --label to name the dataset.");
+                    std::process::exit(1);
+                }
+                match stdin_redirect_filename() {
+                    Some(name) => {
+                        label_derived = name;
+                        &label_derived
+                    }
+                    None => {
+                        eprintln!("error: no label provided. Use --label to name the dataset.");
+                        std::process::exit(1);
+                    }
+                }
+            }
+        },
     };
 
-    let mut body = json!({ "label": label, "source": source });
-    if let Some(tn) = table_name {
-        body["table_name"] = json!(tn);
-    }
-
-    let url = format!("{}/datasets", profile_config.api_url);
     let client = reqwest::blocking::Client::new();
 
-    let resp = match client
-        .post(&url)
-        .header("Authorization", format!("Bearer {api_key}"))
-        .header("X-Workspace-Id", workspace_id)
-        .json(&body)
-        .send()
-    {
-        Ok(r) => r,
-        Err(e) => {
-            eprintln!("error connecting to API: {e}");
-            std::process::exit(1);
-        }
+    let (upload_id, format, upload_id_was_uploaded): (String, &str, bool) = if let Some(id) = upload_id {
+        (id.to_string(), source_format, false)
+    } else {
+        let (id, fmt) = match file {
+            Some(path) => upload_from_file(&client, &api_key, workspace_id, &profile_config.api_url, path),
+            None => {
+                use std::io::IsTerminal;
+                if std::io::stdin().is_terminal() {
+                    eprintln!("error: no input data. Use --file <path>, --upload-id <id>, or pipe data via stdin.");
+                    std::process::exit(1);
+                }
+                upload_from_stdin(&client, &api_key, workspace_id, &profile_config.api_url)
+            }
+        };
+        (id, fmt, true)
     };
 
-    if !resp.status().is_success() {
-        use crossterm::style::Stylize;
-        eprintln!("{}", crate::util::api_error(resp.text().unwrap_or_default()).red());
-        std::process::exit(1);
-    }
+    let source = json!({ "upload_id": upload_id, "format": format });
 
-    let dataset: CreateResponse = match resp.json() {
-        Ok(v) => v,
-        Err(e) => {
-            eprintln!("error parsing response: {e}");
-            std::process::exit(1);
-        }
+    let on_failure: Option<Box<dyn FnOnce()>> = if upload_id_was_uploaded {
+        let uid = upload_id.clone();
+        Some(Box::new(move || {
+            use crossterm::style::Stylize;
+            eprintln!(
+                "{}",
+                format!("Resume dataset creation without re-uploading by passing --upload-id {uid}").yellow()
+            );
+        }))
+    } else {
+        None
     };
 
-    use crossterm::style::Stylize;
-    println!("{}", "Dataset created".green());
-    println!("id:         {}", dataset.id);
-    println!("label:      {}", dataset.label);
-    println!("full_name:  datasets.main.{}", dataset.table_name);
+    create_dataset(workspace_id, label, table_name, source, on_failure);
 }
 
 pub fn create_from_query(
@@ -426,7 +387,14 @@ pub fn create_from_query(
     label: Option<&str>,
     table_name: Option<&str>,
 ) {
-    create_from_source(workspace_id, json!({ "sql": sql }), label, table_name, "--sql");
+    let label = match label {
+        Some(l) => l,
+        None => {
+            eprintln!("error: --label is required when using --sql");
+            std::process::exit(1);
+        }
+    };
+    create_dataset(workspace_id, label, table_name, json!({ "sql": sql }), None);
 }
 
 pub fn create_from_saved_query(
@@ -435,7 +403,14 @@ pub fn create_from_saved_query(
     label: Option<&str>,
     table_name: Option<&str>,
 ) {
-    create_from_source(workspace_id, json!({ "saved_query_id": query_id }), label, table_name, "--query-id");
+    let label = match label {
+        Some(l) => l,
+        None => {
+            eprintln!("error: --label is required when using --query-id");
+            std::process::exit(1);
+        }
+    };
+    create_dataset(workspace_id, label, table_name, json!({ "saved_query_id": query_id }), None);
 }
 
 pub fn list(workspace_id: &str, limit: Option<u32>, offset: Option<u32>, format: &str) {

src/main.rs

@@ -77,7 +77,7 @@ fn main() {
                             } else if let Some(query_id) = query_id {
                                 datasets::create_from_saved_query(&workspace_id, &query_id, label.as_deref(), table_name.as_deref())
                             } else {
-                                datasets::create(&workspace_id, label.as_deref(), table_name.as_deref(), file.as_deref(), upload_id.as_deref(), &format)
+                                datasets::create_from_upload(&workspace_id, label.as_deref(), table_name.as_deref(), file.as_deref(), upload_id.as_deref(), &format)
                             }
                         }
                         None => {