pgcat/src/admin.rs

use bytes::{Buf, BufMut, BytesMut};
use log::{info, trace};
use tokio::net::tcp::OwnedWriteHalf;

use std::collections::HashMap;

use crate::config::{get_config, parse};
use crate::constants::{OID_NUMERIC, OID_TEXT};
use crate::errors::Error;
use crate::messages::write_all_half;
use crate::stats::get_stats;

/// Handle admin client
pub async fn handle_admin(stream: &mut OwnedWriteHalf, mut query: BytesMut) -> Result<(), Error> {
    let code = query.get_u8() as char;

    if code != 'Q' {
        return Err(Error::ProtocolSyncError);
    }

    let len = query.get_i32() as usize;
    let query = String::from_utf8_lossy(&query[..len - 5])
        .to_string()
        .to_ascii_uppercase();

    if query.starts_with("SHOW STATS") {
        trace!("SHOW STATS");
        show_stats(stream).await
    } else if query.starts_with("RELOAD") {
        trace!("RELOAD");
        reload(stream).await
    } else if query.starts_with("SHOW CONFIG") {
        trace!("SHOW CONFIG");
        show_config(stream).await
    } else {
        Err(Error::ProtocolSyncError)
    }
}

/// RELOAD
pub async fn reload(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
    info!("Reloading config");

    let config = get_config();
    let path = config.path.clone().unwrap();

    parse(&path).await?;

    let config = get_config();

    config.show();

    let mut res = BytesMut::new();

    // CommandComplete
    let command_complete = BytesMut::from(&"RELOAD\0"[..]);
    res.put_u8(b'C');
    res.put_i32(command_complete.len() as i32 + 4);
    res.put(command_complete);

    // ReadyForQuery
    res.put_u8(b'Z');
    res.put_i32(5);
    res.put_u8(b'I');

    write_all_half(stream, res).await
}

pub async fn show_config(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
    let guard = get_config();
    let config = &*guard.clone();
    let config: HashMap<String, String> = config.into();
    drop(guard);

    // Configs that cannot be changed dynamically.
    let immutables = ["host", "port", "connect_timeout"];

    // Columns
    let columns = ["key", "value", "default", "changeable"];

    // RowDescription
    let mut row_desc = BytesMut::new();
    row_desc.put_i16(4 as i16); // key, value, default, changeable

    for column in columns {
        row_desc.put_slice(&format!("{}\0", column).as_bytes());

        // Doesn't belong to any table
        row_desc.put_i32(0);

        // Doesn't belong to any table
        row_desc.put_i16(0);

        // Data type
        row_desc.put_i32(OID_TEXT);

        // text size = variable (-1)
        row_desc.put_i16(-1);

        // Type modifier: none that I know
        row_desc.put_i32(-1);

        // Format being used: text (0), binary (1)
        row_desc.put_i16(0);
    }

    // Response data
    let mut res = BytesMut::new();
    res.put_u8(b'T');
    res.put_i32(row_desc.len() as i32 + 4);
    res.put(row_desc);

    // DataRow rows
    for (key, value) in config {
        let mut data_row = BytesMut::new();

        data_row.put_i16(4 as i16); // key, value, default, changeable

        let key_bytes = key.as_bytes();
        let value = value.as_bytes();

        data_row.put_i32(key_bytes.len() as i32);
        data_row.put_slice(&key_bytes);

        data_row.put_i32(value.len() as i32);
        data_row.put_slice(&value);

        data_row.put_i32(1 as i32);
        data_row.put_slice(&"-".as_bytes());

        let changeable = if immutables.iter().filter(|col| *col == &key).count() == 1 {
            "no".as_bytes()
        } else {
            "yes".as_bytes()
        };
        data_row.put_i32(changeable.len() as i32);
        data_row.put_slice(&changeable);

        res.put_u8(b'D');
        res.put_i32(data_row.len() as i32 + 4);
        res.put(data_row);
    }

    res.put_u8(b'C');
    res.put_i32("SHOW CONFIG\0".as_bytes().len() as i32 + 4);
    res.put_slice(&"SHOW CONFIG\0".as_bytes());

    res.put_u8(b'Z');
    res.put_i32(5);
    res.put_u8(b'I');

    write_all_half(stream, res).await
}

/// SHOW STATS
pub async fn show_stats(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
    let columns = [
        "database",
        "total_xact_count",
        "total_query_count",
        "total_received",
        "total_sent",
        "total_xact_time",
        "total_query_time",
        "total_wait_time",
        "avg_xact_count",
        "avg_query_count",
        "avg_recv",
        "avg_sent",
        "avg_xact_time",
        "avg_query_time",
        "avg_wait_time",
    ];

    let stats = get_stats();
    let mut res = BytesMut::new();
    let mut row_desc = BytesMut::new();
    let mut data_row = BytesMut::new();

    // Number of columns: 1
    row_desc.put_i16(columns.len() as i16);
    data_row.put_i16(columns.len() as i16);

    for (i, column) in columns.iter().enumerate() {
        // RowDescription

        // Column name
        row_desc.put_slice(&format!("{}\0", column).as_bytes());

        // Doesn't belong to any table
        row_desc.put_i32(0);

        // Doesn't belong to any table
        row_desc.put_i16(0);

        // Data type
        row_desc.put_i32(if i == 0 { OID_TEXT } else { OID_NUMERIC });

        // Numeric/text size = variable (-1)
        row_desc.put_i16(-1);

        // Type modifier: none that I know
        row_desc.put_i32(-1);

        // Format being used: text (0), binary (1)
        row_desc.put_i16(0);

        // DataRow
        let value = if i == 0 {
            String::from("all shards")
        } else {
            stats.get(&column.to_string()).unwrap_or(&0).to_string()
        };

        data_row.put_i32(value.len() as i32);
        data_row.put_slice(value.as_bytes());
    }

    let command_complete = BytesMut::from(&"SHOW\0"[..]);

    res.put_u8(b'T');
    res.put_i32(row_desc.len() as i32 + 4);
    res.put(row_desc);

    res.put_u8(b'D');
    res.put_i32(data_row.len() as i32 + 4);
    res.put(data_row);

    res.put_u8(b'C');
    res.put_i32(command_complete.len() as i32 + 4);
    res.put(command_complete);

    res.put_u8(b'Z');
    res.put_i32(5);
    res.put_u8(b'I');

    write_all_half(stream, res).await
}