Skip to content

Commit 7987c5f

Browse files
authoredOct 1, 2022
Replace a few types with more developer-friendly names (#182)
* Replace a few types with more developer-friendly names * UserPool -> PoolIdentifier
1 parent 24f5eec commit 7987c5f

File tree

4 files changed

+49
-20
lines changed

4 files changed

+49
-20
lines changed
 

‎src/admin.rs

+6-6
Original file line numberDiff line numberDiff line change
@@ -215,16 +215,16 @@ where
215215

216216
let mut res = BytesMut::new();
217217
res.put(row_description(&columns));
218-
for ((pool_name, username), pool) in get_all_pools() {
218+
for (user_pool, pool) in get_all_pools() {
219219
let def = HashMap::default();
220220
let pool_stats = all_pool_stats
221-
.get(&(pool_name.clone(), username.clone()))
221+
.get(&(user_pool.db.clone(), user_pool.user.clone()))
222222
.unwrap_or(&def);
223223

224224
let pool_config = &pool.settings;
225225
let mut row = vec![
226-
pool_name.clone(),
227-
username.clone(),
226+
user_pool.db.clone(),
227+
user_pool.user.clone(),
228228
pool_config.pool_mode.to_string(),
229229
];
230230
for column in &columns[3..columns.len()] {
@@ -420,7 +420,7 @@ where
420420
let mut res = BytesMut::new();
421421
res.put(row_description(&columns));
422422

423-
for ((db, username), pool) in get_all_pools() {
423+
for (user_pool, pool) in get_all_pools() {
424424
for shard in 0..pool.shards() {
425425
for server in 0..pool.servers(shard) {
426426
let address = pool.address(shard, server);
@@ -429,7 +429,7 @@ where
429429
None => HashMap::new(),
430430
};
431431

432-
let mut row = vec![address.name(), db.clone(), username.clone()];
432+
let mut row = vec![address.name(), user_pool.db.clone(), user_pool.user.clone()];
433433
for column in &columns[3..] {
434434
row.push(stats.get(column.0).unwrap_or(&0).to_string());
435435
}

‎src/client.rs

+2-2
Original file line numberDiff line numberDiff line change
@@ -446,7 +446,7 @@ where
446446
}
447447
// Authenticate normal user.
448448
else {
449-
let pool = match get_pool(pool_name.clone(), username.clone()) {
449+
let pool = match get_pool(&pool_name, &username) {
450450
Some(pool) => pool,
451451
None => {
452452
error_response(
@@ -648,7 +648,7 @@ where
648648
// Get a pool instance referenced by the most up-to-date
649649
// pointer. This ensures we always read the latest config
650650
// when starting a query.
651-
let pool = match get_pool(self.pool_name.clone(), self.username.clone()) {
651+
let pool = match get_pool(&self.pool_name, &self.username) {
652652
Some(pool) => pool,
653653
None => {
654654
error_response(

‎src/pool.rs

+38-9
Original file line numberDiff line numberDiff line change
@@ -19,16 +19,43 @@ use crate::server::Server;
1919
use crate::sharding::ShardingFunction;
2020
use crate::stats::{get_reporter, Reporter};
2121

22+
pub type ProcessId = i32;
23+
pub type SecretKey = i32;
24+
pub type ServerHost = String;
25+
pub type ServerPort = u16;
26+
2227
pub type BanList = Arc<RwLock<Vec<HashMap<Address, NaiveDateTime>>>>;
23-
pub type ClientServerMap = Arc<Mutex<HashMap<(i32, i32), (i32, i32, String, u16)>>>;
24-
pub type PoolMap = HashMap<(String, String), ConnectionPool>;
28+
pub type ClientServerMap =
29+
Arc<Mutex<HashMap<(ProcessId, SecretKey), (ProcessId, SecretKey, ServerHost, ServerPort)>>>;
30+
pub type PoolMap = HashMap<PoolIdentifier, ConnectionPool>;
2531
/// The connection pool, globally available.
2632
/// This is atomic and safe and read-optimized.
2733
/// The pool is recreated dynamically when the config is reloaded.
2834
pub static POOLS: Lazy<ArcSwap<PoolMap>> = Lazy::new(|| ArcSwap::from_pointee(HashMap::default()));
2935
static POOLS_HASH: Lazy<ArcSwap<HashSet<crate::config::Pool>>> =
3036
Lazy::new(|| ArcSwap::from_pointee(HashSet::default()));
3137

38+
/// An identifier for a PgCat pool,
39+
/// a database visible to clients.
40+
#[derive(Hash, Debug, Clone, PartialEq, Eq)]
41+
pub struct PoolIdentifier {
42+
// The name of the database clients want to connect to.
43+
pub db: String,
44+
45+
/// The username the client connects with. Each user gets its own pool.
46+
pub user: String,
47+
}
48+
49+
impl PoolIdentifier {
50+
/// Create a new user/pool identifier.
51+
pub fn new(db: &str, user: &str) -> PoolIdentifier {
52+
PoolIdentifier {
53+
db: db.to_string(),
54+
user: user.to_string(),
55+
}
56+
}
57+
}
58+
3259
/// Pool settings.
3360
#[derive(Clone, Debug)]
3461
pub struct PoolSettings {
@@ -113,14 +140,16 @@ impl ConnectionPool {
113140
// If the pool hasn't changed, get existing reference and insert it into the new_pools.
114141
// We replace all pools at the end, but if the reference is kept, the pool won't get re-created (bb8).
115142
if !changed {
116-
match get_pool(pool_name.clone(), user.username.clone()) {
143+
match get_pool(&pool_name, &user.username) {
117144
Some(pool) => {
118145
info!(
119146
"[pool: {}][user: {}] has not changed",
120147
pool_name, user.username
121148
);
122-
new_pools
123-
.insert((pool_name.clone(), user.username.clone()), pool.clone());
149+
new_pools.insert(
150+
PoolIdentifier::new(&pool_name, &user.username),
151+
pool.clone(),
152+
);
124153
continue;
125154
}
126155
None => (),
@@ -239,7 +268,7 @@ impl ConnectionPool {
239268
};
240269

241270
// There is one pool per database/user pair.
242-
new_pools.insert((pool_name.clone(), user.username.clone()), pool);
271+
new_pools.insert(PoolIdentifier::new(&pool_name, &user.username), pool);
243272
}
244273
}
245274

@@ -603,15 +632,15 @@ impl ManageConnection for ServerPool {
603632
}
604633

605634
/// Get the connection pool
606-
pub fn get_pool(db: String, user: String) -> Option<ConnectionPool> {
607-
match get_all_pools().get(&(db, user)) {
635+
pub fn get_pool(db: &str, user: &str) -> Option<ConnectionPool> {
636+
match get_all_pools().get(&PoolIdentifier::new(&db, &user)) {
608637
Some(pool) => Some(pool.clone()),
609638
None => None,
610639
}
611640
}
612641

613642
/// Get a pointer to all configured pools.
614-
pub fn get_all_pools() -> HashMap<(String, String), ConnectionPool> {
643+
pub fn get_all_pools() -> HashMap<PoolIdentifier, ConnectionPool> {
615644
return (*(*POOLS.load())).clone();
616645
}
617646

‎src/stats.rs

+3-3
Original file line numberDiff line numberDiff line change
@@ -521,11 +521,11 @@ impl Collector {
521521
tokio::time::interval(tokio::time::Duration::from_millis(STAT_PERIOD / 15));
522522
loop {
523523
interval.tick().await;
524-
for ((pool_name, username), _pool) in get_all_pools() {
524+
for (user_pool, _) in get_all_pools() {
525525
let _ = tx.try_send(Event {
526526
name: EventName::UpdateStats {
527-
pool_name,
528-
username,
527+
pool_name: user_pool.db,
528+
username: user_pool.user,
529529
},
530530
value: 0,
531531
});

0 commit comments

Comments
 (0)
Please sign in to comment.