Refactor admin (#52)

src/admin.rs

@@ -5,9 +5,8 @@ use tokio::net::tcp::OwnedWriteHalf;
 use std::collections::HashMap;
 
 use crate::config::{get_config, parse, Role};
-use crate::constants::{OID_INT4, OID_NUMERIC, OID_TEXT};
 use crate::errors::Error;
-use crate::messages::{custom_protocol_response_ok, error_response, write_all_half};
+use crate::messages::*;
 use crate::pool::ConnectionPool;
 use crate::stats::get_stats;
 
@@ -56,115 +55,66 @@ async fn show_databases(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> R
     let config = &*guard.clone();
     drop(guard);
 
-    let columns = [
+    // Columns
-        "name",
+    let columns = vec![
-        "host",
+        ("name", DataType::Text),
-        "port",
+        ("host", DataType::Text),
-        "database",
+        ("port", DataType::Text),
-        "force_user",
+        ("database", DataType::Text),
-        "pool_size",
+        ("force_user", DataType::Text),
-        "min_pool_size",
+        ("pool_size", DataType::Int4),
-        "reserve_pool",
+        ("min_pool_size", DataType::Int4),
-        "pool_mode",
+        ("reserve_pool", DataType::Int4),
-        "max_connections",
+        ("pool_mode", DataType::Text),
-        "current_connections",
+        ("max_connections", DataType::Int4),
-        "paused",
+        ("current_connections", DataType::Int4),
-        "disabled",
+        ("paused", DataType::Int4),
-    ];
+        ("disabled", DataType::Int4),
-
-    let types = [
-        OID_TEXT, OID_TEXT, OID_TEXT, OID_TEXT, OID_TEXT, OID_INT4, OID_INT4, OID_INT4, OID_TEXT,
-        OID_INT4, OID_INT4, OID_INT4, OID_INT4,
     ];
 
     let mut res = BytesMut::new();
-    let mut row_desc = BytesMut::new();
-    row_desc.put_i16(columns.len() as i16);
 
-    for (i, column) in columns.iter().enumerate() {
+    // RowDescription
-        row_desc.put_slice(&format!("{}\0", column).as_bytes());
+    res.put(row_description(&columns));
-
-        // 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(types[i]);
-
-        // text size = variable (-1)
-        row_desc.put_i16(if types[i] == OID_TEXT { -1 } else { 4 });
-
-        // Type modifier: none that I know
-        row_desc.put_i32(-1);
-
-        // Format being used: text (0), binary (1)
-        row_desc.put_i16(0);
-    }
-
-    res.put_u8(b'T');
-    res.put_i32(row_desc.len() as i32 + 4);
-    res.put(row_desc);
 
     for shard in 0..pool.shards() {
         let database_name = &config.shards[&shard.to_string()].database;
         let mut replica_count = 0;
+
         for server in 0..pool.servers(shard) {
-            // DataRow
-            let mut data_row = BytesMut::new();
-            data_row.put_i16(columns.len() as i16);
-
             let address = pool.address(shard, server);
-            let role = address.role.to_string();
+            let name = match address.role {
-            let name = match role.as_ref() {
+                Role::Primary => format!("shard_{}_primary", shard),
-                "primary" => format!("shard_{}_primary", shard),
+
-                "replica" => format!("shard_{}_replica_{}", shard, replica_count),
+                Role::Replica => {
-                _ => unreachable!(),
+                    let name = format!("shard_{}_replica_{}", shard, replica_count);
+                    replica_count += 1;
+                    name
+                }
             };
-            let connections = pool.connections(shard, server);
+            let pool_state = pool.pool_state(shard, server);
 
-            let data = HashMap::from([
+            res.put(data_row(&vec![
-                ("host", address.host.to_string()),
+                name,                                 // name
-                ("port", address.port.to_string()),
+                address.host.to_string(),             // host
-                ("role", role),
+                address.port.to_string(),             // port
-                ("name", name),
+                database_name.to_string(),            // database
-                ("database", database_name.to_string()),
+                config.user.name.to_string(),         // force_user
-                ("force_user", config.user.name.to_string()),
+                config.general.pool_size.to_string(), // pool_size
-                ("pool_size", config.general.pool_size.to_string()),
+                "0".to_string(),                      // min_pool_size
-                ("min_pool_size", "0".to_string()),
+                "0".to_string(),                      // reserve_pool
-                ("reserve_pool", "0".to_string()),
+                config.general.pool_mode.to_string(), // pool_mode
-                ("pool_mode", config.general.pool_mode.to_string()),
+                config.general.pool_size.to_string(), // max_connections
-                // There is only one user support at the moment,
+                pool_state.connections.to_string(),   // current_connections
-                // so max_connections = num of users * pool_size = 1 * pool_size.
+                "0".to_string(),                      // paused
-                ("max_connections", config.general.pool_size.to_string()),
+                "0".to_string(),                      // disabled
-                ("current_connections", connections.connections.to_string()),
+            ]));
-                ("paused", "0".to_string()),
-                ("disabled", "0".to_string()),
-            ]);
-
-            for column in &columns {
-                let value = data[column].as_bytes();
-
-                data_row.put_i32(value.len() as i32);
-                data_row.put_slice(&value);
-            }
-
-            res.put_u8(b'D');
-            res.put_i32(data_row.len() as i32 + 4);
-            res.put(data_row);
-
-            if address.role == Role::Replica {
-                replica_count += 1;
-            }
         }
     }
 
-    let command_complete = BytesMut::from(&"SHOW\0"[..]);
+    res.put(command_complete("SHOW"));
-    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');
@@ -194,10 +144,7 @@ async fn reload(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
     let mut res = BytesMut::new();
 
     // CommandComplete
-    let command_complete = BytesMut::from(&"RELOAD\0"[..]);
+    res.put(command_complete("RELOAD"));
-    res.put_u8(b'C');
-    res.put_i32(command_complete.len() as i32 + 4);
-    res.put(command_complete);
 
     // ReadyForQuery
     res.put_u8(b'Z');
@@ -217,74 +164,31 @@ async fn show_config(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
     let immutables = ["host", "port", "connect_timeout"];
 
     // Columns
-    let columns = ["key", "value", "default", "changeable"];
+    let columns = vec![
-
+        ("key", DataType::Text),
-    // RowDescription
+        ("value", DataType::Text),
-    let mut row_desc = BytesMut::new();
+        ("default", DataType::Text),
-    row_desc.put_i16(4 as i16); // key, value, default, changeable
+        ("changeable", DataType::Text),
-
+    ];
-    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(row_description(&columns));
-    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()
+            "no".to_string()
         } else {
-            "yes".as_bytes()
+            "yes".to_string()
         };
-        data_row.put_i32(changeable.len() as i32);
-        data_row.put_slice(&changeable);
 
-        res.put_u8(b'D');
+        let row = vec![key, value, "-".to_string(), changeable];
-        res.put_i32(data_row.len() as i32 + 4);
+
-        res.put(data_row);
+        res.put(data_row(&row));
     }
 
-    res.put_u8(b'C');
+    res.put(command_complete("SHOW"));
-    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);
@@ -295,81 +199,38 @@ async fn show_config(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
 
 /// SHOW STATS
 async fn show_stats(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
-    let columns = [
+    let columns = vec![
-        "database",
+        ("database", DataType::Text),
-        "total_xact_count",
+        ("total_xact_count", DataType::Numeric),
-        "total_query_count",
+        ("total_query_count", DataType::Numeric),
-        "total_received",
+        ("total_received", DataType::Numeric),
-        "total_sent",
+        ("total_sent", DataType::Numeric),
-        "total_xact_time",
+        ("total_xact_time", DataType::Numeric),
-        "total_query_time",
+        ("total_query_time", DataType::Numeric),
-        "total_wait_time",
+        ("total_wait_time", DataType::Numeric),
-        "avg_xact_count",
+        ("avg_xact_count", DataType::Numeric),
-        "avg_query_count",
+        ("avg_query_count", DataType::Numeric),
-        "avg_recv",
+        ("avg_recv", DataType::Numeric),
-        "avg_sent",
+        ("avg_sent", DataType::Numeric),
-        "avg_xact_time",
+        ("avg_xact_time", DataType::Numeric),
-        "avg_query_time",
+        ("avg_query_time", DataType::Numeric),
-        "avg_wait_time",
+        ("avg_wait_time", DataType::Numeric),
     ];
 
     let stats = get_stats();
     let mut res = BytesMut::new();
-    let mut row_desc = BytesMut::new();
+    res.put(row_description(&columns));
-    let mut data_row = BytesMut::new();
 
-    // Number of columns: 1
+    let mut row = vec![
-    row_desc.put_i16(columns.len() as i16);
+        String::from("all shards"), // TODO: per-database stats,
-    data_row.put_i16(columns.len() as i16);
+    ];
 
-    for (i, column) in columns.iter().enumerate() {
+    for column in &columns[1..] {
-        // RowDescription
+        row.push(stats.get(column.0).unwrap_or(&0).to_string());
-
-        // 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(data_row(&row));
-
+    res.put(command_complete("SHOW"));
-    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);

src/constants.rs

@@ -24,7 +24,4 @@ pub const MESSAGE_TERMINATOR: u8 = 0;
 //
 // Data types
 //
-pub const OID_NUMERIC: i32 = 1700;
-pub const OID_TEXT: i32 = 25;
-pub const OID_INT4: i32 = 23; // int
 pub const _OID_INT8: i32 = 20; // bigint

src/messages.rs

@@ -8,9 +8,26 @@ use tokio::net::{
     TcpStream,
 };
 
+use crate::errors::Error;
 use std::collections::HashMap;
 
-use crate::errors::Error;
+/// Postgres data type mappings
+/// used in RowDescription ('T') message.
+pub enum DataType {
+    Text,
+    Int4,
+    Numeric,
+}
+
+impl From<&DataType> for i32 {
+    fn from(data_type: &DataType) -> i32 {
+        match data_type {
+            DataType::Text => 25,
+            DataType::Int4 => 23,
+            DataType::Numeric => 1700,
+        }
+    }
+}
 
 /// Tell the client that authentication handshake completed successfully.
 pub async fn auth_ok(stream: &mut TcpStream) -> Result<(), Error> {
@@ -259,68 +276,17 @@ pub async fn show_response(
     // 3. CommandComplete
     // 4. ReadyForQuery
 
-    // RowDescription
-    let mut row_desc = BytesMut::new();
-
-    // Number of columns: 1
-    row_desc.put_i16(1);
-
-    // Column name
-    row_desc.put_slice(&format!("{}\0", name).as_bytes());
-
-    // Doesn't belong to any table
-    row_desc.put_i32(0);
-
-    // Doesn't belong to any table
-    row_desc.put_i16(0);
-
-    // Text
-    row_desc.put_i32(25);
-
-    // 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 mut data_row = BytesMut::new();
-
-    // Number of columns
-    data_row.put_i16(1);
-
-    // Size of the column content (length of the string really)
-    data_row.put_i32(value.len() as i32);
-
-    // The content
-    data_row.put_slice(value.as_bytes());
-
-    // CommandComplete
-    let mut command_complete = BytesMut::new();
-
-    // Number of rows returned (just one)
-    command_complete.put_slice(&b"SELECT 1\0"[..]);
-
     // The final messages sent to the client
     let mut res = BytesMut::new();
 
     // RowDescription
-    res.put_u8(b'T');
+    res.put(row_description(&vec![(name, DataType::Text)]));
-    res.put_i32(row_desc.len() as i32 + 4);
-    res.put(row_desc);
 
     // DataRow
-    res.put_u8(b'D');
+    res.put(data_row(&vec![value.to_string()]));
-    res.put_i32(data_row.len() as i32 + 4);
-    res.put(data_row);
 
     // CommandComplete
-    res.put_u8(b'C');
+    res.put(command_complete("SELECT 1"));
-    res.put_i32(command_complete.len() as i32 + 4);
-    res.put(command_complete);
 
     // ReadyForQuery
     res.put_u8(b'Z');
@@ -330,6 +296,77 @@ pub async fn show_response(
     write_all_half(stream, res).await
 }
 
+pub fn row_description(columns: &Vec<(&str, DataType)>) -> BytesMut {
+    let mut res = BytesMut::new();
+    let mut row_desc = BytesMut::new();
+
+    // how many colums we are storing
+    row_desc.put_i16(columns.len() as i16);
+
+    for (name, data_type) in columns {
+        // Column name
+        row_desc.put_slice(&format!("{}\0", name).as_bytes());
+
+        // Doesn't belong to any table
+        row_desc.put_i32(0);
+
+        // Doesn't belong to any table
+        row_desc.put_i16(0);
+
+        // Text
+        row_desc.put_i32(data_type.into());
+
+        // Text size = variable (-1)
+        let type_size = match data_type {
+            DataType::Text => -1,
+            DataType::Int4 => 4,
+            DataType::Numeric => -1,
+        };
+
+        row_desc.put_i16(type_size);
+
+        // Type modifier: none that I know
+        row_desc.put_i32(-1);
+
+        // Format being used: text (0), binary (1)
+        row_desc.put_i16(0);
+    }
+
+    res.put_u8(b'T');
+    res.put_i32(row_desc.len() as i32 + 4);
+    res.put(row_desc);
+
+    res
+}
+
+pub fn data_row(row: &Vec<String>) -> BytesMut {
+    let mut res = BytesMut::new();
+    let mut data_row = BytesMut::new();
+
+    data_row.put_i16(row.len() as i16);
+
+    for column in row {
+        let column = column.as_bytes();
+        data_row.put_i32(column.len() as i32);
+        data_row.put_slice(&column);
+    }
+
+    res.put_u8(b'D');
+    res.put_i32(data_row.len() as i32 + 4);
+    res.put(data_row);
+
+    res
+}
+
+pub fn command_complete(command: &str) -> BytesMut {
+    let cmd = BytesMut::from(format!("{}\0", command).as_bytes());
+    let mut res = BytesMut::new();
+    res.put_u8(b'C');
+    res.put_i32(cmd.len() as i32 + 4);
+    res.put(cmd);
+    res
+}
+
 /// Write all data in the buffer to the TcpStream.
 pub async fn write_all(stream: &mut TcpStream, buf: BytesMut) -> Result<(), Error> {
     match stream.write_all(&buf).await {

src/pool.rs

@@ -337,7 +337,7 @@ impl ConnectionPool {
         self.addresses[shard].len()
     }
 
-    pub fn connections(&self, shard: usize, server: usize) -> bb8::State {
+    pub fn pool_state(&self, shard: usize, server: usize) -> bb8::State {
         self.databases[shard][server].state()
     }