garage/src/k2v-client/bin/k2v-cli.rs

601 lines
14 KiB
Rust
Raw Normal View History

use std::collections::BTreeMap;
use std::process::exit;
2022-07-04 12:00:02 +00:00
use std::time::Duration;
use k2v_client::*;
use garage_util::formater::format_table;
use rusoto_core::credential::AwsCredentials;
use rusoto_core::Region;
use clap::{Parser, Subcommand};
/// K2V command line interface
#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
/// Name of the region to use
#[clap(short, long, env = "AWS_REGION", default_value = "garage")]
region: String,
/// Url of the endpoint to connect to
#[clap(short, long, env = "K2V_ENDPOINT")]
endpoint: String,
/// Access key ID
#[clap(short, long, env = "AWS_ACCESS_KEY_ID")]
key_id: String,
/// Access key ID
#[clap(short, long, env = "AWS_SECRET_ACCESS_KEY")]
secret: String,
/// Bucket name
#[clap(short, long, env = "K2V_BUCKET")]
bucket: String,
#[clap(subcommand)]
command: Command,
}
#[derive(Subcommand, Debug)]
enum Command {
/// Insert a single value
Insert {
/// Partition key to insert to
partition_key: String,
/// Sort key to insert to
sort_key: String,
/// Causality of the insertion
#[clap(short, long)]
causality: Option<String>,
/// Value to insert
#[clap(flatten)]
value: Value,
},
/// Read a single value
Read {
/// Partition key to read from
partition_key: String,
/// Sort key to read from
sort_key: String,
/// Output formating
#[clap(flatten)]
output_kind: ReadOutputKind,
},
/// Watch changes on a single value
PollItem {
/// Partition key of item to watch
partition_key: String,
/// Sort key of item to watch
sort_key: String,
/// Causality information
#[clap(short, long)]
causality: String,
2022-07-04 12:00:02 +00:00
/// Timeout, in seconds
#[clap(short = 'T', long)]
2022-07-04 12:00:02 +00:00
timeout: Option<u64>,
/// Output formating
#[clap(flatten)]
output_kind: ReadOutputKind,
},
/// Watch changes on a range of values
PollRange {
/// Partition key to poll from
partition_key: String,
/// Output only sort keys matching this filter
#[clap(flatten)]
filter: Filter,
/// Marker of data that had previously been seen by a PollRange
#[clap(short = 'S', long)]
seen_marker: Option<String>,
/// Timeout, in seconds
#[clap(short = 'T', long)]
timeout: Option<u64>,
/// Output formating
#[clap(flatten)]
output_kind: BatchOutputKind,
},
/// Delete a single value
Delete {
/// Partition key to delete from
partition_key: String,
/// Sort key to delete from
sort_key: String,
/// Causality information
#[clap(short, long)]
causality: String,
},
/// List partition keys
ReadIndex {
/// Output formating
#[clap(flatten)]
output_kind: BatchOutputKind,
/// Output only partition keys matching this filter
#[clap(flatten)]
filter: Filter,
},
/// Read a range of sort keys
ReadRange {
/// Partition key to read from
partition_key: String,
/// Output formating
#[clap(flatten)]
output_kind: BatchOutputKind,
/// Output only sort keys matching this filter
#[clap(flatten)]
filter: Filter,
},
/// Delete a range of sort keys
DeleteRange {
/// Partition key to delete from
partition_key: String,
/// Output formating
#[clap(flatten)]
output_kind: BatchOutputKind,
/// Delete only sort keys matching this filter
#[clap(flatten)]
filter: Filter,
},
}
/// Where to read a value from
#[derive(Parser, Debug)]
#[clap(group = clap::ArgGroup::new("value").multiple(false).required(true))]
struct Value {
/// Read value from a file. use - to read from stdin
#[clap(short, long, group = "value")]
file: Option<String>,
/// Read a base64 value from commandline
#[clap(short, long, group = "value")]
b64: Option<String>,
/// Read a raw (UTF-8) value from the commandline
#[clap(short, long, group = "value")]
text: Option<String>,
}
impl Value {
async fn to_data(&self) -> Result<Vec<u8>, Error> {
if let Some(ref text) = self.text {
Ok(text.as_bytes().to_vec())
} else if let Some(ref b64) = self.b64 {
base64::decode(b64).map_err(|_| Error::Message("invalid base64 input".into()))
} else if let Some(ref path) = self.file {
use tokio::io::AsyncReadExt;
if path == "-" {
let mut file = tokio::io::stdin();
let mut vec = Vec::new();
file.read_to_end(&mut vec).await?;
Ok(vec)
} else {
let mut file = tokio::fs::File::open(path).await?;
let mut vec = Vec::new();
file.read_to_end(&mut vec).await?;
Ok(vec)
}
} else {
unreachable!("Value must have one option set")
}
}
}
#[derive(Parser, Debug)]
#[clap(group = clap::ArgGroup::new("output-kind").multiple(false).required(false))]
struct ReadOutputKind {
/// Base64 output. Conflicts are line separated, first line is causality token
#[clap(short, long, group = "output-kind")]
b64: bool,
/// Raw output. Conflicts generate error, causality token is not returned
#[clap(short, long, group = "output-kind")]
raw: bool,
/// Human formated output
#[clap(short = 'H', long, group = "output-kind")]
human: bool,
/// JSON formated output
#[clap(short, long, group = "output-kind")]
json: bool,
}
impl ReadOutputKind {
fn display_output(&self, val: CausalValue) -> ! {
use std::io::Write;
if self.json {
let stdout = std::io::stdout();
serde_json::to_writer_pretty(stdout, &val).unwrap();
exit(0);
}
if self.raw {
let mut val = val.value;
if val.len() != 1 {
eprintln!(
"Raw mode can only read non-concurent values, found {} values, expected 1",
val.len()
);
exit(1);
}
let val = val.pop().unwrap();
match val {
K2vValue::Value(v) => {
std::io::stdout().write_all(&v).unwrap();
exit(0);
}
K2vValue::Tombstone => {
eprintln!("Expected value, found tombstone");
exit(2);
}
}
}
let causality: String = val.causality.into();
if self.b64 {
println!("{}", causality);
for val in val.value {
match val {
K2vValue::Value(v) => {
println!("{}", base64::encode(&v))
}
K2vValue::Tombstone => {
println!();
}
}
}
exit(0);
}
// human
println!("causality: {}", causality);
println!("values:");
for val in val.value {
match val {
K2vValue::Value(v) => {
if let Ok(string) = std::str::from_utf8(&v) {
println!(" utf-8: {}", string);
} else {
println!(" base64: {}", base64::encode(&v));
}
}
K2vValue::Tombstone => {
println!(" tombstone");
}
}
}
exit(0);
}
}
#[derive(Parser, Debug)]
#[clap(group = clap::ArgGroup::new("output-kind").multiple(false).required(false))]
struct BatchOutputKind {
/// Human formated output
#[clap(short = 'H', long, group = "output-kind")]
human: bool,
/// JSON formated output
#[clap(short, long, group = "output-kind")]
json: bool,
}
impl BatchOutputKind {
fn display_human_output(&self, values: BTreeMap<String, CausalValue>) -> ! {
for (key, values) in values {
println!("key: {}", key);
let causality: String = values.causality.into();
println!("causality: {}", causality);
for value in values.value {
match value {
K2vValue::Value(v) => {
if let Ok(string) = std::str::from_utf8(&v) {
println!(" value(utf-8): {}", string);
} else {
println!(" value(base64): {}", base64::encode(&v));
}
}
K2vValue::Tombstone => {
println!(" tombstone");
}
}
}
}
exit(0);
}
fn values_json(&self, values: BTreeMap<String, CausalValue>) -> Vec<serde_json::Value> {
values
.into_iter()
.map(|(k, v)| {
let mut value = serde_json::to_value(v).unwrap();
value
.as_object_mut()
.unwrap()
.insert("sort_key".to_owned(), k.into());
value
})
.collect::<Vec<_>>()
}
fn display_poll_range_output(
&self,
seen_marker: String,
values: BTreeMap<String, CausalValue>,
) -> ! {
if self.json {
let json = serde_json::json!({
"values": self.values_json(values),
"seen_marker": seen_marker,
});
let stdout = std::io::stdout();
serde_json::to_writer_pretty(stdout, &json).unwrap();
exit(0)
} else {
println!("seen marker: {}", seen_marker);
self.display_human_output(values)
}
}
fn display_read_range_output(&self, res: PaginatedRange<CausalValue>) -> ! {
if self.json {
let json = serde_json::json!({
"next_key": res.next_start,
"values": self.values_json(res.items),
});
let stdout = std::io::stdout();
serde_json::to_writer_pretty(stdout, &json).unwrap();
exit(0)
} else {
if let Some(next) = res.next_start {
println!("next key: {}", next);
}
self.display_human_output(res.items)
}
}
}
/// Filter for batch operations
#[derive(Parser, Debug)]
#[clap(group = clap::ArgGroup::new("filter").multiple(true).required(true))]
struct Filter {
/// Match only keys starting with this prefix
#[clap(short, long, group = "filter")]
prefix: Option<String>,
/// Match only keys lexicographically after this key (including this key itself)
#[clap(short, long, group = "filter")]
start: Option<String>,
/// Match only keys lexicographically before this key (excluding this key)
#[clap(short, long, group = "filter")]
end: Option<String>,
/// Only match the first X keys
#[clap(short, long)]
limit: Option<u64>,
/// Return keys in reverse order
#[clap(short, long)]
reverse: bool,
/// Return only keys where conflict happened
#[clap(short, long)]
conflicts_only: bool,
/// Also include keys storing only tombstones
#[clap(short, long)]
tombstones: bool,
/// Return any key
#[clap(short, long, group = "filter")]
all: bool,
}
impl Filter {
fn k2v_filter(&self) -> k2v_client::Filter<'_> {
k2v_client::Filter {
start: self.start.as_deref(),
end: self.end.as_deref(),
prefix: self.prefix.as_deref(),
limit: self.limit,
reverse: self.reverse,
}
}
}
#[tokio::main]
async fn main() -> Result<(), Error> {
let args = Args::parse();
let region = Region::Custom {
name: args.region,
endpoint: args.endpoint,
};
let creds = AwsCredentials::new(args.key_id, args.secret, None, None);
let client = K2vClient::new(region, args.bucket, creds, None)?;
match args.command {
Command::Insert {
partition_key,
sort_key,
causality,
value,
} => {
client
.insert_item(
&partition_key,
&sort_key,
value.to_data().await?,
causality.map(Into::into),
)
.await?;
}
Command::Delete {
partition_key,
sort_key,
causality,
} => {
client
.delete_item(&partition_key, &sort_key, causality.into())
.await?;
}
Command::Read {
partition_key,
sort_key,
output_kind,
} => {
let res = client.read_item(&partition_key, &sort_key).await?;
output_kind.display_output(res);
}
Command::PollItem {
partition_key,
sort_key,
causality,
2022-07-04 12:00:02 +00:00
timeout,
output_kind,
} => {
2022-07-04 12:00:02 +00:00
let timeout = timeout.map(Duration::from_secs);
let res_opt = client
2022-07-04 12:00:02 +00:00
.poll_item(&partition_key, &sort_key, causality.into(), timeout)
.await?;
if let Some(res) = res_opt {
output_kind.display_output(res);
} else {
if output_kind.json {
println!("null");
} else {
println!("Delay expired and value didn't change.");
}
}
}
Command::PollRange {
partition_key,
filter,
seen_marker,
timeout,
output_kind,
} => {
if filter.conflicts_only
|| filter.tombstones
|| filter.reverse
|| filter.limit.is_some()
{
return Err(Error::Message(
"limit, reverse, conlicts-only, tombstones are invalid for poll-range".into(),
));
}
let timeout = timeout.map(Duration::from_secs);
let res = client
.poll_range(
&partition_key,
Some(PollRangeFilter {
start: filter.start.as_deref(),
end: filter.end.as_deref(),
prefix: filter.prefix.as_deref(),
}),
seen_marker.as_deref(),
timeout,
)
.await?;
match res {
Some((items, seen_marker)) => {
output_kind.display_poll_range_output(seen_marker, items);
}
None => {
if output_kind.json {
println!("null");
} else {
println!("Delay expired and value didn't change.");
}
}
}
}
Command::ReadIndex {
output_kind,
filter,
} => {
if filter.conflicts_only || filter.tombstones {
return Err(Error::Message(
"conlicts-only and tombstones are invalid for read-index".into(),
));
}
let res = client.read_index(filter.k2v_filter()).await?;
if output_kind.json {
let values = res
.items
.into_iter()
.map(|(k, v)| {
let mut value = serde_json::to_value(v).unwrap();
value
.as_object_mut()
.unwrap()
.insert("sort_key".to_owned(), k.into());
value
})
.collect::<Vec<_>>();
let json = serde_json::json!({
"next_key": res.next_start,
"values": values,
});
let stdout = std::io::stdout();
serde_json::to_writer_pretty(stdout, &json).unwrap();
} else {
if let Some(next) = res.next_start {
println!("next key: {}", next);
}
let mut to_print = Vec::new();
to_print.push(format!("key:\tentries\tconflicts\tvalues\tbytes"));
for (k, v) in res.items {
to_print.push(format!(
"{}\t{}\t{}\t{}\t{}",
k, v.entries, v.conflicts, v.values, v.bytes
));
}
format_table(to_print);
}
}
Command::ReadRange {
partition_key,
output_kind,
filter,
} => {
let op = BatchReadOp {
partition_key: &partition_key,
filter: filter.k2v_filter(),
conflicts_only: filter.conflicts_only,
tombstones: filter.tombstones,
single_item: false,
};
let mut res = client.read_batch(&[op]).await?;
let res = res.pop().unwrap();
output_kind.display_read_range_output(res);
}
Command::DeleteRange {
partition_key,
output_kind,
filter,
} => {
let op = BatchDeleteOp {
partition_key: &partition_key,
prefix: filter.prefix.as_deref(),
start: filter.start.as_deref(),
end: filter.end.as_deref(),
single_item: false,
};
if filter.reverse
|| filter.conflicts_only
|| filter.tombstones
|| filter.limit.is_some()
{
return Err(Error::Message(
"limit, conlicts-only, reverse and tombstones are invalid for delete-range"
.into(),
));
}
let res = client.delete_batch(&[op]).await?;
if output_kind.json {
println!("{}", res[0]);
} else {
println!("deleted {} keys", res[0]);
}
}
}
Ok(())
}