From 32b913af94eab61418f47e6ae9598900afe6eaad Mon Sep 17 00:00:00 2001
From: Lev Kokotov <lev.kokotov@gmail.com>
Date: Thu, 30 Mar 2023 12:09:22 -0700
Subject: [PATCH] update admin

---
 src/admin.rs      | 11 ++++++++++-
 src/prometheus.rs | 10 +++++-----
 src/stats.rs      |  6 ++----
 src/stats/pool.rs |  7 +++++++
 4 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/src/admin.rs b/src/admin.rs
index 4df2206..4cca2ce 100644
--- a/src/admin.rs
+++ b/src/admin.rs
@@ -259,6 +259,7 @@ where
     let columns = vec![
         ("database", DataType::Text),
         ("user", DataType::Text),
+        ("secret", DataType::Text),
         ("pool_mode", DataType::Text),
         ("cl_idle", DataType::Numeric),
         ("cl_active", DataType::Numeric),
@@ -276,10 +277,11 @@ where
     let mut res = BytesMut::new();
     res.put(row_description(&columns));
 
-    for ((_user_pool, _pool), pool_stats) in all_pool_stats {
+    for (_, pool_stats) in all_pool_stats {
         let mut row = vec![
             pool_stats.database(),
             pool_stats.user(),
+            pool_stats.redacted_secret(),
             pool_stats.pool_mode().to_string(),
         ];
         pool_stats.populate_row(&mut row);
@@ -895,13 +897,20 @@ where
     res.put(row_description(&vec![
         ("name", DataType::Text),
         ("pool_mode", DataType::Text),
+        ("secret", DataType::Text),
     ]));
 
     for (user_pool, pool) in get_all_pools() {
         let pool_config = &pool.settings;
+        let redacted_secret = match user_pool.secret {
+            Some(secret) => format!("****{}", &secret[secret.len() - 4..]),
+            None => "<no secret>".to_string(),
+        };
+
         res.put(data_row(&vec![
             user_pool.user.clone(),
             pool_config.pool_mode.to_string(),
+            redacted_secret,
         ]));
     }
 
diff --git a/src/prometheus.rs b/src/prometheus.rs
index 6e578bf..6751790 100644
--- a/src/prometheus.rs
+++ b/src/prometheus.rs
@@ -9,7 +9,7 @@ use std::sync::atomic::Ordering;
 use std::sync::Arc;
 
 use crate::config::Address;
-use crate::pool::get_all_pools;
+use crate::pool::{get_all_pools, PoolIdentifier};
 use crate::stats::{get_pool_stats, get_server_stats, ServerStats};
 
 struct MetricHelpType {
@@ -233,10 +233,10 @@ impl<Value: fmt::Display> PrometheusMetric<Value> {
         Self::from_name(&format!("stats_{}", name), value, labels)
     }
 
-    fn from_pool(pool: &(String, String), name: &str, value: u64) -> Option<PrometheusMetric<u64>> {
+    fn from_pool(pool: &PoolIdentifier, name: &str, value: u64) -> Option<PrometheusMetric<u64>> {
         let mut labels = HashMap::new();
-        labels.insert("pool", pool.0.clone());
-        labels.insert("user", pool.1.clone());
+        labels.insert("pool", pool.db.clone());
+        labels.insert("user", pool.user.clone());
 
         Self::from_name(&format!("pools_{}", name), value, labels)
     }
@@ -294,7 +294,7 @@ fn push_pool_stats(lines: &mut Vec<String>) {
             } else {
                 warn!(
                     "Metric {} not implemented for ({},{})",
-                    name, pool.0, pool.1
+                    name, pool.db, pool.user
                 );
             }
         }
diff --git a/src/stats.rs b/src/stats.rs
index 5b7895b..63488de 100644
--- a/src/stats.rs
+++ b/src/stats.rs
@@ -22,7 +22,7 @@ pub use server::{ServerState, ServerStats};
 /// Convenience types for various stats
 type ClientStatesLookup = HashMap<i32, Arc<ClientStats>>;
 type ServerStatesLookup = HashMap<i32, Arc<ServerStats>>;
-type PoolStatsLookup = HashMap<(String, String), Arc<PoolStats>>;
+type PoolStatsLookup = HashMap<PoolIdentifier, Arc<PoolStats>>;
 
 /// Stats for individual client connections
 /// Used in SHOW CLIENTS.
@@ -83,9 +83,7 @@ impl Reporter {
 
     /// Register a pool with the stats system.
     fn pool_register(&self, identifier: PoolIdentifier, stats: Arc<PoolStats>) {
-        POOL_STATS
-            .write()
-            .insert((identifier.db, identifier.user), stats);
+        POOL_STATS.write().insert(identifier, stats);
     }
 }
 
diff --git a/src/stats/pool.rs b/src/stats/pool.rs
index 1b01ef2..51ec78f 100644
--- a/src/stats/pool.rs
+++ b/src/stats/pool.rs
@@ -102,6 +102,13 @@ impl PoolStats {
         self.identifier.user.clone()
     }
 
+    pub fn redacted_secret(&self) -> String {
+        match self.identifier.secret {
+            Some(ref s) => format!("****{}", &s[s.len() - 4..]),
+            None => "<no secret>".to_string(),
+        }
+    }
+
     pub fn pool_mode(&self) -> PoolMode {
         self.config.pool_mode
     }