Skip to content

Commit eb777eb

Browse files
levkklevkk
authored
Add some test coverage (#622)
1 parent 6762a4e commit eb777eb

File tree

6 files changed

+368
-5
lines changed

6 files changed

+368
-5
lines changed

‎CLAUDE.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
- `cargo check` to test that the code compiles. It shouldn't contain warnings. This is quicker than `cargo build`.
44
- `cargo fmt` to reformat code according to Rust standards.
5-
- `cargo nextest run --test-threads=1 <test name>` to run a specific test
5+
- `cargo nextest run --test-threads=1 <test name>` to run a specific test. Run `pgdog` tests from the `pgdog` directory (`cd pgdog` first).
66
- `cargo nextest run --test-threads=1 --no-fail-fast` to run all tests. Make sure to use `--test-threads=1` because some tests conflict with each other.
77

88
# Code style

‎Cargo.lock‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pgdog/Cargo.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[package]
22
name = "pgdog"
3-
version = "0.1.15"
3+
version = "0.1.16"
44
edition = "2021"
55
description = "Modern PostgreSQL proxy, pooler and load balancer."
66
authors = ["PgDog <hi@pgdog.dev>"]

‎pgdog/src/backend/databases.rs‎

Lines changed: 310 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1187,4 +1187,314 @@ mod tests {
11871187
"Mirror config should not be precomputed when source has no users"
11881188
);
11891189
}
1190+
1191+
#[test]
1192+
fn test_new_pool_fetches_existing_roles_on_reload() {
1193+
use crate::backend::pool::lsn_monitor::LsnStats;
1194+
use crate::backend::replication::publisher::Lsn;
1195+
use std::sync::Arc;
1196+
use tokio::time::Instant;
1197+
1198+
let _lock = lock();
1199+
1200+
let mut config = Config::default();
1201+
config.databases = vec![
1202+
Database {
1203+
name: "testdb".to_string(),
1204+
host: "127.0.0.1".to_string(),
1205+
port: 5432,
1206+
role: Role::Auto,
1207+
shard: 0,
1208+
..Default::default()
1209+
},
1210+
Database {
1211+
name: "testdb".to_string(),
1212+
host: "127.0.0.1".to_string(),
1213+
port: 5433,
1214+
role: Role::Auto,
1215+
shard: 0,
1216+
..Default::default()
1217+
},
1218+
];
1219+
1220+
let users = crate::config::Users {
1221+
users: vec![crate::config::User {
1222+
name: "testuser".to_string(),
1223+
database: "testdb".to_string(),
1224+
password: Some("pass".to_string()),
1225+
..Default::default()
1226+
}],
1227+
..Default::default()
1228+
};
1229+
1230+
let config_and_users = ConfigAndUsers {
1231+
config: config.clone(),
1232+
users: users.clone(),
1233+
config_path: std::path::PathBuf::new(),
1234+
users_path: std::path::PathBuf::new(),
1235+
};
1236+
1237+
let initial_databases = from_config(&config_and_users);
1238+
let cluster = initial_databases.cluster(("testuser", "testdb")).unwrap();
1239+
1240+
for pool in cluster.shards()[0].pools() {
1241+
let lsn_stats = LsnStats {
1242+
replica: pool.addr().database_number == 1,
1243+
lsn: Lsn::from_i64(1000),
1244+
offset_bytes: 1000,
1245+
timestamp: Default::default(),
1246+
fetched: Instant::now(),
1247+
};
1248+
pool.set_lsn_stats(lsn_stats);
1249+
}
1250+
1251+
DATABASES.store(Arc::new(initial_databases));
1252+
1253+
let (_, new_cluster) = new_pool(&users.users[0], &config).unwrap();
1254+
1255+
let roles = new_cluster.shards()[0].current_roles();
1256+
1257+
assert_eq!(roles.len(), 2);
1258+
1259+
assert_eq!(
1260+
roles.get(&0).unwrap().role,
1261+
Role::Primary,
1262+
"database_number 0 should be assigned Primary (replica=false)"
1263+
);
1264+
assert_eq!(
1265+
roles.get(&1).unwrap().role,
1266+
Role::Replica,
1267+
"database_number 1 should be assigned Replica (replica=true)"
1268+
);
1269+
1270+
DATABASES.store(Arc::new(Databases::default()));
1271+
}
1272+
1273+
#[test]
1274+
fn test_new_pool_only_assigns_roles_to_auto() {
1275+
use crate::backend::pool::lsn_monitor::LsnStats;
1276+
use crate::backend::replication::publisher::Lsn;
1277+
use std::sync::Arc;
1278+
use tokio::time::Instant;
1279+
1280+
let _lock = lock();
1281+
1282+
let mut config = Config::default();
1283+
config.databases = vec![
1284+
Database {
1285+
name: "testdb".to_string(),
1286+
host: "127.0.0.1".to_string(),
1287+
port: 5432,
1288+
role: Role::Primary,
1289+
shard: 0,
1290+
..Default::default()
1291+
},
1292+
Database {
1293+
name: "testdb".to_string(),
1294+
host: "127.0.0.1".to_string(),
1295+
port: 5433,
1296+
role: Role::Replica,
1297+
shard: 0,
1298+
..Default::default()
1299+
},
1300+
Database {
1301+
name: "testdb".to_string(),
1302+
host: "127.0.0.1".to_string(),
1303+
port: 5434,
1304+
role: Role::Auto,
1305+
shard: 0,
1306+
..Default::default()
1307+
},
1308+
];
1309+
1310+
let users = crate::config::Users {
1311+
users: vec![crate::config::User {
1312+
name: "testuser".to_string(),
1313+
database: "testdb".to_string(),
1314+
password: Some("pass".to_string()),
1315+
..Default::default()
1316+
}],
1317+
..Default::default()
1318+
};
1319+
1320+
let config_and_users = ConfigAndUsers {
1321+
config: config.clone(),
1322+
users: users.clone(),
1323+
config_path: std::path::PathBuf::new(),
1324+
users_path: std::path::PathBuf::new(),
1325+
};
1326+
1327+
let initial_databases = from_config(&config_and_users);
1328+
let cluster = initial_databases.cluster(("testuser", "testdb")).unwrap();
1329+
1330+
for pool in cluster.shards()[0].pools() {
1331+
let db_num = pool.addr().database_number;
1332+
let lsn_stats = LsnStats {
1333+
replica: db_num != 0,
1334+
lsn: Lsn::from_i64(1000),
1335+
offset_bytes: 1000,
1336+
timestamp: Default::default(),
1337+
fetched: Instant::now(),
1338+
};
1339+
pool.set_lsn_stats(lsn_stats);
1340+
}
1341+
1342+
DATABASES.store(Arc::new(initial_databases));
1343+
1344+
let (_, new_cluster) = new_pool(&users.users[0], &config).unwrap();
1345+
1346+
let roles = new_cluster.shards()[0].current_roles();
1347+
1348+
assert_eq!(roles.len(), 3);
1349+
1350+
assert_eq!(
1351+
roles.get(&0).unwrap().role,
1352+
Role::Primary,
1353+
"Explicit Primary should remain Primary (LSN says replica=false)"
1354+
);
1355+
assert_eq!(
1356+
roles.get(&1).unwrap().role,
1357+
Role::Replica,
1358+
"Explicit Replica should remain Replica (even though LSN says replica=true which would suggest Replica, the explicit config is preserved)"
1359+
);
1360+
assert_eq!(
1361+
roles.get(&2).unwrap().role,
1362+
Role::Replica,
1363+
"Auto role should be assigned Replica based on LSN replica=true"
1364+
);
1365+
1366+
DATABASES.store(Arc::new(Databases::default()));
1367+
}
1368+
1369+
#[test]
1370+
fn test_new_pool_matches_roles_by_database_number() {
1371+
use crate::backend::pool::lsn_monitor::LsnStats;
1372+
use crate::backend::replication::publisher::Lsn;
1373+
use std::sync::Arc;
1374+
use tokio::time::Instant;
1375+
1376+
let _lock = lock();
1377+
1378+
let mut config = Config::default();
1379+
config.databases = vec![
1380+
Database {
1381+
name: "db1".to_string(),
1382+
host: "127.0.0.1".to_string(),
1383+
port: 5432,
1384+
role: Role::Auto,
1385+
shard: 0,
1386+
..Default::default()
1387+
},
1388+
Database {
1389+
name: "db2".to_string(),
1390+
host: "127.0.0.1".to_string(),
1391+
port: 5433,
1392+
role: Role::Auto,
1393+
shard: 0,
1394+
..Default::default()
1395+
},
1396+
Database {
1397+
name: "db1".to_string(),
1398+
host: "127.0.0.1".to_string(),
1399+
port: 5434,
1400+
role: Role::Auto,
1401+
shard: 0,
1402+
..Default::default()
1403+
},
1404+
];
1405+
1406+
let users = crate::config::Users {
1407+
users: vec![
1408+
crate::config::User {
1409+
name: "user1".to_string(),
1410+
database: "db1".to_string(),
1411+
password: Some("pass".to_string()),
1412+
..Default::default()
1413+
},
1414+
crate::config::User {
1415+
name: "user2".to_string(),
1416+
database: "db2".to_string(),
1417+
password: Some("pass".to_string()),
1418+
..Default::default()
1419+
},
1420+
],
1421+
..Default::default()
1422+
};
1423+
1424+
let config_and_users = ConfigAndUsers {
1425+
config: config.clone(),
1426+
users: users.clone(),
1427+
config_path: std::path::PathBuf::new(),
1428+
users_path: std::path::PathBuf::new(),
1429+
};
1430+
1431+
let initial_databases = from_config(&config_and_users);
1432+
1433+
let db1_cluster = initial_databases.cluster(("user1", "db1")).unwrap();
1434+
for pool in db1_cluster.shards()[0].pools() {
1435+
let db_num = pool.addr().database_number;
1436+
let lsn_stats = LsnStats {
1437+
replica: db_num == 2,
1438+
lsn: Lsn::from_i64(1000),
1439+
offset_bytes: 1000,
1440+
timestamp: Default::default(),
1441+
fetched: Instant::now(),
1442+
};
1443+
pool.set_lsn_stats(lsn_stats);
1444+
}
1445+
1446+
let db2_cluster = initial_databases.cluster(("user2", "db2")).unwrap();
1447+
for pool in db2_cluster.shards()[0].pools() {
1448+
let lsn_stats = LsnStats {
1449+
replica: false,
1450+
lsn: Lsn::from_i64(1000),
1451+
offset_bytes: 1000,
1452+
timestamp: Default::default(),
1453+
fetched: Instant::now(),
1454+
};
1455+
pool.set_lsn_stats(lsn_stats);
1456+
}
1457+
1458+
DATABASES.store(Arc::new(initial_databases));
1459+
1460+
let (_, new_db1_cluster) = new_pool(&users.users[0], &config).unwrap();
1461+
let db1_roles = new_db1_cluster.shards()[0].current_roles();
1462+
1463+
assert_eq!(db1_roles.len(), 2);
1464+
assert!(
1465+
db1_roles.get(&0).is_some(),
1466+
"db1 should have database_number 0"
1467+
);
1468+
assert!(
1469+
db1_roles.get(&2).is_some(),
1470+
"db1 should have database_number 2"
1471+
);
1472+
1473+
assert_eq!(
1474+
db1_roles.get(&0).unwrap().role,
1475+
Role::Primary,
1476+
"database_number 0 should be Primary (replica=false)"
1477+
);
1478+
assert_eq!(
1479+
db1_roles.get(&2).unwrap().role,
1480+
Role::Replica,
1481+
"database_number 2 should be Replica (replica=true)"
1482+
);
1483+
1484+
let (_, new_db2_cluster) = new_pool(&users.users[1], &config).unwrap();
1485+
let db2_roles = new_db2_cluster.shards()[0].current_roles();
1486+
1487+
assert_eq!(db2_roles.len(), 1);
1488+
assert!(
1489+
db2_roles.get(&1).is_some(),
1490+
"db2 should have database_number 1"
1491+
);
1492+
assert_eq!(
1493+
db2_roles.get(&1).unwrap().role,
1494+
Role::Primary,
1495+
"database_number 1 should be Primary (replica=false)"
1496+
);
1497+
1498+
DATABASES.store(Arc::new(Databases::default()));
1499+
}
11901500
}

‎pgdog/src/backend/pool/pool_impl.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -421,6 +421,12 @@ impl Pool {
421421
self.lock().config = config;
422422
}
423423

424+
/// Set LSN stats for testing.
425+
#[cfg(test)]
426+
pub(crate) fn set_lsn_stats(&self, stats: LsnStats) {
427+
*self.inner().lsn_stats.write() = stats;
428+
}
429+
424430
/// Fetch OIDs for user-defined data types.
425431
pub fn oids(&self) -> Option<Oids> {
426432
self.lock().oids

0 commit comments

Comments
 (0)