sledcli/src/main.rs

use std::ops::Deref;
use std::path::PathBuf;

use anyhow::Result;

use humansize::{file_size_opts, FileSize};

use rustyline::error::ReadlineError;
use rustyline::Editor;

use structopt::StructOpt;

use sled::IVec;

#[derive(StructOpt, Debug)]
#[structopt(name = "sledcli")]
struct Opt {
	/// Path to Sled database
	#[structopt(name = "path")]
	path: PathBuf,
}

enum DisplayMode {
	TryString,
	HexDump,
	Mixed,
}

struct State {
	db: sled::Db,
	tree: sled::Tree,
	displaymode: DisplayMode,
}

fn main() {
	let opt = Opt::from_args();
	let db = sled::Config::default()
		.path(&opt.path)
		.open()
		.expect("Unable to open database");

	let tree: sled::Tree = db.deref().clone();
	let mut state = State {
		db,
		tree,
		displaymode: DisplayMode::Mixed,
	};

	let mut readline = Editor::<()>::new();
	loop {
		let prefix = match state.displaymode {
			DisplayMode::HexDump => "hex",
			DisplayMode::TryString => "str",
			DisplayMode::Mixed => "mix",
		};
		let prompt = format!("[{}] {}> ", prefix, try_string(&state.tree.name()));
		let lineread = readline.readline(&prompt);
		match lineread {
			Ok(line) => {
				readline.add_history_entry(line.as_str());
				if let Err(e) = do_command(&line, &mut state) {
					println!("Error: {}", e);
				}
			}
			Err(ReadlineError::Interrupted) => {
				println!("^C");
				continue;
			}
			Err(ReadlineError::Eof) => break,
			Err(err) => {
				println!("Readline error: {:?}", err);
				break;
			}
		}
	}
}

fn try_string(input: &sled::IVec) -> String {
	let mut string = String::new();
	utf8::LossyDecoder::new(|s| string.push_str(s)).feed(input);
	string
}

fn do_command(line: &str, state: &mut State) -> Result<()> {
	let parts = line
		.split(' ')
		.filter(|part| part.len() > 0)
		.collect::<Vec<_>>();
	if parts.is_empty() {
		return Ok(());
	}

	match &parts[..] {
		["ls"] => {
			let mut names = state.db.tree_names();
			names.sort();
			for name in names {
				println!("{}", try_string(&name));
			}
		}
		["ll"] => {
			let mut names = state.db.tree_names();
			names.sort();
			let mut total = 0;
			for name in names {
				let nent = state.db.open_tree(&name)?.len();
				total += nent;
				println!("{:8} {}", nent, try_string(&name));
			}
			println!("{:8} TOTAL", total);
		}
		["lu"] => {
			let mut names = state.db.tree_names();
			names.sort();
			let mut total_nent = 0;
			let mut total_size = 0;
			for name in names {
				let tree = state.db.open_tree(&name)?;
				let nent = tree.len();
				let mut size = 0;
				for ent in tree.iter() {
					let (k, v) = ent?;
					size += k.len() + v.len();
				}
				total_nent += nent;
				total_size += size;
				println!(
					"{:8} {:>12} {}",
					nent,
					size.file_size(file_size_opts::CONVENTIONAL).unwrap(),
					try_string(&name)
				);
			}
			println!("{:8} {:>12} TOTAL",
					 total_nent,
					 total_size.file_size(file_size_opts::CONVENTIONAL).unwrap());
		}
		["cd", treename] => {
			if state
				.db
				.tree_names()
				.iter()
				.any(|t| t == treename.as_bytes())
			{
				state.tree = state.db.open_tree(treename.as_bytes())?;
			} else {
				println!("Tree {} does not exist", treename);
			}
		}
		["hex"] => {
			state.displaymode = DisplayMode::HexDump;
		}
		["str"] => {
			state.displaymode = DisplayMode::TryString;
		}
		["mix"] => {
			state.displaymode = DisplayMode::Mixed;
		}
		["keys"] => {
			for (i, pair) in state.tree.iter().enumerate() {
				if i >= 20 {
					println!("...");
					break;
				}
				let (k, _v) = pair?;
				state.displaymode.print_key(&k);
			}
		}
		["pairs"] => {
			for (i, pair) in state.tree.iter().enumerate() {
				if i >= 20 {
					println!("...");
					break;
				}
				let (k, v) = pair?;
				state.displaymode.print_pair(&k, &v);
			}
		}
		bad_cmd => println!("Unrecognized command: {:?}", bad_cmd),
	}
	Ok(())
}

impl DisplayMode {
	fn print_key(&self, k: &IVec) {
		match *self {
			DisplayMode::HexDump => {
				hexdump::hexdump(k);
			}
			_ => {
				println!("{}", try_string(k));
			}
		}
	}

	fn print_pair(&self, k: &IVec, v: &IVec) {
		match *self {
			DisplayMode::HexDump => {
				hexdump::hexdump(k);
				hexdump::hexdump(v);
				println!();
			}
			DisplayMode::Mixed => {
				println!("{}", try_string(k));
				hexdump::hexdump(v);
				println!();
			}
			DisplayMode::TryString => {
				println!("{}\t{}", try_string(k), try_string(v));
			}
		}
	}
}