feat(desktop): add database info and SQL statement markers

This commit is contained in:
shen authored and GitHub committed 2026-07-15 18:35:34 +08:00
1 parent 83f7604a11
commit 3748d83cde
60 files changed
+2403 -134

No files matched your search

+1
View File
@@ -624,6 +624,7 @@ mod tests {
read_only: false,
is_production: false,
production_databases: vec![],
database_info: None,
}
}
+3
View File
@@ -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());
+99 -7
View File
@@ -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,
}
}
+5
View File
@@ -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)]
+50 -1
View File
@@ -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()))
+1
View File
@@ -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);
+129 -2
View File
@@ -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]
+1
View File
@@ -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
+1
View File
@@ -599,6 +599,7 @@ mod tests {
read_only,
is_production: false,
production_databases: Vec::new(),
database_info: None,
}
}
+1
View File
@@ -172,6 +172,7 @@ mod tests {
read_only: false,
is_production: false,
production_databases: vec![],
database_info: None,
}
}
+2
View File
@@ -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(
+1
View File
@@ -612,6 +612,7 @@ mod tests {
read_only: false,
is_production: false,
production_databases: vec!["prod_app".to_string()],
database_info: None,
}
}
+40 -16
View File
@@ -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(
+1
View File
@@ -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,
}
}
+62 -2
View File
@@ -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");
+1
View File
@@ -4789,6 +4789,7 @@ mod tests {
read_only: false,
is_production: false,
production_databases: vec![],
database_info: None,
}
}