zkli 0.1.0

Cli around zookeeper
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

use std::{
    io::{stdout, Write},
    time::Duration,
};

use clap::Parser;
use miette::{IntoDiagnostic, Result};
use zookeeper::ZooKeeper;

#[derive(Debug, Parser)]
#[clap(about = "Cli around zookeeper")]
struct Options {
    #[clap(long, short, default_value_t = String::from("localhost:2181/"))]
    pub addr: String,

    #[clap(subcommand)]
    pub command: Command,
}

#[derive(Debug, Parser)]
enum Command {
    Ls { path: Option<String> },
    Tree { path: Option<String> },
    Cat { file: String },
    Rm { file: String },
}

fn main() -> Result<()> {
    let opt = Options::parse();
    let mut log_builder = env_logger::Builder::new();
    log_builder.parse_filters("info");

    log_builder.init();

    log::info!("Connecting to {}", opt.addr);
    let zk = ZooKeeper::connect(&opt.addr, Duration::from_secs(1), |_| ()).into_diagnostic()?;
    log::info!("Connected");

    match opt.command {
        Command::Ls { path } => {
            let mut path = path.unwrap_or(String::from("/"));
            sanitize_path(&mut path);
            let mut children = zk.get_children(&path, false).into_diagnostic()?;
            children.sort();
            for child in children {
                if path == "/" {
                    path = format!("");
                }
                let stat = zk
                    .exists(&format!("{}/{}", path, child), false)
                    .into_diagnostic()?
                    .unwrap();
                if stat.num_children == 0 {
                    print!("{child} ");
                } else {
                    print!("{child}/ ");
                }
            }
            println!();
        }
        Command::Tree { path } => {
            let mut path = path.unwrap_or(String::from("/"));
            sanitize_path(&mut path);
            tree(&zk, &path, 0)?;
        }
        Command::Cat { mut file } => {
            sanitize_path(&mut file);
            let (data, _) = zk.get_data(&file, false).into_diagnostic()?;
            stdout().write_all(&data).into_diagnostic()?;
        }
        Command::Rm { mut file } => {
            sanitize_path(&mut file);
            zk.delete(&file, None).into_diagnostic()?;
            log::info!("Successfully deleted");
        }
    }
    Ok(())
}

fn tree(zk: &ZooKeeper, mut path: &str, depth: usize) -> Result<()> {
    println!("{}{}", "  ".repeat(depth), path);
    let mut children = zk.get_children(&path, false).into_diagnostic()?;
    if path == "/" {
        path = ""
    }
    children.sort();
    for child in children {
        let stat = zk
            .exists(&format!("{}/{}", path, child), false)
            .into_diagnostic()?
            .unwrap();
        if stat.num_children == 0 {
            println!("{}{child}", "  ".repeat(depth + 1));
        } else {
            tree(zk, &format!("{path}/{child}"), depth + 1)?;
        }
    }
    Ok(())
}

fn sanitize_path(path: &mut String) {
    if !path.starts_with("/") {
        log::warn!(
            "Invalid path, adding a `/` to the beginning of your path: `{path}` => `/{path}`"
        );
        *path = format!("/{path}");
    }
    if path.ends_with("/") && *path != "/" {
        *path = path.trim_end_matches("/").to_string();
        log::warn!("Invalid path, removing the `/` at the end of your path: `{path}/` => `{path}`");
    }
}