feat(zookeeper): add agent client browsing and node management

crates/dbx-core/src/agent_catalog.rs

@@ -248,6 +248,13 @@ const AGENT_CATALOG: &[AgentCatalogEntry] = &[
         profiles: &[],
     },
     AgentCatalogEntry { db_type: DatabaseType::Etcd, key: "etcd", label: "etcd", store_visible: true, profiles: &[] },
+    AgentCatalogEntry {
+        db_type: DatabaseType::ZooKeeper,
+        key: "zookeeper",
+        label: "Apache ZooKeeper",
+        store_visible: true,
+        profiles: &[],
+    },
     AgentCatalogEntry {
         db_type: DatabaseType::MongoDb,
         key: "mongodb",

crates/dbx-core/src/agent_connection.rs

@@ -35,6 +35,11 @@ pub fn agent_connect_params(config: &ConnectionConfig, host: &str, port: u16, da
     };
     let etcd_endpoints =
         if config.db_type == DatabaseType::Etcd { normalize_etcd_endpoints(config, host, port) } else { String::new() };
+    let zookeeper_connect_string = if config.db_type == DatabaseType::ZooKeeper {
+        normalize_zookeeper_connect_string(config, host, port)
+    } else {
+        String::new()
+    };
     let (agent_host, agent_port) = if is_h2_file_connection(config) { ("", 0) } else { (host, port) };
 
     serde_json::json!({
@@ -51,6 +56,7 @@ pub fn agent_connect_params(config: &ConnectionConfig, host: &str, port: u16, da
         "client_cert_path": config.client_cert_path,
         "client_key_path": config.client_key_path,
         "etcd_endpoints": etcd_endpoints,
+        "zookeeper_connect_string": zookeeper_connect_string,
         "gbase_server": config.gbase_server,
         "informix_server": config.informix_server,
         "jdbc_driver_class": config.jdbc_driver_class.as_deref().unwrap_or(""),
@@ -511,6 +517,16 @@ fn normalize_etcd_endpoints(config: &ConnectionConfig, host: &str, port: u16) ->
     format!("{scheme}://{host}:{port}")
 }
 
+fn normalize_zookeeper_connect_string(config: &ConnectionConfig, host: &str, port: u16) -> String {
+    config
+        .connection_string
+        .as_deref()
+        .map(str::trim)
+        .filter(|value| !value.is_empty())
+        .map(str::to_string)
+        .unwrap_or_else(|| format!("{host}:{port}"))
+}
+
 fn normalize_agent_url_params(params: Option<&str>) -> &str {
     params.unwrap_or("").trim().trim_start_matches(['?', '&'])
 }
@@ -659,6 +675,27 @@ mod tests {
         assert_eq!(params["connection_string"], "");
     }
 
+    #[test]
+    fn zookeeper_agent_params_preserve_configured_connect_string() {
+        let mut cfg = config(DatabaseType::ZooKeeper, None);
+        cfg.connection_string = Some("zk-1:2181,zk-2:2181/app".to_string());
+
+        let params = agent_connect_params(&cfg, "127.0.0.1", 2181, "");
+
+        assert_eq!(params["connection_string"], "zk-1:2181,zk-2:2181/app");
+        assert_eq!(params["zookeeper_connect_string"], "zk-1:2181,zk-2:2181/app");
+    }
+
+    #[test]
+    fn zookeeper_agent_params_fall_back_to_host_port_connect_string() {
+        let cfg = config(DatabaseType::ZooKeeper, None);
+
+        let params = agent_connect_params(&cfg, "zk.local", 2281, "");
+
+        assert_eq!(params["connection_string"], "");
+        assert_eq!(params["zookeeper_connect_string"], "zk.local:2281");
+    }
+
     #[test]
     fn mongo_auth_hint_preserves_original_error() {
         let err = "Agent RPC error: Exception authenticating MongoCredential{mechanism=SCRAM-SHA-1, userName='rwuser', source='admin'}";

crates/dbx-core/src/agent_kv.rs

@@ -25,6 +25,16 @@ pub struct KvKeyMetadata {
     pub version: Option<i64>,
     pub lease: Option<i64>,
     pub value_size: Option<u64>,
+    pub czxid: Option<i64>,
+    pub mzxid: Option<i64>,
+    pub pzxid: Option<i64>,
+    pub ctime: Option<i64>,
+    pub mtime: Option<i64>,
+    pub cversion: Option<i64>,
+    pub aversion: Option<i64>,
+    pub ephemeral_owner: Option<i64>,
+    pub data_length: Option<u64>,
+    pub num_children: Option<i64>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
@@ -74,12 +84,46 @@ pub struct KvPutRequest {
     pub value: KvValue,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub lease: Option<i64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub write_mode: Option<KvWriteMode>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub create_mode: Option<KvCreateMode>,
+}
+
+#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
+#[serde(rename_all = "camelCase")]
+pub struct KvPutOptions {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub lease: Option<i64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub write_mode: Option<KvWriteMode>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub create_mode: Option<KvCreateMode>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+#[serde(rename_all = "snake_case")]
+pub enum KvWriteMode {
+    Upsert,
+    Create,
+    Update,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+#[serde(rename_all = "snake_case")]
+pub enum KvCreateMode {
+    Persistent,
+    Ephemeral,
+    PersistentSequential,
+    EphemeralSequential,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
 #[serde(rename_all = "camelCase")]
 pub struct KvPutResponse {
     pub revision: Option<i64>,
+    pub key: Option<String>,
+    pub created_key: Option<String>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
@@ -109,7 +153,18 @@ pub fn kv_get_params(key: &str) -> serde_json::Value {
 }
 
 pub fn kv_put_params(key: &str, value: KvValue, lease: Option<i64>) -> serde_json::Value {
-    serde_json::to_value(KvPutRequest { key: key.to_string(), value, lease }).expect("KV put request should serialize")
+    kv_put_params_with_options(key, value, KvPutOptions { lease, ..KvPutOptions::default() })
+}
+
+pub fn kv_put_params_with_options(key: &str, value: KvValue, options: KvPutOptions) -> serde_json::Value {
+    serde_json::to_value(KvPutRequest {
+        key: key.to_string(),
+        value,
+        lease: options.lease,
+        write_mode: options.write_mode,
+        create_mode: options.create_mode,
+    })
+    .expect("KV put request should serialize")
 }
 
 pub fn kv_delete_params(key: &str) -> serde_json::Value {
@@ -138,7 +193,17 @@ pub async fn kv_put_core(
     value: KvValue,
     lease: Option<i64>,
 ) -> Result<KvPutResponse, String> {
-    call_agent_kv(state, connection_id, AgentKvMethod::Put, kv_put_params(key, value, lease)).await
+    kv_put_core_with_options(state, connection_id, key, value, KvPutOptions { lease, ..KvPutOptions::default() }).await
+}
+
+pub async fn kv_put_core_with_options(
+    state: &AppState,
+    connection_id: &str,
+    key: &str,
+    value: KvValue,
+    options: KvPutOptions,
+) -> Result<KvPutResponse, String> {
+    call_agent_kv(state, connection_id, AgentKvMethod::Put, kv_put_params_with_options(key, value, options)).await
 }
 
 pub async fn kv_delete_core(state: &AppState, connection_id: &str, key: &str) -> Result<KvDeleteResponse, String> {
@@ -211,6 +276,41 @@ mod tests {
         assert_eq!(kv_delete_params("/app/name"), serde_json::json!({ "key": "/app/name" }));
     }
 
+    #[test]
+    fn serializes_kv_put_options_without_changing_default_shape() {
+        let value = KvValue { encoding: KvValueEncoding::Utf8, data: "dbx".to_string() };
+        assert_eq!(
+            kv_put_params("/app/name", value.clone(), None),
+            serde_json::json!({
+                "key": "/app/name",
+                "value": {
+                    "encoding": "utf8",
+                    "data": "dbx"
+                }
+            })
+        );
+        assert_eq!(
+            kv_put_params_with_options(
+                "/jobs/job-",
+                value,
+                KvPutOptions {
+                    write_mode: Some(KvWriteMode::Create),
+                    create_mode: Some(KvCreateMode::EphemeralSequential),
+                    ..KvPutOptions::default()
+                },
+            ),
+            serde_json::json!({
+                "key": "/jobs/job-",
+                "value": {
+                    "encoding": "utf8",
+                    "data": "dbx"
+                },
+                "writeMode": "create",
+                "createMode": "ephemeral_sequential"
+            })
+        );
+    }
+
     #[test]
     fn decodes_kv_list_prefix_response() {
         let decoded: KvListPrefixResponse = serde_json::from_value(serde_json::json!({
@@ -232,4 +332,49 @@ mod tests {
         assert_eq!(decoded.continuation.as_deref(), Some("next-token"));
         assert_eq!(decoded.revision, Some(9));
     }
+
+    #[test]
+    fn decodes_kv_put_response_with_created_key() {
+        let decoded: KvPutResponse = serde_json::from_value(serde_json::json!({
+            "key": "/jobs/job-0000000001",
+            "createdKey": "/jobs/job-0000000001"
+        }))
+        .unwrap();
+
+        assert_eq!(decoded.revision, None);
+        assert_eq!(decoded.key.as_deref(), Some("/jobs/job-0000000001"));
+        assert_eq!(decoded.created_key.as_deref(), Some("/jobs/job-0000000001"));
+    }
+
+    #[test]
+    fn decodes_zookeeper_metadata_fields() {
+        let decoded: KvGetResponse = serde_json::from_value(serde_json::json!({
+            "found": true,
+            "key": "/admin",
+            "value": {
+                "encoding": "utf8",
+                "data": ""
+            },
+            "metadata": {
+                "czxid": 27,
+                "mzxid": 27,
+                "pzxid": 39825,
+                "ctime": 1780674584000_i64,
+                "mtime": 1780674585000_i64,
+                "cversion": 5,
+                "aversion": 0,
+                "ephemeralOwner": 0,
+                "dataLength": 0,
+                "numChildren": 5
+            }
+        }))
+        .unwrap();
+
+        let metadata = decoded.metadata.unwrap();
+        assert_eq!(metadata.czxid, Some(27));
+        assert_eq!(metadata.mzxid, Some(27));
+        assert_eq!(metadata.pzxid, Some(39825));
+        assert_eq!(metadata.ephemeral_owner, Some(0));
+        assert_eq!(metadata.num_children, Some(5));
+    }
 }

crates/dbx-core/src/connection.rs

@@ -117,6 +117,7 @@ macro_rules! agent_connection_pool_database_type {
             | DatabaseType::Xugu
             | DatabaseType::Iotdb
             | DatabaseType::Etcd
+            | DatabaseType::ZooKeeper
             | DatabaseType::Iris
             | DatabaseType::Access
     };
@@ -1286,6 +1287,12 @@ impl AppState {
                     let timeout = crate::db::connection_timeout();
                     match agent.validate_connection(Some(timeout)).await {
                         Ok(_) => false,
+                        Err(err) if is_agent_validate_connection_unsupported(&err) => {
+                            log::debug!(
+                                "Agent connection pool '{pool_key}' does not support validate_connection; keeping pool"
+                            );
+                            false
+                        }
                         Err(err) => {
                             log::warn!("Agent connection pool '{pool_key}' is stale: {err}");
                             true
@@ -1876,6 +1883,7 @@ async fn ping_keepalive_target(target: &mut KeepaliveTarget, timeout: Duration)
             };
             match client.validate_connection(Some(timeout)).await {
                 Ok(_) => Ok(()),
+                Err(err) if is_agent_validate_connection_unsupported(&err) => Ok(()),
                 Err(err) => {
                     client.kill();
                     Err(err)
@@ -1885,6 +1893,11 @@ async fn ping_keepalive_target(target: &mut KeepaliveTarget, timeout: Duration)
     }
 }
 
+fn is_agent_validate_connection_unsupported(err: &str) -> bool {
+    let lower = err.to_ascii_lowercase();
+    lower.contains("validate_connection") && (lower.contains("unknown method") || lower.contains("method not found"))
+}
+
 fn connection_remote_endpoint(config: &ConnectionConfig) -> (String, u16) {
     if config.db_type == DatabaseType::MongoDb {
         config
@@ -2480,6 +2493,7 @@ mod tests {
         assert!(uses_bare_mysql_pool(&DatabaseType::ManticoreSearch));
         assert!(!uses_bare_mysql_pool(&DatabaseType::Databend));
         assert!(database_capabilities::is_agent_type(&DatabaseType::Databend));
+        assert!(super::uses_agent_connection_pool(&DatabaseType::ZooKeeper));
     }
 
     #[test]
@@ -3304,6 +3318,18 @@ mod tests {
         let _ = std::fs::remove_dir_all(dir);
     }
 
+    #[test]
+    fn agent_validate_connection_unknown_method_is_not_stale() {
+        assert!(super::is_agent_validate_connection_unsupported(
+            "Agent RPC error (-1): Unknown method: validate_connection"
+        ));
+        assert!(super::is_agent_validate_connection_unsupported(
+            "Agent RPC error (-32601): Method not found: validate_connection"
+        ));
+        assert!(!super::is_agent_validate_connection_unsupported("Agent RPC error (-1): Connection timed out"));
+        assert!(!super::is_agent_validate_connection_unsupported("Agent RPC error (-1): Unknown method: kv_put"));
+    }
+
     #[cfg(feature = "duckdb-bundled")]
     #[tokio::test]
     async fn duckdb_client_session_reuses_base_pool_to_avoid_file_locks() {

crates/dbx-core/src/models/connection.rs

@@ -309,6 +309,8 @@ pub enum DatabaseType {
     Xugu,
     Iotdb,
     Etcd,
+    #[serde(rename = "zookeeper")]
+    ZooKeeper,
     Nacos,
     #[serde(rename = "iris")]
     Iris,
@@ -789,6 +791,9 @@ impl ConnectionConfig {
             DatabaseType::Etcd => {
                 format!("etcd://{host}:{port}")
             }
+            DatabaseType::ZooKeeper => {
+                format!("zookeeper://{host}:{port}")
+            }
             DatabaseType::Iris => format!("iris://{host}:{port}{db_part}"),
             DatabaseType::InfluxDb => {
                 let scheme = if self.ssl { "https" } else { "http" };
@@ -991,6 +996,13 @@ impl ConnectionConfig {
                     format!("etcd://{}:{}@{host}:{port}", username, password)
                 }
             }
+            DatabaseType::ZooKeeper => {
+                if self.username.is_empty() {
+                    format!("zookeeper://{host}:{port}")
+                } else {
+                    format!("zookeeper://{}:{}@{host}:{port}", username, password)
+                }
+            }
             DatabaseType::Iris => {
                 format!("iris://{}:{}@{host}:{port}{db_part}", username, password)
             }
@@ -1571,6 +1583,27 @@ mod tests {
         config
     }
 
+    #[test]
+    fn zookeeper_database_type_uses_stable_wire_name() {
+        assert_eq!(serde_json::to_string(&DatabaseType::ZooKeeper).unwrap(), "\"zookeeper\"");
+        assert_eq!(serde_json::from_str::<DatabaseType>("\"zookeeper\"").unwrap(), DatabaseType::ZooKeeper);
+    }
+
+    #[test]
+    fn zookeeper_connection_url_uses_zookeeper_scheme() {
+        let mut config = mysql_config("", "", None);
+        config.db_type = DatabaseType::ZooKeeper;
+        config.host = "zk.local".to_string();
+        config.port = 2181;
+
+        assert_eq!(config.connection_url_with_host("zk.local", 2181), "zookeeper://zk.local:2181");
+
+        config.username = "digest".to_string();
+        config.password = "secret".to_string();
+
+        assert_eq!(config.connection_url_with_host("zk.local", 2181), "zookeeper://digest:secret@zk.local:2181");
+    }
+
     #[test]
     fn legacy_single_ssh_config_migrates_to_transport_layer() {
         let config: ConnectionConfig = serde_json::from_value(serde_json::json!({