From 26f75f8d5d655825d2dd83005672b543a26705bc Mon Sep 17 00:00:00 2001
From: Lev Kokotov <levkk@users.noreply.github.com>
Date: Sun, 27 Feb 2022 10:21:24 -0800
Subject: [PATCH] admin RELOAD (#49)

* admin RELOAD

* test
---
 .circleci/run_tests.sh |  3 +++
 src/admin.rs           | 35 ++++++++++++++++++++++++++++++++++-
 src/config.rs          |  7 ++++++-
 3 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/.circleci/run_tests.sh b/.circleci/run_tests.sh
index 93827f9..abc1a7f 100644
--- a/.circleci/run_tests.sh
+++ b/.circleci/run_tests.sh
@@ -86,6 +86,9 @@ sed -i 's/pool_mode = "transaction"/pool_mode = "session"/' pgcat.toml
 # Reload config test
 kill -SIGHUP $(pgrep pgcat)
 
+# Reload again with the admin database
+psql -h 127.0.0.1 -p 6432 -d pgbouncer -c 'RELOAD' > /dev/null
+
 # Prepared statements that will only work in session mode
 pgbench -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol prepared
 
diff --git a/src/admin.rs b/src/admin.rs
index a19febd..97346ed 100644
--- a/src/admin.rs
+++ b/src/admin.rs
@@ -1,7 +1,8 @@
 use bytes::{Buf, BufMut, BytesMut};
-use log::trace;
+use log::{info, trace};
 use tokio::net::tcp::OwnedWriteHalf;
 
+use crate::config::{get_config, parse};
 use crate::constants::{OID_NUMERIC, OID_TEXT};
 use crate::errors::Error;
 use crate::messages::write_all_half;
@@ -23,11 +24,43 @@ pub async fn handle_admin(stream: &mut OwnedWriteHalf, mut query: BytesMut) -> R
     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 {
         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
+}
+
 /// SHOW STATS
 pub async fn show_stats(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
     let columns = [
diff --git a/src/config.rs b/src/config.rs
index 754dc4f..3cb413a 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -134,6 +134,7 @@ impl Default for QueryRouter {
 
 #[derive(Deserialize, Debug, Clone)]
 pub struct Config {
+    pub path: Option<String>,
     pub general: General,
     pub user: User,
     pub shards: HashMap<String, Shard>,
@@ -143,6 +144,7 @@ pub struct Config {
 impl Default for Config {
     fn default() -> Config {
         Config {
+            path: Some(String::from("pgcat.toml")),
             general: General::default(),
             user: User::default(),
             shards: HashMap::from([(String::from("1"), Shard::default())]),
@@ -189,7 +191,7 @@ pub async fn parse(path: &str) -> Result<(), Error> {
         }
     };
 
-    let config: Config = match toml::from_str(&contents) {
+    let mut config: Config = match toml::from_str(&contents) {
         Ok(config) => config,
         Err(err) => {
             error!("Could not parse config file: {}", err.to_string());
@@ -279,6 +281,8 @@ pub async fn parse(path: &str) -> Result<(), Error> {
         }
     };
 
+    config.path = Some(path.to_string());
+
     CONFIG.store(Arc::new(config.clone()));
 
     Ok(())
@@ -296,5 +300,6 @@ mod test {
         assert_eq!(get_config().shards["1"].servers[0].0, "127.0.0.1");
         assert_eq!(get_config().shards["0"].servers[0].2, "primary");
         assert_eq!(get_config().query_router.default_role, "any");
+        assert_eq!(get_config().path, Some("pgcat.toml".to_string()));
     }
 }