Skip to main content

datalab_cli/commands/
create_document.rs

1use crate::cache::Cache;
2use crate::client::DatalabClient;
3use crate::error::{DatalabError, Result};
4use crate::output::Progress;
5use base64::Engine;
6use clap::Args;
7use reqwest::multipart::Form;
8use serde_json::json;
9use std::fs;
10use std::path::PathBuf;
11
12#[derive(Args, Debug)]
13pub struct CreateDocumentArgs {
14    /// Markdown content or path to markdown file
15    #[arg(long, value_name = "MARKDOWN")]
16    pub markdown: String,
17
18    /// Output format (currently only docx supported)
19    #[arg(
20        long,
21        default_value = "docx",
22        value_name = "FORMAT",
23        help_heading = "Output Options"
24    )]
25    pub output_format: String,
26
27    /// Skip local cache lookup
28    #[arg(long, help_heading = "Cache Options")]
29    pub skip_cache: bool,
30
31    /// Write created document to file (binary output)
32    #[arg(long, short, value_name = "FILE", help_heading = "Output Options")]
33    pub output: Option<PathBuf>,
34
35    /// Request timeout in seconds
36    #[arg(
37        long,
38        default_value = "300",
39        value_name = "SECS",
40        help_heading = "Advanced Options"
41    )]
42    pub timeout: u64,
43}
44
45impl CreateDocumentArgs {
46    fn get_markdown(&self) -> Result<String> {
47        let md_path = PathBuf::from(&self.markdown);
48        if md_path.exists() {
49            Ok(fs::read_to_string(&md_path)?)
50        } else {
51            Ok(self.markdown.clone())
52        }
53    }
54
55    fn to_cache_params(&self, markdown: &str) -> serde_json::Value {
56        use sha2::{Digest, Sha256};
57        let mut hasher = Sha256::new();
58        hasher.update(markdown.as_bytes());
59        let markdown_hash = hex::encode(hasher.finalize());
60
61        json!({
62            "markdown_hash": markdown_hash,
63            "output_format": self.output_format,
64        })
65    }
66}
67
68pub async fn execute(args: CreateDocumentArgs, progress: &Progress) -> Result<()> {
69    let client = DatalabClient::new(Some(args.timeout))?;
70    let cache = Cache::new()?;
71
72    let markdown = args.get_markdown()?;
73
74    progress.start("create-document", None);
75
76    let cache_params = args.to_cache_params(&markdown);
77    let cache_key = Cache::generate_key(None, None, "create-document", &cache_params);
78
79    if !args.skip_cache {
80        if let Some(cached) = cache.get(&cache_key) {
81            progress.cache_hit(&cache_key);
82            output_result(&cached, args.output.as_ref())?;
83            return Ok(());
84        }
85    }
86
87    let form = Form::new()
88        .text("markdown", markdown)
89        .text("output_format", args.output_format.clone());
90
91    let result = client
92        .submit_and_poll("create-document", form, progress)
93        .await?;
94
95    cache.set(&cache_key, &result, "create-document", None, None)?;
96
97    output_result(&result, args.output.as_ref())?;
98
99    Ok(())
100}
101
102fn output_result(result: &serde_json::Value, output_file: Option<&PathBuf>) -> Result<()> {
103    if let Some(path) = output_file {
104        if let Some(base64_data) = result.get("output_base64").and_then(|v| v.as_str()) {
105            let decoded = base64::engine::general_purpose::STANDARD
106                .decode(base64_data)
107                .map_err(|e| DatalabError::InvalidInput(format!("Invalid base64: {}", e)))?;
108            fs::write(path, &decoded)?;
109
110            let mut meta = result.clone();
111            meta.as_object_mut().map(|o| o.remove("output_base64"));
112            println!("{}", serde_json::to_string_pretty(&meta)?);
113        } else {
114            println!("{}", serde_json::to_string_pretty(result)?);
115        }
116    } else {
117        println!("{}", serde_json::to_string_pretty(result)?);
118    }
119
120    Ok(())
121}