Writing Command Line Programs in Rust

こんにちは、次世代システム研究室のN.M.です。

Introduction

Rust is a systems programming language that combines safety, speed, and concurrency. In this blog post, we’ll take a brief tour of writing a command line program in Rust and introduce some useful packages for creating powerful CLI applications.

The purpose of this post is to give the reader a feel for what a CLI program written in Rust can look like. I wrote the program in an attempt to learn Rust as well as to improve some existing session switching and time management shell scripts. Also, these days Rust is being used more and more for CLI programs (eg. https://towardsdatascience.com/awesome-rust-powered-command-line-utilities-b5359c38692), and I wanted to discover why for myself.

 

The App

This is some of the functionality.

What’s happening

• We start in a tmux session called rust-work(master)
  • There are 2 panes, to the left is a pane for editing our rust program and to the right a pane for terminal interaction
  • rust-work(master) is our tmux session name. It is made from the last directory in the project path and the git branch of the project, if present.
  • The most recent two entries in an sqlite table called tasks are output. The tasks table is how we classify our work. Currently the last 2 tasks are shown: 32|rust-work(master) and 31|hello_world(master)
• I invoke the tmux fzf custom command popup to switch sessions (code courtesy of Waylon Walker: https://waylonwalker.com/tmux-fzf-session-jump/)
  • I choose a session to switch to: hello-world(master)
  • I have a tmux hook configured to invoke my rust session switcher program
    • • set-hook -g client-session-changed 'run "tmux display-popup -E \"~/bin/rust-work clock-task '\''#S'\''\""'
      • This calls my rust program that handles clocking in and out of tasks and tells tmux to run it inside a new popup
      • rust-work clock-task clocks out of the previous task (rust-work(master)), and clocks into the new task (hello-world(master))
        • This was the same task that was shown in the beginning
• I invoke the tmux fzf session jump command again
  • I choose a session to switch to: hello_world_server(master)
  • This time an entry in the task table does not exist for this session name, so rust-work clock-task prompts the user for necessary information to make a new entry into the task table. The user selects from existing departments and task types and the new task is inserted into the task table.
  • As, before the previous task is clocked out and the next task is clocked in
• I switch back to to the session we started from at the beginning, rust-work(master) and query the new hello_world_server(master) task and the work_log entry that references that task.

Structure of the CLI rust-work Project

When building a command line program in Rust, organizing your code into a clear directory structure is essential. For example, the sample CLI project has the following directory structure:

.
├── Cargo.lock
├── Cargo.toml
├── diesel.toml
└── src
├── args.rs
├── db
│ ├── models.rs
│ └── schema.rs
├── db.rs
├── main.rs
├── ops
│ └── tasks.rs
└── ops.rs

This structure separates the argument parsing, database handling, and operations into their respective modules, making it easy to understand and maintain.

Leveraging External Libraries

To enhance your CLI application, you can utilize external libraries for various functionalities like user prompts, directory navigation, and more.

Argument Parsing with Clap

• https://github.com/clap-rs/clap

Clap is a popular Rust crate for parsing command line arguments. It provides a simple and efficient way to define and handle arguments in your CLI application. To use Clap, add it as a dependency in your Cargo.toml file and import it in your args.rs module.

Here’s a simple example of how to define command line arguments using Clap:

use clap::{
    Args,
    Parser,
    Subcommand
};

#[derive(Parser,Debug)]
#[clap(author, version, about)]
pub struct RustWorkArgs {
    /// Specify the SubCommand
    #[clap(subcommand)]
    pub task: Task,
}

#[derive(Debug, Subcommand)]
pub enum Task {
    /// Clock in to new task and clock out of last task
    ClockTask(ClockTask),
    /// Create a new session or change to an existing session
    Session(Session)
}

#[derive(Debug, Args)]
pub struct ClockTask {
   pub task: String
}

#[derive(Debug, Args)]
pub struct Session {
    pub task: Option<String>
}

 

This gives me 2 subtasks in my program: session and clock-task.

As shown by automatically generated help:

➜ rust-work
Usage: rust-work <COMMAND>

Commands:
  clock-task  Clock in to new task and clock out of last task
  session     Create a new session or change to an existing session
  help        Print this message or the help of the given subcommand(s)

Options:
  -h, --help     Print help
  -V, --version  Print version

 

Database Access with Diesel

• https://diesel.rs/

Diesel is an ORM (Object-Relational Mapping) and query builder for Rust. It helps you manage your database schema and perform database operations with ease. In our project, we use Diesel for SQLite database management.

To set up Diesel, add it to your Cargo.toml and create a diesel.toml configuration file. Then, create models.rs and schema.rs modules in the db directory to define your database models and schema. schema.rs can be generated automatically using the diesel CLI.

~/.asdf/installs/rust/1.66.1/bin/diesel print-schema > src/db/schema.rs

I use asdf to handle versions of my development language tools, and when I installed the diesel CLI via Rust cargo, diesel was automatically installed under the appropriate rust shim directory.

Here’s an example of how to establish a connection to the SQLite database using Diesel:

pub fn establish_connection() -> SqliteConnection {
    dotenv().ok();
    let database_url = match env::var("DATABASE_URL") {
        Ok(val) => val,
        Err(_) => format!("{}/.data/work.db", env::var("HOME").unwrap()),
    };
    SqliteConnection::establish(&database_url)
        .unwrap_or_else(|_| panic!("Error connecting to {}", database_url))
}

In the case of sqlite3 the DATABASE_URL environment variable is just a path to the sqlite3 file. This is defined in a .env file:

DATABASE_URL=/Users/usr0101345/.data/work.db

I had to write no SQL in this project has it was all handled by Diesels ORM functionality.

pub fn switch_tasks(task_name: &str) {
    use crate::db::schema::current_task::dsl::current_task;
    use crate::db::schema::task::*;
    let connection = &mut crate::db::establish_connection();
    let task_res = task.filter(name.eq(task_name)).first::<Task>(connection);

    match task_res {
        Ok(task_row) => {
            let current_task_row = current_task
                .first::<CurrentTask>(connection)
                .expect("Error loading current_task");

            if current_task_row.task_id == task_row.id {
                return;
            }
            
            make_task_current(current_task_row, &task_row, connection);
        }
        Err(err) => {
            if err.eq(&diesel::result::Error::NotFound) {

                match Confirm::with_theme(&ColorfulTheme::default()).with_prompt(
                    format!(
                        "No task was found for {}.\nDo you want to create a new task?",
                        &task_name
                    )).interact_opt().unwrap() {
                    Some(true) => {
                        let department_item = select_department(connection);

                        let dept_id = match department_item {
                            Some(d) => d.id,
                            None => {
                                return;
                            }
                        };

                        let task_type_item = select_task_type(connection);
                        let task_tp_id = match task_type_item {
                                Some(tt) => Some(tt.id),
                                None => {
                                    return;
                                }
                            };

                        let current_dir = env::current_dir().expect("Failed to get the current directory");

                        // Convert the PathBuf to a &str
                        let current_dir_str = current_dir
                            .to_str()
                            .expect("Failed to convert the current directory to a string");

                        let new_task = NewTask {
                            name: task_name.to_string(),
                            project_root: current_dir_str.to_string(),
                            due_date: None,
                            department_id: dept_id,
                            task_type_id: task_tp_id,
                            jira_issue_id: None,
                        };

                        let inserted: Vec<Task> = insert_into(task)
                            .values(&new_task)
                            .get_results(connection)
                            .unwrap();

                        let current_task_row = current_task
                            .first::<CurrentTask>(connection)
                            .expect("Error loading current_task");


                        make_task_current(current_task_row, &inserted[0], connection);
                    },
                    Some(false) => println!("nevermind then :("),
                    None => println!("I see you found another way out!")
                }
            } else {
                eprintln!("Error occurred selecting task: {}, {}", task_name, err);
            }
        }
    }
}

fn make_task_current(current_task_row: CurrentTask, task_row: &Task, connection: &mut SqliteConnection) {
    use crate::db::schema::current_task::dsl::current_task;

    let now: DateTime<Utc> = Utc::now();
    let local_time = now.with_timezone(&Tokyo);
    let formatted = local_time.format("%Y-%m-%dT%H:%M:%S").to_string();

    let new_work_log = NewWorkLog {
        task_id: current_task_row.task_id,
        clock_in: current_task_row.clock_in.unwrap(),
        clock_out: formatted.clone(),
        description: "None".to_string(),
    };

    let _num_inserted = insert_into(work_log)
        .values(&new_work_log)
        .execute(connection)
        .unwrap();

    let new_current_task = CurrentTask {
        id: 1,
        task_id: task_row.id,
        project_root: task_row.project_root.clone(),
        clock_in: Some(formatted),
    };

    diesel::update(current_task.find(1))
        .set(&new_current_task)
        .execute(connection)
        .expect("Error updating current_task");
}

This code shows all aspects of DB CRUD for diesel and sqlite. Most of the code is fairly self explanatory. The retrieval of inserted on line 61, was necessary to get the auto-incremented Primary Key value of the new task. For this to work you need the diesel feature "returning_clauses_for_sqlite_3_35" inside your Cargo.toml:

diesel = { version = "2.0.3", features = ["sqlite", "returning_clauses_for_sqlite_3_35"] }

dialoguer

• https://github.com/console-rs/dialoguer

A library for creating CLI user prompts and progress bars.

See line 22, from the above code, for an example of using this lib.

skim

• https://github.com/lotabout/skim

The skim selector is a fuzzy finder library that provides a convenient interface for users to quickly search and select items from a list. In this case, it’s used to display the filtered list of project directories to the user, allowing them to easily select the desired directory. Skim’s various options, such as multi-select and reverse order, are also configured in this file.

 

#[derive(Clone)]
struct MyItem {
    id: i32,
    inner: String,
}

impl SkimItem for MyItem {
    fn text(&self) -> Cow {
        Cow::Borrowed(&self.inner)
    }
}

fn select_department(connection: &mut SqliteConnection) -> Option<MyItem> {
    let options = SkimOptionsBuilder::default()
        .multi(false)
        .reverse(true)
        .build()
        .unwrap();

    let (tx_item, rx_item): (SkimItemSender, SkimItemReceiver) = unbounded();
    let departments = department
        .load::<Department>(connection)
        .expect("Failed to retrieve departments");

    for d in departments {
        let _ = tx_item.send(Arc::new(MyItem {
            id: d.id,
            inner: d.name,
        }));
    }
    drop(tx_item); // so that skim could know when to stop waiting for more items.
    if let Some(out) = Skim::run_with(&options, Some(rx_item)) {
        match out.final_key {
            // Create a new item based on the query
            Key::Enter => {
                let immutable: &MyItem = (*out.selected_items[0])
                    .as_any() // cast to Any
                    .downcast_ref::<MyItem>() // downcast to concrete type
                    .expect("something wrong with downcast");
                Some(immutable.clone())
            }
            _ => None,
        }
    } else {
        None
    }
} 

This is the part that does the display of selects with fuzzy finding for the user to select the department. Selection of task type is very similar. One thing to note is the asynchronous population of items in the select, this makes for fast display if you have many items.

chrono

• https://docs.rs/chrono/latest/chrono/

A standard library for handling date and time operations.

Get a local date like:

let now: DateTime<Utc> = Utc::now();
let local_time = now.with_timezone(&Tokyo);
let formatted = local_time.format("%Y-%m-%dT%H:%M:%S").to_string();

 

.z directory and z command

• https://github.com/agkozak/zsh-z

The .z file is a data file created by the z command, which is a command line tool for quickly navigating to frequently used directories. The z command tracks your most commonly visited directories and allows you to jump to them using fuzzy matching. I find this command very useful and use it in place of cd. z uses a file ~/.z to store the recency of visited files. The z command line tool is a zsh plugin, to install, simply list it in your .zshrc plugins section, as listed below along with some other of my favorite plugins.

 

plugins=(
  zsh-vi-mode
  git
  zsh-autosuggestions
  z
)

 

In tasks.rs, the .z file is read to obtain a list of directories, which are then used for further processing.

pub async fn new_session(task_dir: Option<String>) {
  let z_file_path = format!("{}/.z", env::var("HOME").unwrap());
  let z_file = File::open(z_file_path).unwrap();
  const PIPE: char = '|';

  let directories = BufReader::new(z_file)
      .lines()
      .map(|line| line.unwrap())
      .filter_map(|line| match line.contains(PIPE) {
          true => Some(line.split(PIPE).next().unwrap().to_owned()),
          false => None,
      })
      .filter(|dir| {
          let path = PathBuf::from(dir);
          path.is_dir()
              && PROJECT_FILES
                  .iter()
                  .any(|proj_file| path.join(proj_file).exists())
      });

 

Filtering project directories

• https://doc.rust-lang.org/std/iter/struct.Filter.html

In the new_session function, shown above, the list of directories obtained from the .z file is filtered to only include project directories. A project directory is identified by the presence of specific files, such as Cargo.toml, package.json, .proj_id, .git, and README.md. This filtering ensures that the skim selector only displays relevant project directories.

 

Invoking tmux

• https://doc.rust-lang.org/std/process/struct.Command.html

tmux is a terminal multiplexer, allowing users to create, access, and control multiple terminal sessions from a single screen. In tasks.rs, the to_session function uses tmux to create a new session or attach to an existing one, based on the selected project directory. It also handles switching between existing sessions, providing seamless navigation between projects.

async fn to_session(selection: &str) {
    let session_name = session_name(selection);

    let tmux_has_session = Command::new("tmux")
        .arg("has-session")
        .arg("-t")
        .arg(&session_name)
        .output()
        .expect("failed to execute process");

    let _result = match env::var("TMUX") {
        Ok(_in_tmux) => {
            if !tmux_has_session.status.success() {
                // Can't create an attached session from within tmux
                let _tmux_new_session = Command::new("tmux")
                    .arg("new-session")
                    .arg("-d")
                    .arg("-s")
                    .arg(&session_name)
                    .arg("-c")
                    .arg(selection)
                    .output()
                    .expect("failed to execute process");

                // println!("tmux: {:?}", _tmux_new_session);
            }

            Command::new("tmux")
                .arg("switch-client")
                .arg("-t")
                .arg(&session_name)
                .output()
                .expect("failed to execute process");
        }
        Err(_error) => {
            if tmux_has_session.status.success() {
                let _tmux_attach_session = Command::new("tmux")
                    .arg("attach-session")
                    .arg("-t")
                    .arg(&session_name)
                    .stdin(Stdio::inherit())
                    .output()
                    .expect("failed to execute process");

                // println!("tmux: {:?}", _tmux_attach_session);
            } else {
                let _tmux_new_session = Command::new("tmux")
                    .arg("new-session")
                    .arg("-s")
                    .arg(&session_name)
                    .stdin(Stdio::inherit())
                    .arg("-c")
                    .arg(selection)
                    .output()
                    .expect("failed to execute process");
                // println!("tmux: {:?}", _tmux_new_session);
            }
        }
    };
}

The code above handles invoking tmux via the Command library. tmux is invoked from outside tmux, whether a tmux server is currently running or not, or from within tmux. In both cases if the target session already exists it is attached to, otherwise a new session is created.

There are some tmux libraries, which wrap the tmux commands in Rust such as tmux-interface: https://docs.rs/tmux_interface/latest/tmux_interface/

But I since these libraries call tmux as I am doing above using the Command lib, I thought it better to learn the Command api, which can be used for all shell commands.

Summary

tasks.rs implements essential functionalities for the command line program using a combination of helpful libraries and tools. By utilizing the power of the z command, Skim, tmux, and Diesel, this blog gives an idea how to create a feature-rich CLI application in Rust that simplifies project time tracking and session navigation in a fast, reliable, manner.

Links to Past Similar articles

• https://recruit.gmo.jp/engineer/jisedai/blog/the-power-of-tmux-in-scripting/
• https://recruit.gmo.jp/engineer/jisedai/blog/save-development-time-with-tmux-and-tmuxinator/

Additional References

• https://diesel.rs/guides/all-about-inserts.html
• https://rust-cli.github.io/book/index.html
• https://www.youtube.com/results?search_query=rust+command+line+programs

 

次世代システム研究室では、グループ全体のインテグレーションを支援してくれるアーキテクトを募集しています。インフラ設計、構築経験者の方、次世代システム研究室にご興味を持って頂ける方がいらっしゃいましたら、ぜひ募集職種一覧からご応募をお願いします。