diff --git a/crates/dbx-core/src/agent_manager.rs b/crates/dbx-core/src/agent_manager.rs index e5729d613e..0ce4731179 100644 --- a/crates/dbx-core/src/agent_manager.rs +++ b/crates/dbx-core/src/agent_manager.rs @@ -5,7 +5,7 @@ use serde::{Deserialize, Serialize}; use tokio::sync::Mutex; use crate::database_capabilities; -use crate::db::agent_driver::AgentDriverClient; +use crate::db::agent_driver::{AgentDriverClient, AgentMethod}; use crate::models::connection::DatabaseType; pub const DEFAULT_JRE_KEY: &str = "17"; @@ -389,6 +389,16 @@ impl AgentManager { } } + pub async fn call_daemon_method( + &self, + db_type: &DatabaseType, + driver_profile: Option<&str>, + method: AgentMethod, + params: serde_json::Value, + ) -> Result { + self.call_daemon(db_type, driver_profile, method.as_str(), params).await + } + pub async fn download_file(url: &str, dest: &Path) -> Result<(), String> { if let Some(parent) = dest.parent() { std::fs::create_dir_all(parent).map_err(|e| e.to_string())?; diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 7c2394aa85..36c52b9aba 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -5,6 +5,7 @@ use tokio::sync::RwLock; use crate::database_capabilities; use crate::db; +use crate::db::agent_driver::AgentMethod; use crate::db::proxy_tunnel::ProxyTunnelManager; use crate::db::ssh_tunnel::TunnelManager; use crate::external; @@ -203,10 +204,7 @@ impl AppState { log::info!("Native MongoDB driver failed ({native_err}), falling back to agent driver"); let connect_params = serde_json::json!({ "connection": agent_connect_params(&db_config, &host, port, db_config.effective_database().unwrap_or("")) }); let mut client = self.agent_manager.spawn(&DatabaseType::MongoDb, None).await?; - client - .call::("connect", connect_params) - .await - .map_err(|err| mongo_legacy_error_with_auth_hint(&err))?; + client.connect(connect_params).await.map_err(|err| mongo_legacy_error_with_auth_hint(&err))?; PoolKind::Agent(Arc::new(tokio::sync::Mutex::new(client))) } else { return Err(native_err); @@ -274,8 +272,8 @@ impl AppState { let mut client = self.agent_manager.spawn(&db_config.db_type, db_config.driver_profile.as_deref()).await?; client - .call::( - "connect", + .call_method::( + AgentMethod::Connect, agent_connect_params(&db_config, &host, port, db_config.effective_database().unwrap_or("")), ) .await?; diff --git a/crates/dbx-core/src/db/agent_driver.rs b/crates/dbx-core/src/db/agent_driver.rs index a05671ff5a..16bde6ec10 100644 --- a/crates/dbx-core/src/db/agent_driver.rs +++ b/crates/dbx-core/src/db/agent_driver.rs @@ -72,6 +72,37 @@ impl AgentCapability { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AgentMethod { + Handshake, + Connect, + TestConnection, + ListDatabases, + ListSchemas, + ExecuteQuery, + ExecuteQueryPage, + FetchQueryPage, + Disconnect, + Shutdown, +} + +impl AgentMethod { + pub fn as_str(self) -> &'static str { + match self { + Self::Handshake => "handshake", + Self::Connect => "connect", + Self::TestConnection => "test_connection", + Self::ListDatabases => "list_databases", + Self::ListSchemas => "list_schemas", + Self::ExecuteQuery => "execute_query", + Self::ExecuteQueryPage => "execute_query_page", + Self::FetchQueryPage => "fetch_query_page", + Self::Disconnect => "disconnect", + Self::Shutdown => "shutdown", + } + } +} + struct StderrTail { lines: VecDeque, capacity: usize, @@ -246,8 +277,28 @@ impl AgentDriverClient { result.map_err(|e| self.format_agent_process_error(&e)) } + pub async fn call_method( + &mut self, + method: AgentMethod, + params: Value, + ) -> Result { + self.call(method.as_str(), params).await + } + + pub async fn connect(&mut self, params: Value) -> Result { + self.call_method(AgentMethod::Connect, params).await + } + + pub async fn test_connection(&mut self, params: Value) -> Result { + self.call_method(AgentMethod::TestConnection, params).await + } + + pub async fn disconnect(&mut self) -> Result { + self.call_method(AgentMethod::Disconnect, serde_json::json!({})).await + } + pub async fn try_optional_handshake(&mut self, app_version: &str) -> Option { - match self.call::("handshake", agent_handshake_params(app_version)).await { + match self.call_method::(AgentMethod::Handshake, agent_handshake_params(app_version)).await { Ok(handshake) => { log::info!( "[agent] handshake complete: protocol={}, agent_protocol={}, capabilities={:?}", @@ -271,7 +322,7 @@ impl AgentDriverClient { /// Send a shutdown message to the agent and wait for the process to exit. pub async fn shutdown(&mut self) { // Try to send a shutdown RPC; ignore errors if the agent is already gone - let shutdown_result: Result = self.call("shutdown", Value::Null).await; + let shutdown_result: Result = self.call_method(AgentMethod::Shutdown, Value::Null).await; if let Err(e) = &shutdown_result { log::warn!("Agent shutdown RPC failed: {e}"); } @@ -419,7 +470,7 @@ impl Drop for AgentDriverClient { mod tests { use super::{ agent_handshake_params, agent_java_args, agent_proxy_env_vars, format_agent_process_error, - is_unsupported_handshake_error, read_agent_line, AgentCapability, AgentHandshake, StderrTail, + is_unsupported_handshake_error, read_agent_line, AgentCapability, AgentHandshake, AgentMethod, StderrTail, AGENT_PROTOCOL_VERSION, }; use std::io::Cursor; @@ -527,6 +578,20 @@ mod tests { assert_eq!(AgentCapability::ALL.len(), 7); } + #[test] + fn defines_agent_protocol_methods() { + assert_eq!(AgentMethod::Handshake.as_str(), "handshake"); + assert_eq!(AgentMethod::Connect.as_str(), "connect"); + assert_eq!(AgentMethod::TestConnection.as_str(), "test_connection"); + assert_eq!(AgentMethod::ListDatabases.as_str(), "list_databases"); + assert_eq!(AgentMethod::ListSchemas.as_str(), "list_schemas"); + assert_eq!(AgentMethod::ExecuteQuery.as_str(), "execute_query"); + assert_eq!(AgentMethod::ExecuteQueryPage.as_str(), "execute_query_page"); + assert_eq!(AgentMethod::FetchQueryPage.as_str(), "fetch_query_page"); + assert_eq!(AgentMethod::Disconnect.as_str(), "disconnect"); + assert_eq!(AgentMethod::Shutdown.as_str(), "shutdown"); + } + #[test] fn checks_handshake_capability_support() { let handshake = AgentHandshake { diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index d90c8e3a38..5f4d7cb14a 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -8,6 +8,7 @@ pub use dbx_core::connection::{ }; use dbx_core::database_capabilities; use dbx_core::db; +use dbx_core::db::agent_driver::AgentMethod; use dbx_core::models::connection::{rewrite_jdbc_url_host, ConnectionConfig, DatabaseType}; fn mongo_legacy_connect_params(config: &ConnectionConfig, host: &str, port: u16) -> serde_json::Value { @@ -173,10 +174,10 @@ pub async fn test_connection(state: State<'_, Arc>, config: Connection let am = &state.agent_manager; let mut client = am.spawn(&config.db_type, config.driver_profile.as_deref()).await?; client - .call::("connect", mongo_legacy_connect_params(&config, &host, port)) + .connect(mongo_legacy_connect_params(&config, &host, port)) .await .map_err(|err| mongo_legacy_error_with_auth_hint(&err))?; - client.call::("disconnect", serde_json::json!({})).await.ok(); + client.disconnect().await.ok(); Ok("Connection successful (via legacy driver)".to_string()) } else { Err(native_err) @@ -205,10 +206,10 @@ pub async fn test_connection(state: State<'_, Arc>, config: Connection db_type if database_capabilities::is_agent_type(&db_type) => { state .agent_manager - .call_daemon::( + .call_daemon_method::( &config.db_type, config.driver_profile.as_deref(), - "test_connection", + AgentMethod::TestConnection, agent_connect_params(&config, &host, port, config.database.as_deref().unwrap_or("")), ) .await?; @@ -290,9 +291,7 @@ pub async fn connect_db(state: State<'_, Arc>, config: ConnectionConfi log::info!("Native MongoDB driver failed ({native_err}), falling back to agent driver"); let mut client = state.agent_manager.spawn(&db_config.db_type, db_config.driver_profile.as_deref()).await?; - client - .call::("connect", mongo_legacy_connect_params(&db_config, &host, port)) - .await?; + client.connect(mongo_legacy_connect_params(&db_config, &host, port)).await?; PoolKind::Agent(std::sync::Arc::new(tokio::sync::Mutex::new(client))) } else { return Err(native_err); @@ -330,8 +329,8 @@ pub async fn connect_db(state: State<'_, Arc>, config: ConnectionConfi db_type if database_capabilities::is_agent_type(&db_type) => { let mut client = state.agent_manager.spawn(&db_config.db_type, db_config.driver_profile.as_deref()).await?; client - .call::( - "connect", + .call_method::( + AgentMethod::Connect, agent_connect_params(&db_config, &host, port, db_config.effective_database().unwrap_or("")), ) .await?;