Add new edit command for editing workloads

src/config.rs

@@ -274,6 +274,14 @@ impl ValueEnum for Consistency {
     }
 }
 
+#[derive(Parser, Debug, Serialize, Deserialize)]
+#[command(next_line_help = true)]
+pub struct EditCommand {
+    /// Path to the workload definition file.
+    #[clap(name = "workload", required = true, value_name = "PATH")]
+    pub workload: PathBuf,
+}
+
 #[derive(Parser, Debug, Serialize, Deserialize)]
 #[command(next_line_help = true)]
 pub struct SchemaCommand {
@@ -501,6 +509,19 @@ pub struct PlotCommand {
 #[derive(Parser, Debug)]
 #[allow(clippy::large_enum_variant)]
 pub enum Command {
+    /// Opens the specified workload script file for editing.
+    ///
+    /// Searches for the script on the workload search path. Workload files
+    /// are first searched in the current working directory, next in the paths
+    /// specified by LATTE_WORKLOAD_PATH environment variable. If the variable
+    /// is not defined, `/usr/share/latte/workloads` and `.local/share/latte/workloads`
+    /// are searched.
+    ///
+    /// Opens the editor pointed by LATTE_EDITOR or EDITOR environment variable.
+    /// If no variable is set, tries to launch vi.
+    ///
+    Edit(EditCommand),
+
     /// Creates the database schema by invoking the `schema` function of the workload script.
     ///
     /// The function should remove the old schema if present.

src/error.rs

@@ -39,6 +39,9 @@ pub enum LatteError {
 
     #[error(display = "Interrupted")]
     Interrupted(Box<BenchmarkStats>),
+
+    #[error(display = "Failed to launch external editor {}: {}", _0, _1)]
+    ExternalEditorLaunch(String, std::io::Error),
 }
 
 pub type Result<T> = std::result::Result<T, LatteError>;

src/main.rs

@@ -1,3 +1,4 @@
+use std::env;
 use std::fs::File;
 use std::io::{stdout, Write};
 use std::path::{Path, PathBuf};
@@ -15,8 +16,8 @@ use tokio::runtime::{Builder, Runtime};
 use config::RunCommand;
 
 use crate::config::{
-    AppConfig, Command, ConnectionConf, HdrCommand, Interval, LoadCommand, SchemaCommand,
-    ShowCommand,
+    AppConfig, Command, ConnectionConf, EditCommand, HdrCommand, Interval, LoadCommand,
+    SchemaCommand, ShowCommand,
 };
 use crate::context::*;
 use crate::context::{CassError, CassErrorKind, Context, SessionStats};
@@ -371,6 +372,7 @@ async fn export_hdr_log(conf: HdrCommand) -> Result<()> {
 
 async fn async_main(command: Command) -> Result<()> {
     match command {
+        Command::Edit(config) => edit(config)?,
         Command::Schema(config) => schema(config).await?,
         Command::Load(config) => load(config).await?,
         Command::Run(config) => run(config).await?,
@@ -381,6 +383,26 @@ async fn async_main(command: Command) -> Result<()> {
     Ok(())
 }
 
+fn edit(config: EditCommand) -> Result<()> {
+    let workload = find_workload(&config.workload)
+        .canonicalize()
+        .unwrap_or_else(|_| config.workload.to_path_buf());
+    File::open(&workload).map_err(|err| LatteError::ScriptRead(workload.clone(), err))?;
+    edit_workload(workload)
+}
+
+fn edit_workload(workload: PathBuf) -> Result<()> {
+    let editor = env::var("LATTE_EDITOR")
+        .or_else(|_| env::var("EDITOR"))
+        .unwrap_or("vi".to_string());
+    std::process::Command::new(&editor)
+        .current_dir(workload.parent().unwrap_or(Path::new(".")))
+        .arg(workload)
+        .status()
+        .map_err(|e| LatteError::ExternalEditorLaunch(editor, e))?;
+    Ok(())
+}
+
 fn init_runtime(thread_count: usize) -> std::io::Result<Runtime> {
     if thread_count == 1 {
         Builder::new_current_thread().enable_all().build()