From d64f6793c1e993c30b79735ac810769f81c048f5 Mon Sep 17 00:00:00 2001
From: Pradeep Chhetri <30620077+chhetripradeep@users.noreply.github.com>
Date: Wed, 17 Aug 2022 01:01:46 +0800
Subject: [PATCH] Minor cleanup in admin command (#126)

* Minor cleanup in admin command

* Typo correction

* fix when the admin query is ending with semicolon
---
 examples/docker/pgcat.toml |  5 +++
 pgcat.toml                 |  2 +-
 src/admin.rs               | 65 +++++++++++++++++++++++---------------
 3 files changed, 45 insertions(+), 27 deletions(-)

diff --git a/examples/docker/pgcat.toml b/examples/docker/pgcat.toml
index 3c74df3..85a3f46 100644
--- a/examples/docker/pgcat.toml
+++ b/examples/docker/pgcat.toml
@@ -92,10 +92,14 @@ password = "postgres"
 # is the sum of pool_size across all users.
 pool_size = 9
 
+# Maximum query duration. Dangerous, but protects against DBs that died in a non-obvious way.
+statement_timeout = 0
+
 [pools.sharded.users.1]
 username = "postgres"
 password = "postgres"
 pool_size = 21
+statement_timeout = 15000
 
 # Shard 0
 [pools.sharded.shards.0]
@@ -133,6 +137,7 @@ sharding_function = "pg_bigint_hash"
 username = "postgres"
 password = "postgres"
 pool_size = 5
+statement_timeout = 0
 
 [pools.simple_db.shards.0]
 servers = [
diff --git a/pgcat.toml b/pgcat.toml
index 2976118..9125afd 100644
--- a/pgcat.toml
+++ b/pgcat.toml
@@ -92,7 +92,7 @@ password = "sharding_user"
 # is the sum of pool_size across all users.
 pool_size = 9
 
-# Maximum query duration. Dangerous, but protetcts against DBs that died and a non-obvious way.
+# Maximum query duration. Dangerous, but protects against DBs that died in a non-obvious way.
 statement_timeout = 0
 
 [pools.sharded_db.users.1]
diff --git a/src/admin.rs b/src/admin.rs
index b82b294..6a79e49 100644
--- a/src/admin.rs
+++ b/src/admin.rs
@@ -44,32 +44,45 @@ where
 
     trace!("Admin query: {}", query);
 
-    if query.starts_with("SHOW STATS") {
-        trace!("SHOW STATS");
-        show_stats(stream).await
-    } else if query.starts_with("RELOAD") {
-        trace!("RELOAD");
-        reload(stream, client_server_map).await
-    } else if query.starts_with("SHOW CONFIG") {
-        trace!("SHOW CONFIG");
-        show_config(stream).await
-    } else if query.starts_with("SHOW DATABASES") {
-        trace!("SHOW DATABASES");
-        show_databases(stream).await
-    } else if query.starts_with("SHOW POOLS") {
-        trace!("SHOW POOLS");
-        show_pools(stream).await
-    } else if query.starts_with("SHOW LISTS") {
-        trace!("SHOW LISTS");
-        show_lists(stream).await
-    } else if query.starts_with("SHOW VERSION") {
-        trace!("SHOW VERSION");
-        show_version(stream).await
-    } else if query.starts_with("SET ") {
-        trace!("SET");
-        ignore_set(stream).await
-    } else {
-        error_response(stream, "Unsupported query against the admin database").await
+    let query_parts: Vec<&str> = query.trim_end_matches(';').split_whitespace().collect();
+
+    match query_parts[0] {
+        "RELOAD" => {
+            trace!("RELOAD");
+            reload(stream, client_server_map).await
+        }
+        "SET" => {
+            trace!("SET");
+            ignore_set(stream).await
+        }
+        "SHOW" => match query_parts[1] {
+            "CONFIG" => {
+                trace!("SHOW CONFIG");
+                show_config(stream).await
+            }
+            "DATABASES" => {
+                trace!("SHOW DATABASES");
+                show_databases(stream).await
+            }
+            "LISTS" => {
+                trace!("SHOW LISTS");
+                show_lists(stream).await
+            }
+            "POOLS" => {
+                trace!("SHOW POOLS");
+                show_pools(stream).await
+            }
+            "STATS" => {
+                trace!("SHOW STATS");
+                show_stats(stream).await
+            }
+            "VERSION" => {
+                trace!("SHOW VERSION");
+                show_version(stream).await
+            }
+            _ => error_response(stream, "Unsupported SHOW query against the admin database").await,
+        },
+        _ => error_response(stream, "Unsupported query against the admin database").await,
     }
 }