xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor(agent): route core calls through protocol methods

This commit is contained in:
t8y2 2026-05-19 15:38:03 +08:00
parent db10d92d13
commit 5ce6a08f05
4 changed files with 91 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
crates/dbx-core/src/agent_manager.rs
View File

@ -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<T: serde::de::DeserializeOwned + Send + 'static>(
&self,
db_type: &DatabaseType,
driver_profile: Option<&str>,
method: AgentMethod,
params: serde_json::Value,
) -> Result<T, String> {
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())?;

10
crates/dbx-core/src/connection.rs
View File

@ -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::<serde_json::Value>("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::<serde_json::Value>(
"connect",
.call_method::<serde_json::Value>(
AgentMethod::Connect,
agent_connect_params(&db_config, &host, port, db_config.effective_database().unwrap_or("")),
)
.await?;

71
crates/dbx-core/src/db/agent_driver.rs
View File

@ -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<String>,
capacity: usize,
@ -246,8 +277,28 @@ impl AgentDriverClient {
result.map_err(|e| self.format_agent_process_error(&e))
}
pub async fn call_method<T: DeserializeOwned + Send + 'static>(
&mut self,
method: AgentMethod,
params: Value,
) -> Result<T, String> {
self.call(method.as_str(), params).await
}
pub async fn connect(&mut self, params: Value) -> Result<Value, String> {
self.call_method(AgentMethod::Connect, params).await
}
pub async fn test_connection(&mut self, params: Value) -> Result<Value, String> {
self.call_method(AgentMethod::TestConnection, params).await
}
pub async fn disconnect(&mut self) -> Result<Value, String> {
self.call_method(AgentMethod::Disconnect, serde_json::json!({})).await
}
pub async fn try_optional_handshake(&mut self, app_version: &str) -> Option<AgentHandshake> {
match self.call::<AgentHandshake>("handshake", agent_handshake_params(app_version)).await {
match self.call_method::<AgentHandshake>(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<Value, String> = self.call("shutdown", Value::Null).await;
let shutdown_result: Result<Value, String> = 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 {

17
src-tauri/src/commands/connection.rs
View File

@ -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<AppState>>, config: Connection
let am = &state.agent_manager;
let mut client = am.spawn(&config.db_type, config.driver_profile.as_deref()).await?;
client
.call::<serde_json::Value>("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::<serde_json::Value>("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<AppState>>, config: Connection
db_type if database_capabilities::is_agent_type(&db_type) => {
state
.agent_manager
.call_daemon::<serde_json::Value>(
.call_daemon_method::<serde_json::Value>(
&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<AppState>>, 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::<serde_json::Value>("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<AppState>>, 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::<serde_json::Value>(
"connect",
.call_method::<serde_json::Value>(
AgentMethod::Connect,
agent_connect_params(&db_config, &host, port, db_config.effective_database().unwrap_or("")),
)
.await?;