Files
dbx/src-tauri/src/db/mysql.rs
T

230 lines
8.2 KiB
Rust

use sqlx::mysql::{MySqlPool, MySqlPoolOptions, MySqlRow};
use sqlx::{Column, Executor, Row};
use std::time::{Duration, Instant};
use super::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryResult, TableInfo, TriggerInfo};
fn get_str(row: &MySqlRow, idx: usize) -> String {
row.try_get::<String, _>(idx)
.or_else(|_| row.try_get::<Vec<u8>, _>(idx).map(|b| String::from_utf8_lossy(&b).to_string()))
.unwrap_or_default()
}
fn get_str_by_name(row: &MySqlRow, name: &str) -> String {
row.try_get::<String, _>(name)
.or_else(|_| row.try_get::<Vec<u8>, _>(name).map(|b| String::from_utf8_lossy(&b).to_string()))
.unwrap_or_default()
}
fn get_opt_str(row: &MySqlRow, name: &str) -> Option<String> {
row.try_get::<Option<String>, _>(name).ok().flatten()
.or_else(|| row.try_get::<Option<Vec<u8>>, _>(name).ok().flatten().map(|b| String::from_utf8_lossy(&b).to_string()))
}
pub async fn connect(url: &str) -> Result<MySqlPool, String> {
MySqlPoolOptions::new()
.max_connections(5)
.acquire_timeout(Duration::from_secs(10))
.idle_timeout(Duration::from_secs(300))
.connect(url)
.await
.map_err(|e| format!("MySQL connection failed: {e}"))
}
pub async fn list_databases(pool: &MySqlPool) -> Result<Vec<DatabaseInfo>, String> {
let rows: Vec<MySqlRow> = sqlx::query("SELECT SCHEMA_NAME FROM information_schema.SCHEMATA ORDER BY SCHEMA_NAME")
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
Ok(rows.iter().map(|row| DatabaseInfo { name: get_str(row, 0) }).collect())
}
pub async fn list_tables(pool: &MySqlPool, database: &str) -> Result<Vec<TableInfo>, String> {
let rows: Vec<MySqlRow> = sqlx::query(
"SELECT TABLE_NAME, TABLE_TYPE FROM information_schema.TABLES WHERE TABLE_SCHEMA = ? ORDER BY TABLE_NAME",
)
.bind(database)
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
Ok(rows
.iter()
.map(|row| TableInfo {
name: get_str_by_name(row, "TABLE_NAME"),
table_type: get_str_by_name(row, "TABLE_TYPE"),
})
.collect())
}
pub async fn get_columns(
pool: &MySqlPool,
database: &str,
table: &str,
) -> Result<Vec<ColumnInfo>, String> {
let rows: Vec<MySqlRow> = sqlx::query(
"SELECT c.COLUMN_NAME, c.DATA_TYPE, c.IS_NULLABLE, c.COLUMN_DEFAULT, c.EXTRA, \
CASE WHEN kcu.COLUMN_NAME IS NOT NULL THEN 1 ELSE 0 END AS IS_PK \
FROM information_schema.COLUMNS c \
LEFT JOIN information_schema.KEY_COLUMN_USAGE kcu \
ON c.TABLE_SCHEMA = kcu.TABLE_SCHEMA \
AND c.TABLE_NAME = kcu.TABLE_NAME \
AND c.COLUMN_NAME = kcu.COLUMN_NAME \
AND kcu.CONSTRAINT_NAME = 'PRIMARY' \
WHERE c.TABLE_SCHEMA = ? AND c.TABLE_NAME = ? \
ORDER BY c.ORDINAL_POSITION",
)
.bind(database)
.bind(table)
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
Ok(rows
.iter()
.map(|row| ColumnInfo {
name: get_str_by_name(row, "COLUMN_NAME"),
data_type: get_str_by_name(row, "DATA_TYPE"),
is_nullable: get_str_by_name(row, "IS_NULLABLE") == "YES",
column_default: get_opt_str(row, "COLUMN_DEFAULT"),
is_primary_key: row.get::<i32, _>("IS_PK") == 1,
extra: get_opt_str(row, "EXTRA"),
})
.collect())
}
pub async fn execute_query(pool: &MySqlPool, sql: &str) -> Result<QueryResult, String> {
let start = Instant::now();
let trimmed = sql.trim().to_uppercase();
if trimmed.starts_with("SELECT") || trimmed.starts_with("SHOW") || trimmed.starts_with("DESCRIBE") || trimmed.starts_with("EXPLAIN") {
let desc = pool.describe(sql).await.map_err(|e| e.to_string())?;
let columns: Vec<String> = desc.columns().iter().map(|c| c.name().to_string()).collect();
let rows: Vec<MySqlRow> = sqlx::query(sql)
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
let result_rows: Vec<Vec<serde_json::Value>> = rows
.iter()
.map(|row| {
(0..row.len())
.map(|i| {
row.try_get::<String, _>(i)
.map(serde_json::Value::String)
.or_else(|_| row.try_get::<i64, _>(i).map(|v| serde_json::Value::Number(v.into())))
.or_else(|_| row.try_get::<f64, _>(i).map(|v| {
serde_json::Number::from_f64(v)
.map(serde_json::Value::Number)
.unwrap_or(serde_json::Value::Null)
}))
.or_else(|_| row.try_get::<bool, _>(i).map(serde_json::Value::Bool))
.unwrap_or(serde_json::Value::Null)
})
.collect()
})
.collect();
Ok(QueryResult {
columns,
rows: result_rows,
affected_rows: 0,
execution_time_ms: start.elapsed().as_millis(),
truncated: false,
})
} else {
let result = sqlx::query(sql)
.execute(pool)
.await
.map_err(|e| e.to_string())?;
Ok(QueryResult {
columns: vec![],
rows: vec![],
affected_rows: result.rows_affected(),
execution_time_ms: start.elapsed().as_millis(),
truncated: false,
})
}
}
pub async fn list_indexes(pool: &MySqlPool, database: &str, table: &str) -> Result<Vec<IndexInfo>, String> {
let rows: Vec<MySqlRow> = sqlx::query(
"SELECT INDEX_NAME, GROUP_CONCAT(COLUMN_NAME ORDER BY SEQ_IN_INDEX) AS columns, \
NOT NON_UNIQUE AS is_unique, INDEX_NAME = 'PRIMARY' AS is_primary \
FROM information_schema.STATISTICS \
WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ? \
GROUP BY INDEX_NAME, NON_UNIQUE \
ORDER BY INDEX_NAME",
)
.bind(database)
.bind(table)
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
Ok(rows
.iter()
.map(|row| {
let cols_str = get_str_by_name(row, "columns");
IndexInfo {
name: get_str_by_name(row, "INDEX_NAME"),
columns: cols_str.split(',').map(|s| s.to_string()).collect(),
is_unique: row.get::<bool, _>("is_unique"),
is_primary: row.get::<bool, _>("is_primary"),
}
})
.collect())
}
pub async fn list_foreign_keys(pool: &MySqlPool, database: &str, table: &str) -> Result<Vec<ForeignKeyInfo>, String> {
let rows: Vec<MySqlRow> = sqlx::query(
"SELECT kcu.CONSTRAINT_NAME, kcu.COLUMN_NAME, \
kcu.REFERENCED_TABLE_NAME, kcu.REFERENCED_COLUMN_NAME \
FROM information_schema.KEY_COLUMN_USAGE kcu \
WHERE kcu.TABLE_SCHEMA = ? AND kcu.TABLE_NAME = ? \
AND kcu.REFERENCED_TABLE_NAME IS NOT NULL \
ORDER BY kcu.CONSTRAINT_NAME",
)
.bind(database)
.bind(table)
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
Ok(rows
.iter()
.map(|row| ForeignKeyInfo {
name: get_str_by_name(row, "CONSTRAINT_NAME"),
column: get_str_by_name(row, "COLUMN_NAME"),
ref_table: get_str_by_name(row, "REFERENCED_TABLE_NAME"),
ref_column: get_str_by_name(row, "REFERENCED_COLUMN_NAME"),
})
.collect())
}
pub async fn list_triggers(pool: &MySqlPool, database: &str, table: &str) -> Result<Vec<TriggerInfo>, String> {
let rows: Vec<MySqlRow> = sqlx::query(
"SELECT TRIGGER_NAME, EVENT_MANIPULATION, ACTION_TIMING \
FROM information_schema.TRIGGERS \
WHERE TRIGGER_SCHEMA = ? AND EVENT_OBJECT_TABLE = ? \
ORDER BY TRIGGER_NAME",
)
.bind(database)
.bind(table)
.fetch_all(pool)
.await
.map_err(|e| e.to_string())?;
Ok(rows
.iter()
.map(|row| TriggerInfo {
name: get_str_by_name(row, "TRIGGER_NAME"),
event: get_str_by_name(row, "EVENT_MANIPULATION"),
timing: get_str_by_name(row, "ACTION_TIMING"),
})
.collect())
}