feat(desktop): add database info and SQL statement markers
This commit is contained in:
60 files changed
+2403
-134
No files matched your search
@@ -624,6 +624,7 @@ mod tests {
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1099,6 +1099,7 @@ mod tests {
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1160,6 +1161,7 @@ mod tests {
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1269,6 +1271,7 @@ mod tests {
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
};
|
||||
scrub_connection_secrets(&mut config);
|
||||
assert!(config.password.is_empty());
|
||||
|
||||
@@ -22,8 +22,8 @@ use crate::db::http_tunnel::HttpTunnelManager;
|
||||
use crate::db::proxy_tunnel::ProxyTunnelManager;
|
||||
use crate::db::ssh_tunnel::TunnelManager;
|
||||
use crate::models::connection::{
|
||||
parse_jdbc_host_port, parse_mongo_first_host, rewrite_jdbc_url_host, ConnectionConfig, DatabaseType,
|
||||
TransportLayerConfig,
|
||||
database_info_from_protocol_value, parse_jdbc_host_port, parse_mongo_first_host, rewrite_jdbc_url_host,
|
||||
ConnectionConfig, ConnectionTestResult, DatabaseConnectionInfo, DatabaseType, TransportLayerConfig,
|
||||
};
|
||||
use crate::path_utils::expand_tilde;
|
||||
use crate::plugins::{PluginDriverSession, PluginRegistry, PluginRuntimeEnv};
|
||||
@@ -110,6 +110,12 @@ pub enum PoolKind {
|
||||
Nacos,
|
||||
}
|
||||
|
||||
enum ConnectionDatabaseInfoSource {
|
||||
Agent(Arc<tokio::sync::Mutex<db::agent_driver::AgentDriverClient>>),
|
||||
ExternalDriver { config: Arc<ConnectionConfig>, session: Arc<PluginDriverSession> },
|
||||
NativeMysql(db::mysql::MySqlPool),
|
||||
}
|
||||
|
||||
/// Held connection for a manual transaction session
|
||||
pub enum TxnConnection {
|
||||
Postgres(Box<deadpool_postgres::Object>),
|
||||
@@ -654,9 +660,18 @@ impl AppState {
|
||||
}
|
||||
|
||||
pub async fn test_external_driver(&self, driver_id: &str, config: &ConnectionConfig) -> Result<String, String> {
|
||||
self.test_external_driver_with_info(driver_id, config).await.map(|result| result.message)
|
||||
}
|
||||
|
||||
pub async fn test_external_driver_with_info(
|
||||
&self,
|
||||
driver_id: &str,
|
||||
config: &ConnectionConfig,
|
||||
) -> Result<ConnectionTestResult, String> {
|
||||
let params = serde_json::json!({ "connection": config });
|
||||
let env = self.external_driver_runtime_env(driver_id)?;
|
||||
self.plugins
|
||||
let response = self
|
||||
.plugins
|
||||
.invoke_driver_with_env_and_timeout::<serde_json::Value>(
|
||||
driver_id,
|
||||
"testConnection",
|
||||
@@ -665,7 +680,8 @@ impl AppState {
|
||||
Some(external_driver_connect_timeout(config)),
|
||||
)
|
||||
.await?;
|
||||
Ok("Connection successful".to_string())
|
||||
Ok(ConnectionTestResult::success("Connection successful")
|
||||
.with_database_info(database_info_from_protocol_value(&response)))
|
||||
}
|
||||
|
||||
pub async fn external_driver_pool(&self, driver_id: &str, config: &ConnectionConfig) -> Result<PoolKind, String> {
|
||||
@@ -685,6 +701,18 @@ impl AppState {
|
||||
port: u16,
|
||||
connect_timeout: Duration,
|
||||
) -> Result<String, String> {
|
||||
self.test_sqlserver_connection_with_legacy_fallback_with_info(config, host, port, connect_timeout)
|
||||
.await
|
||||
.map(|result| result.message)
|
||||
}
|
||||
|
||||
pub async fn test_sqlserver_connection_with_legacy_fallback_with_info(
|
||||
&self,
|
||||
config: &ConnectionConfig,
|
||||
host: &str,
|
||||
port: u16,
|
||||
connect_timeout: Duration,
|
||||
) -> Result<ConnectionTestResult, String> {
|
||||
match db::sqlserver::connect_with_port_explicit(
|
||||
host,
|
||||
port,
|
||||
@@ -696,7 +724,7 @@ impl AppState {
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(_) => Ok("Connection successful".to_string()),
|
||||
Ok(_) => Ok(ConnectionTestResult::success("Connection successful")),
|
||||
Err(native_error)
|
||||
if db::sqlserver::sqlserver_legacy_compatibility_enabled(config.url_params.as_deref()) =>
|
||||
{
|
||||
@@ -708,7 +736,7 @@ impl AppState {
|
||||
.spawn(&legacy_config.db_type, legacy_config.driver_profile.as_deref())
|
||||
.await
|
||||
.map_err(|err| sqlserver_legacy_agent_error(&native_error, &err))?;
|
||||
client
|
||||
let response = client
|
||||
.call_method_with_timeout::<serde_json::Value>(
|
||||
AgentMethod::TestConnection,
|
||||
connect_params,
|
||||
@@ -717,7 +745,8 @@ impl AppState {
|
||||
.await
|
||||
.map_err(|err| sqlserver_legacy_agent_error(&native_error, &err))?;
|
||||
client.disconnect().await.ok();
|
||||
Ok("Connection successful (via SQL Server legacy compatibility driver)".to_string())
|
||||
Ok(ConnectionTestResult::success("Connection successful (via SQL Server legacy compatibility driver)")
|
||||
.with_database_info(database_info_from_protocol_value(&response)))
|
||||
}
|
||||
Err(err) => Err(err),
|
||||
}
|
||||
@@ -2346,6 +2375,68 @@ impl AppState {
|
||||
Ok(Some(info.identifier_quote))
|
||||
}
|
||||
|
||||
pub async fn connection_database_info(
|
||||
&self,
|
||||
connection_id: &str,
|
||||
database: Option<&str>,
|
||||
) -> Result<Option<DatabaseConnectionInfo>, String> {
|
||||
let config = self
|
||||
.configs
|
||||
.read()
|
||||
.await
|
||||
.get(connection_id)
|
||||
.cloned()
|
||||
.ok_or_else(|| format!("Connection config not found: {connection_id}"))?;
|
||||
let pool_key = self.get_or_create_pool(connection_id, database).await?;
|
||||
let source = {
|
||||
let connections = self.connections.read().await;
|
||||
match connections.get(&pool_key) {
|
||||
Some(PoolKind::Agent(client)) => Some(ConnectionDatabaseInfoSource::Agent(client.clone())),
|
||||
Some(PoolKind::ExternalDriver { config, session, .. }) => {
|
||||
Some(ConnectionDatabaseInfoSource::ExternalDriver {
|
||||
config: config.clone(),
|
||||
session: session.clone(),
|
||||
})
|
||||
}
|
||||
Some(PoolKind::Mysql(pool, _)) => Some(ConnectionDatabaseInfoSource::NativeMysql(pool.clone())),
|
||||
_ => None,
|
||||
}
|
||||
};
|
||||
|
||||
match source {
|
||||
Some(ConnectionDatabaseInfoSource::Agent(client)) => {
|
||||
let mut agent = client.lock().await;
|
||||
Ok(agent.connection_info(Some(db::connection_timeout())).await?.database_info)
|
||||
}
|
||||
Some(ConnectionDatabaseInfoSource::ExternalDriver { config, session }) => {
|
||||
let response = session
|
||||
.invoke_with_timeout::<serde_json::Value>(
|
||||
"connectionInfo",
|
||||
serde_json::json!({ "connection": config.as_ref() }),
|
||||
Some(db::connection_timeout()),
|
||||
)
|
||||
.await?;
|
||||
Ok(database_info_from_protocol_value(&response))
|
||||
}
|
||||
Some(ConnectionDatabaseInfoSource::NativeMysql(pool)) => {
|
||||
db::mysql::database_connection_info(&pool, db::mysql::protocol_product_name(&config)).await.map(Some)
|
||||
}
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn save_connection_database_info(
|
||||
&self,
|
||||
connection_id: &str,
|
||||
database_info: Option<DatabaseConnectionInfo>,
|
||||
) -> Result<(), String> {
|
||||
self.storage.save_connection_database_info(connection_id, database_info.clone()).await?;
|
||||
if let Some(config) = self.configs.write().await.get_mut(connection_id) {
|
||||
config.database_info = database_info;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn reset_connection_transport(&self, connection_id: &str) {
|
||||
let layer_count = {
|
||||
let configs = self.configs.read().await;
|
||||
@@ -3382,6 +3473,7 @@ mod tests {
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -782,6 +782,7 @@ mod tests {
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -11,6 +11,8 @@ use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::models::connection::DatabaseConnectionInfo;
|
||||
|
||||
type PendingAgentResponse = tokio::sync::oneshot::Sender<Result<Value, String>>;
|
||||
|
||||
pub struct AgentRuntimeClient {
|
||||
@@ -481,7 +483,10 @@ impl AgentMethod {
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AgentConnectionInfo {
|
||||
#[serde(default)]
|
||||
pub identifier_quote: String,
|
||||
#[serde(default)]
|
||||
pub database_info: Option<DatabaseConnectionInfo>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
|
||||
@@ -12,7 +12,7 @@ use std::time::Duration;
|
||||
use std::time::Instant;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::models::connection::DatabaseType;
|
||||
use crate::models::connection::{ConnectionConfig, DatabaseConnectionInfo, DatabaseType};
|
||||
use crate::sql::starts_with_executable_sql_keyword;
|
||||
use crate::types::{
|
||||
ColumnInfo, CompletionAssistantCandidate, CompletionAssistantCandidateKind, CompletionAssistantMatchMode,
|
||||
@@ -103,6 +103,55 @@ fn first_nonempty_str_by_name(row: &mysql_async::Row, names: &[&str]) -> String
|
||||
String::new()
|
||||
}
|
||||
|
||||
fn nonblank(value: String) -> Option<String> {
|
||||
let trimmed = value.trim();
|
||||
if trimmed.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(trimmed.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
async fn query_first_nonblank_string(conn: &mut mysql_async::Conn, sql: &str) -> Option<String> {
|
||||
match conn.query_first::<String, _>(sql).await {
|
||||
Ok(Some(value)) => nonblank(value),
|
||||
Ok(None) => None,
|
||||
Err(error) => {
|
||||
log::debug!("Failed to read optional MySQL database information with `{sql}`: {error}");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn database_connection_info(
|
||||
pool: &MySqlPool,
|
||||
product_name: impl Into<String>,
|
||||
) -> Result<DatabaseConnectionInfo, String> {
|
||||
let product_name = nonblank(product_name.into()).unwrap_or_else(|| "MySQL".to_string());
|
||||
let mut conn = get_conn_with_health_check(pool).await?;
|
||||
|
||||
Ok(DatabaseConnectionInfo {
|
||||
product_name: Some(product_name),
|
||||
product_version: query_first_nonblank_string(&mut conn, "SELECT VERSION()").await,
|
||||
current_database: query_first_nonblank_string(&mut conn, "SELECT COALESCE(DATABASE(), '')").await,
|
||||
server_comment: query_first_nonblank_string(&mut conn, "SELECT @@version_comment").await,
|
||||
server_charset: query_first_nonblank_string(&mut conn, "SELECT @@character_set_server").await,
|
||||
server_collation: query_first_nonblank_string(&mut conn, "SELECT @@collation_server").await,
|
||||
..DatabaseConnectionInfo::default()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn protocol_product_name(config: &ConnectionConfig) -> String {
|
||||
config.driver_label.as_deref().map(str::trim).filter(|value| !value.is_empty()).map(str::to_string).unwrap_or_else(
|
||||
|| match config.db_type {
|
||||
DatabaseType::Doris => "Doris".to_string(),
|
||||
DatabaseType::StarRocks => "StarRocks".to_string(),
|
||||
DatabaseType::ManticoreSearch => "Manticore Search".to_string(),
|
||||
_ => "MySQL".to_string(),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
fn get_opt_metadata_string(row: &mysql_async::Row, name: &str) -> Option<String> {
|
||||
get_opt_str(row, name)
|
||||
.or_else(|| row_get::<NaiveDateTime, _>(row, name).map(|value| value.to_string()))
|
||||
|
||||
@@ -3204,6 +3204,7 @@ mod tests {
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
};
|
||||
|
||||
assert_eq!(redis_database_index(&config), 4);
|
||||
|
||||
@@ -3,6 +3,71 @@ use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use std::collections::HashMap;
|
||||
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum IdentifierCase {
|
||||
Lower,
|
||||
Upper,
|
||||
Mixed,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct DatabaseConnectionInfo {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub product_name: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub product_version: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub current_database: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub server_comment: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub server_charset: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub server_collation: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub unquoted_identifier_case: Option<IdentifierCase>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub quoted_identifier_case: Option<IdentifierCase>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub driver_name: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub driver_version: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub jdbc_version: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ConnectionTestResult {
|
||||
pub message: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub database_info: Option<DatabaseConnectionInfo>,
|
||||
}
|
||||
|
||||
impl ConnectionTestResult {
|
||||
pub fn success(message: impl Into<String>) -> Self {
|
||||
Self { message: message.into(), database_info: None }
|
||||
}
|
||||
|
||||
pub fn with_database_info(mut self, database_info: Option<DatabaseConnectionInfo>) -> Self {
|
||||
self.database_info = database_info;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct DatabaseInfoEnvelope {
|
||||
#[serde(default)]
|
||||
database_info: Option<DatabaseConnectionInfo>,
|
||||
}
|
||||
|
||||
pub fn database_info_from_protocol_value(value: &Value) -> Option<DatabaseConnectionInfo> {
|
||||
serde_json::from_value::<DatabaseInfoEnvelope>(value.clone()).ok()?.database_info
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, PartialEq)]
|
||||
pub struct ConnectionConfig {
|
||||
pub id: String,
|
||||
@@ -100,6 +165,9 @@ pub struct ConnectionConfig {
|
||||
/// Database-level production markers for multi-database connections.
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub production_databases: Vec<String>,
|
||||
/// Metadata captured from the latest successful connection test for this saved config.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub database_info: Option<DatabaseConnectionInfo>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
@@ -546,6 +614,8 @@ struct ConnectionConfigData {
|
||||
pub is_production: bool,
|
||||
#[serde(default)]
|
||||
pub production_databases: Vec<String>,
|
||||
#[serde(default)]
|
||||
pub database_info: Option<DatabaseConnectionInfo>,
|
||||
}
|
||||
|
||||
impl From<ConnectionConfigData> for ConnectionConfig {
|
||||
@@ -599,6 +669,7 @@ impl From<ConnectionConfigData> for ConnectionConfig {
|
||||
read_only: data.read_only,
|
||||
is_production: data.is_production,
|
||||
production_databases: data.production_databases,
|
||||
database_info: data.database_info,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1889,8 +1960,9 @@ fn bracket_ipv6(host: &str) -> String {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
default_query_timeout_secs, default_redis_key_separator, default_ssh_connect_timeout_secs, ConnectionConfig,
|
||||
DatabaseType, ProxyTunnelConfig, ProxyType, TransportLayerConfig,
|
||||
database_info_from_protocol_value, default_query_timeout_secs, default_redis_key_separator,
|
||||
default_ssh_connect_timeout_secs, ConnectionConfig, ConnectionTestResult, DatabaseConnectionInfo, DatabaseType,
|
||||
IdentifierCase, ProxyTunnelConfig, ProxyType, TransportLayerConfig,
|
||||
};
|
||||
use std::str::FromStr;
|
||||
|
||||
@@ -1899,6 +1971,41 @@ mod tests {
|
||||
assert_eq!(default_query_timeout_secs(), 60);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn connection_test_result_uses_camel_case_and_omits_missing_details() {
|
||||
let result =
|
||||
ConnectionTestResult::success("Connection successful").with_database_info(Some(DatabaseConnectionInfo {
|
||||
product_name: Some("H2".to_string()),
|
||||
unquoted_identifier_case: Some(IdentifierCase::Upper),
|
||||
..DatabaseConnectionInfo::default()
|
||||
}));
|
||||
|
||||
let value = serde_json::to_value(result).unwrap();
|
||||
assert_eq!(value["message"], "Connection successful");
|
||||
assert_eq!(value["databaseInfo"]["productName"], "H2");
|
||||
assert_eq!(value["databaseInfo"]["unquotedIdentifierCase"], "upper");
|
||||
assert!(value["databaseInfo"].get("driverName").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn protocol_database_info_parser_accepts_details_and_legacy_responses() {
|
||||
let parsed = database_info_from_protocol_value(&serde_json::json!({
|
||||
"ok": true,
|
||||
"databaseInfo": {
|
||||
"productName": "H2",
|
||||
"currentDatabase": "app",
|
||||
"serverCharset": "utf8mb4",
|
||||
"jdbcVersion": "4.2"
|
||||
}
|
||||
}))
|
||||
.unwrap();
|
||||
assert_eq!(parsed.product_name.as_deref(), Some("H2"));
|
||||
assert_eq!(parsed.current_database.as_deref(), Some("app"));
|
||||
assert_eq!(parsed.server_charset.as_deref(), Some("utf8mb4"));
|
||||
assert_eq!(parsed.jdbc_version.as_deref(), Some("4.2"));
|
||||
assert_eq!(database_info_from_protocol_value(&serde_json::json!({ "ok": true })), None);
|
||||
}
|
||||
|
||||
fn mysql_config(username: &str, password: &str, database: Option<&str>) -> ConnectionConfig {
|
||||
ConnectionConfig {
|
||||
id: "id".to_string(),
|
||||
@@ -1949,6 +2056,7 @@ mod tests {
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2020,6 +2128,25 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
assert!(config.agent_java_options.is_empty());
|
||||
assert_eq!(config.database_info, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn connection_config_database_info_survives_json_round_trip() {
|
||||
let mut config = mysql_config("root", "secret", Some("app"));
|
||||
config.database_info = Some(DatabaseConnectionInfo {
|
||||
product_name: Some("MySQL".to_string()),
|
||||
product_version: Some("8.4.0".to_string()),
|
||||
current_database: Some("app".to_string()),
|
||||
server_charset: Some("utf8mb4".to_string()),
|
||||
..DatabaseConnectionInfo::default()
|
||||
});
|
||||
|
||||
let value = serde_json::to_value(&config).unwrap();
|
||||
assert_eq!(value["database_info"]["productVersion"], "8.4.0");
|
||||
|
||||
let restored: ConnectionConfig = serde_json::from_value(value).unwrap();
|
||||
assert_eq!(restored.database_info, config.database_info);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -147,6 +147,7 @@ mod tests {
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
};
|
||||
cfg.redis_key_separator = ":".to_string();
|
||||
cfg
|
||||
|
||||
@@ -599,6 +599,7 @@ mod tests {
|
||||
read_only,
|
||||
is_production: false,
|
||||
production_databases: Vec::new(),
|
||||
database_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -172,6 +172,7 @@ mod tests {
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -216,6 +216,7 @@ mod tests {
|
||||
read_only: true,
|
||||
is_production: false,
|
||||
production_databases: Vec::new(),
|
||||
database_info: None,
|
||||
};
|
||||
cfg.read_only = true;
|
||||
state.configs.write().await.insert(cfg.id.clone(), cfg);
|
||||
@@ -284,6 +285,7 @@ mod tests {
|
||||
read_only: true,
|
||||
is_production: false,
|
||||
production_databases: Vec::new(),
|
||||
database_info: None,
|
||||
};
|
||||
state.configs.write().await.insert(cfg.id.clone(), cfg);
|
||||
let err = nacos_rollback_config_core(
|
||||
|
||||
@@ -612,6 +612,7 @@ mod tests {
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec!["prod_app".to_string()],
|
||||
database_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -46,19 +46,30 @@ pub enum PoolErrorAction {
|
||||
|
||||
/// A multi-statement result with metadata intended for query clients.
|
||||
///
|
||||
/// `execution_error` is emitted only for synthesized MySQL-protocol errors so
|
||||
/// clients can distinguish them from a successful result column named `Error`.
|
||||
/// `execution_error` is emitted for synthesized per-statement errors so clients
|
||||
/// can distinguish them from a successful result column named `Error`.
|
||||
/// `statement_index` is emitted only after a concrete statement starts running.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct ExecuteMultiResult {
|
||||
#[serde(flatten)]
|
||||
pub result: db::QueryResult,
|
||||
#[serde(skip_serializing_if = "is_false")]
|
||||
pub execution_error: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub statement_index: Option<usize>,
|
||||
}
|
||||
|
||||
impl ExecuteMultiResult {
|
||||
fn execution_error(result: db::QueryResult) -> Self {
|
||||
Self { result, execution_error: true }
|
||||
Self { result, execution_error: true, statement_index: None }
|
||||
}
|
||||
|
||||
fn execution_error_with_index(result: db::QueryResult, statement_index: usize) -> Self {
|
||||
Self { result, execution_error: true, statement_index: Some(statement_index) }
|
||||
}
|
||||
|
||||
fn success_with_index(result: db::QueryResult, statement_index: usize) -> Self {
|
||||
Self { result, execution_error: false, statement_index: Some(statement_index) }
|
||||
}
|
||||
|
||||
fn into_query_result(self) -> db::QueryResult {
|
||||
@@ -68,7 +79,7 @@ impl ExecuteMultiResult {
|
||||
|
||||
impl From<db::QueryResult> for ExecuteMultiResult {
|
||||
fn from(result: db::QueryResult) -> Self {
|
||||
Self { result, execution_error: false }
|
||||
Self { result, execution_error: false, statement_index: None }
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1918,9 +1929,9 @@ pub async fn execute_multi_core_with_options_for_client(
|
||||
}
|
||||
|
||||
let mut results = Vec::with_capacity(statements.len());
|
||||
for stmt in &statements {
|
||||
for (statement_index, stmt) in statements.iter().enumerate() {
|
||||
if is_canceled(&cancel_token) {
|
||||
results.push(error_query_result(canceled_error()));
|
||||
results.push(ExecuteMultiResult::execution_error(error_query_result(canceled_error())));
|
||||
break;
|
||||
}
|
||||
match execute_sql_statement_with_options(
|
||||
@@ -1934,10 +1945,10 @@ pub async fn execute_multi_core_with_options_for_client(
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(r) => results.push(r),
|
||||
Ok(r) => results.push(ExecuteMultiResult::success_with_index(r, statement_index)),
|
||||
Err(e) => {
|
||||
let action = query_pool_error_action(db_type, stmt, &e);
|
||||
results.push(error_query_result(e));
|
||||
results.push(ExecuteMultiResult::execution_error_with_index(error_query_result(e), statement_index));
|
||||
if !should_continue_batch_after_error(options.continue_on_error, action) {
|
||||
break;
|
||||
}
|
||||
@@ -1945,7 +1956,7 @@ pub async fn execute_multi_core_with_options_for_client(
|
||||
}
|
||||
}
|
||||
|
||||
Ok(results.into_iter().map(Into::into).collect())
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
trait MysqlBatchStatementExecutor {
|
||||
@@ -1989,17 +2000,17 @@ where
|
||||
E: MysqlBatchStatementExecutor,
|
||||
{
|
||||
let mut results = Vec::with_capacity(statements.len());
|
||||
for statement in statements {
|
||||
for (statement_index, statement) in statements.iter().enumerate() {
|
||||
if is_canceled(&cancel_token) {
|
||||
results.push(ExecuteMultiResult::execution_error(error_query_result(canceled_error())));
|
||||
return (results, None);
|
||||
}
|
||||
|
||||
match executor.execute_statement(statement).await {
|
||||
Ok(result) => results.push(result.into()),
|
||||
Ok(result) => results.push(ExecuteMultiResult::success_with_index(result, statement_index)),
|
||||
Err(err) => {
|
||||
let action = pool_error_action(db_type, &err);
|
||||
results.push(ExecuteMultiResult::execution_error(error_query_result(err)));
|
||||
results.push(ExecuteMultiResult::execution_error_with_index(error_query_result(err), statement_index));
|
||||
// Statement errors are safe to collect, but connection-level failures leave
|
||||
// the protocol state unusable and must still trigger pool cleanup.
|
||||
if !should_continue_batch_after_error(continue_on_error, action) {
|
||||
@@ -3209,6 +3220,7 @@ mod tests {
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3336,6 +3348,8 @@ mod tests {
|
||||
|
||||
assert_eq!(executor.executed, vec!["first", "fails"]);
|
||||
assert_eq!(results.len(), 2);
|
||||
assert_eq!(results[0].statement_index, Some(0));
|
||||
assert_eq!(results[1].statement_index, Some(1));
|
||||
assert!(results[1].execution_error);
|
||||
assert_eq!(error_action, Some(PoolErrorAction::Keep));
|
||||
}
|
||||
@@ -3374,6 +3388,10 @@ mod tests {
|
||||
|
||||
assert_eq!(executor.executed, statements);
|
||||
assert_eq!(results.len(), 3);
|
||||
assert_eq!(
|
||||
results.iter().map(|result| result.statement_index).collect::<Vec<_>>(),
|
||||
vec![Some(0), Some(1), Some(2)]
|
||||
);
|
||||
assert!(results[1].execution_error);
|
||||
assert_eq!(error_action, None);
|
||||
}
|
||||
@@ -3412,19 +3430,24 @@ mod tests {
|
||||
|
||||
assert_eq!(executor.executed, vec!["first", "disconnects"]);
|
||||
assert_eq!(results.len(), 2);
|
||||
assert_eq!(results.iter().map(|result| result.statement_index).collect::<Vec<_>>(), vec![Some(0), Some(1)]);
|
||||
assert!(results[1].execution_error);
|
||||
assert_eq!(error_action, Some(PoolErrorAction::ReconnectAndRetry));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn execute_multi_result_serializes_error_marker_only_for_synthesized_errors() {
|
||||
fn execute_multi_result_serializes_client_metadata_only_when_present() {
|
||||
let success = serde_json::to_value(ExecuteMultiResult::from(empty_query_result(0))).unwrap();
|
||||
assert!(success.get("execution_error").is_none());
|
||||
assert!(success.get("statement_index").is_none());
|
||||
|
||||
let failure =
|
||||
serde_json::to_value(ExecuteMultiResult::execution_error(error_query_result("failed".to_string())))
|
||||
.unwrap();
|
||||
let failure = serde_json::to_value(ExecuteMultiResult::execution_error_with_index(
|
||||
error_query_result("failed".to_string()),
|
||||
2,
|
||||
))
|
||||
.unwrap();
|
||||
assert_eq!(failure.get("execution_error"), Some(&serde_json::Value::Bool(true)));
|
||||
assert_eq!(failure.get("statement_index"), Some(&serde_json::json!(2)));
|
||||
assert_eq!(failure.get("columns"), Some(&serde_json::json!(["Error"])));
|
||||
}
|
||||
|
||||
@@ -4222,6 +4245,7 @@ mod tests {
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
};
|
||||
|
||||
let params = external_driver_query_params(
|
||||
|
||||
@@ -2501,6 +2501,7 @@ mod tests {
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -191,6 +191,7 @@ mod tests {
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -14,7 +14,7 @@ use crate::connection_secrets::{
|
||||
};
|
||||
use crate::db::sqlite::{connect_path_create_if_missing, SqliteHandle};
|
||||
use crate::history::{HistoryEntry, MAX_HISTORY};
|
||||
use crate::models::connection::{ConnectionConfig, DatabaseType, TransportLayerConfig};
|
||||
use crate::models::connection::{ConnectionConfig, DatabaseConnectionInfo, DatabaseType, TransportLayerConfig};
|
||||
use crate::saved_sql::{SavedSqlFile, SavedSqlFolder, SavedSqlLibrary};
|
||||
|
||||
const SSH_TUNNEL_SECRET_PREFIX: &str = "ssh_tunnels.";
|
||||
@@ -1390,6 +1390,30 @@ impl Storage {
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn save_connection_database_info(
|
||||
&self,
|
||||
connection_id: &str,
|
||||
database_info: Option<DatabaseConnectionInfo>,
|
||||
) -> Result<(), String> {
|
||||
let connection_id = connection_id.to_string();
|
||||
self.with_conn(move |conn| {
|
||||
let json = conn
|
||||
.query_row("SELECT config_json FROM connections WHERE id = ?1", [&connection_id], |row| {
|
||||
row.get::<_, String>(0)
|
||||
})
|
||||
.optional()
|
||||
.map_err(|error| error.to_string())?
|
||||
.ok_or_else(|| format!("Connection config not found: {connection_id}"))?;
|
||||
let mut config: ConnectionConfig = serde_json::from_str(&json).map_err(|error| error.to_string())?;
|
||||
config.database_info = database_info;
|
||||
let json = serde_json::to_string(&config).map_err(|error| error.to_string())?;
|
||||
conn.execute("UPDATE connections SET config_json = ?1 WHERE id = ?2", params![json, connection_id])
|
||||
.map(|_| ())
|
||||
.map_err(|error| error.to_string())
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn load_connections(&self) -> Result<Vec<ConnectionConfig>, String> {
|
||||
let rows: Vec<(String, String)> = self
|
||||
.with_conn(|conn| {
|
||||
@@ -2507,7 +2531,9 @@ mod tests {
|
||||
use crate::connection_secrets::{
|
||||
MQ_AUTH_PASSWORD_KEY, MQ_AUTH_TOKEN_KEY, MQ_TOKEN_SIGNING_KEY, NACOS_AUTH_PASSWORD_KEY,
|
||||
};
|
||||
use crate::models::connection::{ConnectionConfig, DatabaseType, SshTunnelConfig, TransportLayerConfig};
|
||||
use crate::models::connection::{
|
||||
ConnectionConfig, DatabaseConnectionInfo, DatabaseType, SshTunnelConfig, TransportLayerConfig,
|
||||
};
|
||||
use crate::saved_sql::SavedSqlFile;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
@@ -2637,6 +2663,7 @@ mod tests {
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2698,6 +2725,7 @@ mod tests {
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2824,6 +2852,38 @@ mod tests {
|
||||
assert!(!target_dir.join("dbx.db").exists());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn save_connections_preserves_database_info() {
|
||||
let path = temp_db_path("database-info");
|
||||
let storage = Storage::open(&path).await.unwrap();
|
||||
let mut config = mq_connection("database-info", "mq-secret");
|
||||
config.database_info = Some(DatabaseConnectionInfo {
|
||||
product_name: Some("MySQL".to_string()),
|
||||
product_version: Some("8.4.0".to_string()),
|
||||
current_database: Some("app".to_string()),
|
||||
..DatabaseConnectionInfo::default()
|
||||
});
|
||||
|
||||
storage.save_connections(std::slice::from_ref(&config)).await.unwrap();
|
||||
|
||||
let raw_json = raw_connection_json(&storage, "database-info").await;
|
||||
assert!(raw_json.contains("8.4.0"));
|
||||
let loaded = storage.load_connections().await.unwrap();
|
||||
assert_eq!(loaded[0].database_info, config.database_info);
|
||||
|
||||
let updated_info = DatabaseConnectionInfo {
|
||||
product_name: Some("MySQL".to_string()),
|
||||
product_version: Some("8.4.1".to_string()),
|
||||
..DatabaseConnectionInfo::default()
|
||||
};
|
||||
storage.save_connection_database_info("database-info", Some(updated_info.clone())).await.unwrap();
|
||||
let loaded = storage.load_connections().await.unwrap();
|
||||
assert_eq!(loaded[0].database_info, Some(updated_info));
|
||||
assert_eq!(mq_token(&loaded[0]), Some("mq-secret"));
|
||||
|
||||
let _ = std::fs::remove_file(path);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn save_connections_moves_mq_auth_token_to_secret_table_and_restores_it() {
|
||||
let path = temp_db_path("mq-token-secrets");
|
||||
|
||||
@@ -4789,6 +4789,7 @@ mod tests {
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -64,6 +64,7 @@ fn live_postgres_config(
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -58,6 +58,7 @@ fn postgres_test_config(id: &str, database: &str) -> ConnectionConfig {
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -61,6 +61,7 @@ fn live_sqlserver_config(id: &str, database: &str) -> dbx_core::models::connecti
|
||||
read_only: false,
|
||||
is_production: false,
|
||||
production_databases: vec![],
|
||||
database_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in new issue
Block a user