diff --git a/mise.toml b/mise.toml index 0d2f0ed0..fe216b4a 100644 --- a/mise.toml +++ b/mise.toml @@ -1,3 +1,4 @@ [tools] +"rust" = { version = "1.91.0" } "cargo:cargo-nextest" = "latest" "cargo:cargo-watch" = "latest" \ No newline at end of file diff --git a/pgdog/src/backend/pool/lb/mod.rs b/pgdog/src/backend/pool/lb/mod.rs index 784f192e..be9f6591 100644 --- a/pgdog/src/backend/pool/lb/mod.rs +++ b/pgdog/src/backend/pool/lb/mod.rs @@ -277,7 +277,10 @@ impl LoadBalancer { let primary_reads = match self.rw_split { IncludePrimary => true, IncludePrimaryIfReplicaBanned => candidates.iter().any(|target| target.ban.banned()), - ExcludePrimary => false, + // we read from the primary if we have no replicas + ExcludePrimary => !candidates + .iter() + .any(|target| target.role() == Role::Replica), }; if !primary_reads { diff --git a/pgdog/src/backend/pool/lb/test.rs b/pgdog/src/backend/pool/lb/test.rs index 31d29a48..0a32c1ae 100644 --- a/pgdog/src/backend/pool/lb/test.rs +++ b/pgdog/src/backend/pool/lb/test.rs @@ -393,6 +393,42 @@ async fn test_read_write_split_include_primary() { replicas.shutdown(); } +#[tokio::test] +async fn test_read_write_split_exclude_primary_no_replicas() { + let primary_config = create_test_pool_config("127.0.0.1", 5432); + let primary_pool = Pool::new(&primary_config); + primary_pool.launch(); + + let replica_configs = []; + + let replicas = LoadBalancer::new( + &Some(primary_pool), + &replica_configs, + LoadBalancingStrategy::RoundRobin, + ReadWriteSplit::ExcludePrimary, + ); + replicas.launch(); + + let request = Request::default(); + + // Try getting connections multiple times and we have primary in the set + let mut used_pool_ids = HashSet::new(); + for _ in 0..2 { + let conn = replicas.get(&request).await.unwrap(); + used_pool_ids.insert(conn.pool.id()); + } + + // Should use only primary + assert_eq!(used_pool_ids.len(), 1); + + // Verify primary pool ID is in the set of used pools + let primary_id = replicas.primary().unwrap().id(); + assert!(used_pool_ids.contains(&primary_id)); + + // Shutdown + replicas.shutdown(); +} + #[tokio::test] async fn test_read_write_split_exclude_primary_no_primary() { // Test exclude primary setting when no primary exists