summaryrefslogtreecommitdiffstats
path: root/src/pg_top.rs
blob: e45f492ba5dab8f80ca6c436772bcad19754b9db (plain)
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
113
extern crate clap;
extern crate termion;

mod display;
mod os;
mod postgresql;

use clap::{App, Arg};
use std::env;
use std::io::{stdout, Write};
use std::thread;
use std::time::Duration;
use termion::raw::IntoRawMode;

pub struct PGTOP {
    pub new: usize,
    pub old: usize,
    pub cpu: os::CPUSTAT,
}

fn main() {
    let matches = App::new("pg_top")
        .version("5.0.0")
        .about("'top' for PostgreSQL")
        .arg(
            Arg::with_name("DELAY")
                .help("set delay between screen updates")
                .default_value("1")
                .takes_value(true),
        )
        .arg(
            Arg::with_name("COUNT")
                .help("change number of displays to show")
                .takes_value(true),
        )
        .get_matches();

    let delay = matches.value_of("DELAY").unwrap().parse::<u64>().unwrap();
    let mut count = if matches.value_of("COUNT").is_none() {
        1
    } else {
        matches.value_of("COUNT").unwrap().parse::<u64>().unwrap()
    };

    let mut conninfo = String::new();

    // Host is required.
    let pghost = if env::var("PGHOST").is_ok() {
        env::var("PGHOST").unwrap()
    } else {
        String::from("/tmp,/var/run/postgresql")
    };
    conninfo.push_str("host=");
    conninfo.push_str(&pghost);

    // User is required.
    let pguser = if env::var("PGUSER").is_ok() {
        env::var("PGUSER").unwrap()
    } else {
        env::var("USER").unwrap()
    };
    conninfo.push_str(" user=");
    conninfo.push_str(&pguser);

    // Database is not required and will default to PGUSER.
    if env::var("PGDATABASE").is_ok() {
        conninfo.push_str(" database=");
        env::var("PGDATABASE").unwrap();
    }

    // Port is not required.
    if env::var("PGPORT").is_ok() {
        conninfo.push_str(" port=");
        env::var("PGPORT").unwrap();
    }

    let mut pg_top = PGTOP {
        new: 1,
        old: 0,
        cpu: os::CPUSTAT {
            stat: [[0; 5], [0; 5]],
        },
    };

    let mut stdout = stdout().into_raw_mode().unwrap();
    write!(stdout, "{}{}", termion::clear::All, termion::cursor::Hide).unwrap();

    loop {
        let mut client = postgresql::connect(&conninfo);
        let mut processes = postgresql::get_processes(&mut client);
        let terminal_size = termion::terminal_size().unwrap();

        display::loadavg(terminal_size.0 as usize);
        processes = display::process_summary(processes, terminal_size.0 as usize);
        pg_top = display::stat_cpu(pg_top, terminal_size.0 as usize);
        display::processes(processes, terminal_size.0 as usize);
        write!(stdout, "{}", termion::clear::AfterCursor,).unwrap();
        stdout.flush().unwrap();

        count = count - 1;
        if count == 0 {
            write!(
                stdout,
                "{}{}",
                termion::cursor::Goto(1, terminal_size.1),
                termion::cursor::Show,
            )
            .unwrap();
            break;
        }
        thread::sleep(Duration::from_secs(delay));
    }
}