From 7c0a8ed01d96d491d4a87274cf205e1aeb451201 Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Tue, 26 May 2026 14:59:17 +0800
Subject: [PATCH] fix: scope postgres-like agent connections to selected
 database

---
 crates/dbx-core/src/connection.rs | 70 +++++++++++++++++++++++++++++++
 1 file changed, 70 insertions(+)

diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs
index c2b866cce..f93b784af 100644
--- a/crates/dbx-core/src/connection.rs
+++ b/crates/dbx-core/src/connection.rs
@@ -634,6 +634,8 @@ pub fn agent_connect_params(config: &ConnectionConfig, host: &str, port: u16, da
         config.connection_url_with_host(host, port)
     } else if config.db_type == DatabaseType::Oracle {
         oracle_jdbc_connection_string(config, host, port, database)
+    } else if matches!(config.db_type, DatabaseType::Kingbase | DatabaseType::Highgo | DatabaseType::Vastbase) {
+        postgres_like_agent_jdbc_connection_string(config, host, port, database)
     } else if config.db_type == DatabaseType::SapHana {
         sap_hana_jdbc_connection_string(config, host, port, database)
     } else {
@@ -709,6 +711,22 @@ fn oracle_jdbc_connection_string(config: &ConnectionConfig, host: &str, port: u1
     }
 }
 
+fn postgres_like_agent_jdbc_connection_string(
+    config: &ConnectionConfig,
+    host: &str,
+    port: u16,
+    database: &str,
+) -> String {
+    let scheme = match config.db_type {
+        DatabaseType::Kingbase => "kingbase8",
+        DatabaseType::Highgo => "highgo",
+        DatabaseType::Vastbase => "vastbase",
+        _ => unreachable!("postgres-like agent JDBC URL requested for {:?}", config.db_type),
+    };
+    let base = format!("jdbc:{scheme}://{host}:{port}/{}", database.trim());
+    append_agent_url_params(base, config.url_params.as_deref())
+}
+
 pub fn should_retry_oracle_with_10g_driver(config: &ConnectionConfig, err: &str) -> bool {
     if config.db_type != DatabaseType::Oracle {
         return false;
@@ -753,6 +771,15 @@ fn sap_hana_jdbc_connection_string(config: &ConnectionConfig, host: &str, port:
     }
 }
 
+fn append_agent_url_params(base: String, params: Option<&str>) -> String {
+    let params = params.unwrap_or("").trim().trim_start_matches(['?', '&']);
+    if params.is_empty() {
+        return base;
+    }
+    let separator = if base.contains('?') { '&' } else { '?' };
+    format!("{base}{separator}{params}")
+}
+
 fn duckdb_paths_match(left: &str, right: &str) -> bool {
     let left = expand_tilde(left);
     let right = expand_tilde(right);
@@ -952,6 +979,49 @@ mod tests {
         assert_eq!(params["connection_string"], "jdbc:oracle:thin:@//oracle.example.com:1521/ORCLPDB1");
     }
 
+    #[test]
+    fn agent_connect_params_build_postgres_like_agent_connection_string_for_selected_database() {
+        let cases = [
+            (
+                DatabaseType::Kingbase,
+                "kingbase.example.com",
+                54321,
+                "jdbc:kingbase8://kingbase.example.com:54321/platform_face_jgj",
+                "jdbc:kingbase8://kingbase.example.com:54321/platform_face_freezer_jgj?sslmode=disable",
+            ),
+            (
+                DatabaseType::Highgo,
+                "highgo.example.com",
+                5866,
+                "jdbc:highgo://highgo.example.com:5866/highgo",
+                "jdbc:highgo://highgo.example.com:5866/platform_face_freezer_jgj?sslmode=disable",
+            ),
+            (
+                DatabaseType::Vastbase,
+                "vastbase.example.com",
+                5432,
+                "jdbc:vastbase://vastbase.example.com:5432/postgres",
+                "jdbc:vastbase://vastbase.example.com:5432/platform_face_freezer_jgj?sslmode=disable",
+            ),
+        ];
+
+        for (db_type, host, port, stale_connection_string, expected_connection_string) in cases {
+            let mut config = mysql_config(Some("platform_face_jgj"));
+            config.db_type = db_type;
+            config.host = host.to_string();
+            config.port = port;
+            config.username = "system".to_string();
+            config.password = "secret".to_string();
+            config.url_params = Some("sslmode=disable".to_string());
+            config.connection_string = Some(stale_connection_string.to_string());
+
+            let params = agent_connect_params(&config, host, port, "platform_face_freezer_jgj");
+
+            assert_eq!(params["database"], "platform_face_freezer_jgj");
+            assert_eq!(params["connection_string"], expected_connection_string);
+        }
+    }
+
     #[test]
     fn agent_connect_params_build_oracle_sid_connection_string() {
         let mut config = mysql_config(Some("ORCL"));