diff --git a/src/config.rs b/src/config.rs index e8d42df..8d9230c 100644 --- a/src/config.rs +++ b/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. diff --git a/src/error.rs b/src/error.rs index 5e7bbb4..1047a42 100644 --- a/src/error.rs +++ b/src/error.rs @@ -39,6 +39,9 @@ pub enum LatteError { #[error(display = "Interrupted")] Interrupted(Box), + + #[error(display = "Failed to launch external editor {}: {}", _0, _1)] + ExternalEditorLaunch(String, std::io::Error), } pub type Result = std::result::Result; diff --git a/src/main.rs b/src/main.rs index 1f0691f..b1ec1d2 100644 --- a/src/main.rs +++ b/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 { if thread_count == 1 { Builder::new_current_thread().enable_all().build()