From 75c2dc5e17b959290a655bdfed7275592535516c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9D=8E=E8=BE=BE=E5=8D=83?=
 <128981031+LDQONJ@users.noreply.github.com>
Date: Wed, 26 Aug 2026 14:43:52 +0800
Subject: [PATCH] fix(mongo): do not force directConnection when loadBalanced

---
 crates/dbx-core/src/db/mongo_driver.rs   |  2 +-
 crates/dbx-core/src/models/connection.rs | 52 ++++++++++++++++++++++--
 2 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/crates/dbx-core/src/db/mongo_driver.rs b/crates/dbx-core/src/db/mongo_driver.rs
index a33b71d229..138e8dffd7 100644
--- a/crates/dbx-core/src/db/mongo_driver.rs
+++ b/crates/dbx-core/src/db/mongo_driver.rs
@@ -111,7 +111,7 @@ pub async fn connect_with_oidc(
         // set discovery. This is essential when connecting through a TCP proxy
         // or NAT where the driver would otherwise receive internal IPs from
         // the replica set handshake and fail to connect.
-        if !is_multi_host {
+        if !is_multi_host && options.load_balanced != Some(true) && options.direct_connection.is_none() {
             options.direct_connection = Some(true);
         }
         Client::with_options(options).map_err(|e| format!("MongoDB connection failed: {e}"))
diff --git a/crates/dbx-core/src/models/connection.rs b/crates/dbx-core/src/models/connection.rs
index fc04c98bf1..c4d0d39fa4 100644
--- a/crates/dbx-core/src/models/connection.rs
+++ b/crates/dbx-core/src/models/connection.rs
@@ -1011,7 +1011,10 @@ impl ConnectionConfig {
                     return cs;
                 }
                 let mut suffix = if params.is_empty() { String::new() } else { format!("?{params}") };
-                if is_tunneled && !suffix.contains("directConnection=") {
+                if is_tunneled
+                    && !suffix.contains("directConnection=")
+                    && !mongo_url_params_have_load_balanced_true(&suffix)
+                {
                     if suffix.is_empty() {
                         suffix = "?directConnection=true".to_string();
                     } else {
@@ -1179,7 +1182,10 @@ impl ConnectionConfig {
                     return cs;
                 }
                 let mut suffix = if params.is_empty() { String::new() } else { format!("?{params}") };
-                if is_tunneled && !suffix.contains("directConnection=") {
+                if is_tunneled
+                    && !suffix.contains("directConnection=")
+                    && !mongo_url_params_have_load_balanced_true(&suffix)
+                {
                     if suffix.is_empty() {
                         suffix = "?directConnection=true".to_string();
                     } else {
@@ -2042,6 +2048,24 @@ fn mongo_url_param_is_direct_connection_true(part: &str) -> bool {
         && percent_decode_str(value).decode_utf8_lossy().eq_ignore_ascii_case("true")
 }
 
+fn mongo_url_param_is_load_balanced_true(part: &str) -> bool {
+    let Some((key, value)) = part.split_once('=') else {
+        return false;
+    };
+    percent_decode_str(key).decode_utf8_lossy().eq_ignore_ascii_case("loadBalanced")
+        && percent_decode_str(value).decode_utf8_lossy().eq_ignore_ascii_case("true")
+}
+
+fn mongo_uri_has_load_balanced_true(uri: &str) -> bool {
+    uri.split_once('?')
+        .map(|(_, query)| query.split('#').next().unwrap_or("").split('&').any(mongo_url_param_is_load_balanced_true))
+        .unwrap_or(false)
+}
+
+fn mongo_url_params_have_load_balanced_true(params: &str) -> bool {
+    params.trim_start_matches('?').split('&').any(mongo_url_param_is_load_balanced_true)
+}
+
 fn normalize_postgres_url_params(value: &str, force_tls: bool) -> String {
     let value = value.trim_start_matches('?');
 
@@ -2284,7 +2308,9 @@ fn rewrite_mongo_uri_host(uri: &str, new_host: &str, new_port: u16) -> String {
 
     let mut result = format!("mongodb://{creds_prefix}{new_host}:{new_port}{after_hosts}");
 
-    if !result.contains("directConnection=") {
+    // loadBalanced=true requires directConnection to be unset or false, so a
+    // tunneled load-balanced endpoint must keep its original option set.
+    if !result.contains("directConnection=") && !mongo_uri_has_load_balanced_true(&result) {
         if result.contains('?') {
             result.push_str("&directConnection=true");
         } else {
@@ -4223,6 +4249,26 @@ mod tests {
         );
     }
 
+    #[test]
+    fn mongodb_tunneled_connection_string_keeps_load_balanced_without_direct_connection() {
+        let mut config = mongodb_config("root", "secret", Some("admin"));
+        config.connection_string = Some("mongodb://read:pass@lb.example.net:27017/admin?loadBalanced=true".to_string());
+
+        let url = config.connection_url_with_host("127.0.0.1", 54321);
+
+        assert_eq!(url, "mongodb://read:pass@127.0.0.1:54321/admin?loadBalanced=true");
+    }
+
+    #[test]
+    fn mongodb_tunneled_form_url_keeps_load_balanced_without_direct_connection() {
+        let mut config = mongodb_config("root", "secret", Some("admin"));
+        config.url_params = Some("loadBalanced=true&authSource=admin".to_string());
+
+        let url = config.connection_url_with_host("127.0.0.1", 54321);
+
+        assert_eq!(url, "mongodb://root:secret@127.0.0.1:54321/admin?loadBalanced=true&authSource=admin");
+    }
+
     #[test]
     fn mongodb_connection_string_unchanged_when_not_tunneled() {
         let mut config = mongodb_config("root", "secret", Some("admin"));