478 lines
16 KiB
Rust
478 lines
16 KiB
Rust
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashSet;
|
|
use std::io::Write;
|
|
use tokio::sync::RwLock;
|
|
|
|
use crate::models::connection::DatabaseType;
|
|
|
|
static EXPORT_CANCELLED: std::sync::LazyLock<RwLock<HashSet<String>>> =
|
|
std::sync::LazyLock::new(|| RwLock::new(HashSet::new()));
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub struct DatabaseExportRequest {
|
|
pub export_id: String,
|
|
pub connection_id: String,
|
|
pub database: String,
|
|
pub schema: String,
|
|
pub file_path: String,
|
|
#[serde(default)]
|
|
pub selected_tables: Vec<String>,
|
|
pub include_structure: bool,
|
|
pub include_data: bool,
|
|
pub include_objects: bool,
|
|
pub batch_size: usize,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub struct ExportProgress {
|
|
pub export_id: String,
|
|
pub current_object: String,
|
|
pub object_index: usize,
|
|
pub total_objects: usize,
|
|
pub rows_exported: u64,
|
|
pub total_rows: Option<u64>,
|
|
pub status: ExportStatus,
|
|
pub error: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub enum ExportStatus {
|
|
Running,
|
|
Done,
|
|
Error,
|
|
Cancelled,
|
|
}
|
|
|
|
pub async fn is_export_cancelled(export_id: &str) -> bool {
|
|
EXPORT_CANCELLED.read().await.contains(export_id)
|
|
}
|
|
|
|
pub async fn set_export_cancelled(export_id: &str) {
|
|
EXPORT_CANCELLED.write().await.insert(export_id.to_string());
|
|
}
|
|
|
|
pub async fn clear_export_cancelled(export_id: &str) {
|
|
EXPORT_CANCELLED.write().await.remove(export_id);
|
|
}
|
|
|
|
pub async fn export_database_sql_core(
|
|
state: &crate::connection::AppState,
|
|
request: &DatabaseExportRequest,
|
|
on_progress: impl Fn(ExportProgress),
|
|
) -> Result<(), String> {
|
|
// 1. Get database type
|
|
let db_type = state
|
|
.configs
|
|
.read()
|
|
.await
|
|
.get(&request.connection_id)
|
|
.map(|c| c.db_type.clone())
|
|
.ok_or_else(|| format!("Connection config not found: {}", request.connection_id))?;
|
|
|
|
// 2. Get pool
|
|
let pool_key = state.get_or_create_pool(&request.connection_id, Some(&request.database)).await?;
|
|
|
|
// 3. List tables
|
|
let all_tables =
|
|
crate::schema::list_tables_core(state, &request.connection_id, &request.database, &request.schema, None, None)
|
|
.await?;
|
|
let all_tables = filter_selected_table_infos(all_tables, &request.selected_tables);
|
|
|
|
// 4. Create file
|
|
let mut file = std::fs::File::create(&request.file_path).map_err(|e| format!("Failed to write file: {e}"))?;
|
|
|
|
// 5. Write header
|
|
let timestamp = chrono::Local::now().format("%Y-%m-%d %H:%M:%S");
|
|
writeln!(file, "-- Database export: {}", request.database).map_err(|e| format!("Failed to write file: {e}"))?;
|
|
writeln!(file, "-- Date: {timestamp}").map_err(|e| format!("Failed to write file: {e}"))?;
|
|
writeln!(file, "-- Generated by DBX").map_err(|e| format!("Failed to write file: {e}"))?;
|
|
writeln!(file).map_err(|e| format!("Failed to write file: {e}"))?;
|
|
|
|
// 6. For MySQL: disable foreign key checks
|
|
if matches!(db_type, DatabaseType::Mysql) {
|
|
writeln!(file, "SET FOREIGN_KEY_CHECKS = 0;\n").map_err(|e| format!("Failed to write file: {e}"))?;
|
|
}
|
|
|
|
// 7. Separate tables and views
|
|
let tables: Vec<_> = all_tables.iter().filter(|t| t.table_type != "VIEW").collect();
|
|
let views: Vec<_> = all_tables.iter().filter(|t| t.table_type == "VIEW").collect();
|
|
|
|
// 8. Calculate total objects
|
|
let mut total_objects = tables.len() + views.len();
|
|
|
|
// We'll add procedures/functions count later if include_objects
|
|
let mut procedures: Vec<String> = Vec::new();
|
|
let mut functions: Vec<String> = Vec::new();
|
|
|
|
if request.include_objects && request.selected_tables.is_empty() {
|
|
if let Ok(objects) =
|
|
crate::schema::list_objects_core(state, &request.connection_id, &request.database, &request.schema).await
|
|
{
|
|
for obj in &objects {
|
|
let ot = obj.object_type.to_uppercase();
|
|
if ot.contains("PROCEDURE") {
|
|
procedures.push(obj.name.clone());
|
|
} else if ot.contains("FUNCTION") {
|
|
functions.push(obj.name.clone());
|
|
}
|
|
}
|
|
}
|
|
total_objects += procedures.len() + functions.len();
|
|
}
|
|
|
|
let mut object_index: usize = 0;
|
|
|
|
// Export tables
|
|
let batch_size = if request.batch_size == 0 { 1000 } else { request.batch_size };
|
|
|
|
for table_info in &tables {
|
|
// Check cancellation
|
|
if is_export_cancelled(&request.export_id).await {
|
|
on_progress(ExportProgress {
|
|
export_id: request.export_id.clone(),
|
|
current_object: table_info.name.clone(),
|
|
object_index,
|
|
total_objects,
|
|
rows_exported: 0,
|
|
total_rows: None,
|
|
status: ExportStatus::Cancelled,
|
|
error: None,
|
|
});
|
|
return Err("Export cancelled".to_string());
|
|
}
|
|
|
|
let table_name = &table_info.name;
|
|
|
|
// Emit Running progress
|
|
on_progress(ExportProgress {
|
|
export_id: request.export_id.clone(),
|
|
current_object: table_name.clone(),
|
|
object_index,
|
|
total_objects,
|
|
rows_exported: 0,
|
|
total_rows: None,
|
|
status: ExportStatus::Running,
|
|
error: None,
|
|
});
|
|
|
|
// Export structure
|
|
if request.include_structure {
|
|
match crate::schema::get_table_ddl_core(
|
|
state,
|
|
&request.connection_id,
|
|
&request.database,
|
|
&request.schema,
|
|
table_name,
|
|
)
|
|
.await
|
|
{
|
|
Ok(ddl) => {
|
|
writeln!(file, "{};\n", ddl).map_err(|e| format!("Failed to write file: {e}"))?;
|
|
}
|
|
Err(e) => {
|
|
writeln!(file, "-- ERROR exporting table {table_name}: {e}")
|
|
.map_err(|e| format!("Failed to write file: {e}"))?;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Export data
|
|
if request.include_data {
|
|
// Get columns
|
|
let col_names: Vec<String> = match crate::schema::get_columns_core(
|
|
state,
|
|
&request.connection_id,
|
|
&request.database,
|
|
&request.schema,
|
|
table_name,
|
|
)
|
|
.await
|
|
{
|
|
Ok(cols) => cols.iter().map(|c| c.name.clone()).collect(),
|
|
Err(e) => {
|
|
writeln!(file, "-- ERROR exporting table {table_name}: {e}")
|
|
.map_err(|e| format!("Failed to write file: {e}"))?;
|
|
object_index += 1;
|
|
continue;
|
|
}
|
|
};
|
|
|
|
if !col_names.is_empty() {
|
|
// Get row count
|
|
let count_query = crate::transfer::count_sql(table_name, &request.schema, &db_type);
|
|
let total_rows = match crate::transfer::execute_on_pool(state, &pool_key, &count_query).await {
|
|
Ok(result) => result.rows.first().and_then(|r| r.first()).and_then(|v| match v {
|
|
serde_json::Value::Number(n) => n.as_u64(),
|
|
serde_json::Value::String(s) => s.parse::<u64>().ok(),
|
|
_ => None,
|
|
}),
|
|
Err(_) => None,
|
|
};
|
|
|
|
// Loop batches
|
|
let mut offset: u64 = 0;
|
|
let mut rows_exported: u64 = 0;
|
|
|
|
loop {
|
|
// Check cancellation between batches
|
|
if is_export_cancelled(&request.export_id).await {
|
|
on_progress(ExportProgress {
|
|
export_id: request.export_id.clone(),
|
|
current_object: table_name.clone(),
|
|
object_index,
|
|
total_objects,
|
|
rows_exported,
|
|
total_rows,
|
|
status: ExportStatus::Cancelled,
|
|
error: None,
|
|
});
|
|
return Err("Export cancelled".to_string());
|
|
}
|
|
|
|
let sql = crate::transfer::pagination_sql(
|
|
&col_names,
|
|
table_name,
|
|
&request.schema,
|
|
&db_type,
|
|
offset,
|
|
batch_size,
|
|
);
|
|
|
|
let result = match crate::transfer::execute_on_pool(state, &pool_key, &sql).await {
|
|
Ok(r) => r,
|
|
Err(e) => {
|
|
writeln!(file, "-- ERROR exporting data for table {table_name}: {e}")
|
|
.map_err(|e| format!("Failed to write file: {e}"))?;
|
|
break;
|
|
}
|
|
};
|
|
|
|
let row_count = result.rows.len();
|
|
if row_count == 0 {
|
|
break;
|
|
}
|
|
|
|
let insert_sql = crate::transfer::generate_insert(
|
|
&col_names,
|
|
&result.rows,
|
|
table_name,
|
|
&request.schema,
|
|
&db_type,
|
|
);
|
|
|
|
if !insert_sql.is_empty() {
|
|
writeln!(file, "{};\n", insert_sql).map_err(|e| format!("Failed to write file: {e}"))?;
|
|
}
|
|
|
|
rows_exported += row_count as u64;
|
|
offset += row_count as u64;
|
|
|
|
on_progress(ExportProgress {
|
|
export_id: request.export_id.clone(),
|
|
current_object: table_name.clone(),
|
|
object_index,
|
|
total_objects,
|
|
rows_exported,
|
|
total_rows,
|
|
status: ExportStatus::Running,
|
|
error: None,
|
|
});
|
|
|
|
if row_count < batch_size {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
object_index += 1;
|
|
}
|
|
|
|
// Export views (if include_objects)
|
|
if request.include_objects {
|
|
for view_info in &views {
|
|
if is_export_cancelled(&request.export_id).await {
|
|
return Err("Export cancelled".to_string());
|
|
}
|
|
|
|
let view_name = &view_info.name;
|
|
|
|
on_progress(ExportProgress {
|
|
export_id: request.export_id.clone(),
|
|
current_object: view_name.clone(),
|
|
object_index,
|
|
total_objects,
|
|
rows_exported: 0,
|
|
total_rows: None,
|
|
status: ExportStatus::Running,
|
|
error: None,
|
|
});
|
|
|
|
match crate::schema::get_object_source_core(
|
|
state,
|
|
&request.connection_id,
|
|
&request.database,
|
|
&request.schema,
|
|
view_name,
|
|
crate::db::ObjectSourceKind::View,
|
|
)
|
|
.await
|
|
{
|
|
Ok(obj_source) => {
|
|
if !obj_source.source.is_empty() {
|
|
writeln!(file, "{};\n", obj_source.source).map_err(|e| format!("Failed to write file: {e}"))?;
|
|
}
|
|
}
|
|
Err(e) => {
|
|
writeln!(file, "-- ERROR exporting view {view_name}: {e}")
|
|
.map_err(|e| format!("Failed to write file: {e}"))?;
|
|
}
|
|
}
|
|
|
|
object_index += 1;
|
|
}
|
|
|
|
// Export procedures
|
|
for proc_name in &procedures {
|
|
if is_export_cancelled(&request.export_id).await {
|
|
return Err("Export cancelled".to_string());
|
|
}
|
|
|
|
on_progress(ExportProgress {
|
|
export_id: request.export_id.clone(),
|
|
current_object: proc_name.clone(),
|
|
object_index,
|
|
total_objects,
|
|
rows_exported: 0,
|
|
total_rows: None,
|
|
status: ExportStatus::Running,
|
|
error: None,
|
|
});
|
|
|
|
match crate::schema::get_object_source_core(
|
|
state,
|
|
&request.connection_id,
|
|
&request.database,
|
|
&request.schema,
|
|
proc_name,
|
|
crate::db::ObjectSourceKind::Procedure,
|
|
)
|
|
.await
|
|
{
|
|
Ok(obj_source) => {
|
|
if !obj_source.source.is_empty() {
|
|
writeln!(file, "{};\n", obj_source.source).map_err(|e| format!("Failed to write file: {e}"))?;
|
|
}
|
|
}
|
|
Err(e) => {
|
|
writeln!(file, "-- ERROR exporting procedure {proc_name}: {e}")
|
|
.map_err(|e| format!("Failed to write file: {e}"))?;
|
|
}
|
|
}
|
|
|
|
object_index += 1;
|
|
}
|
|
|
|
// Export functions
|
|
for func_name in &functions {
|
|
if is_export_cancelled(&request.export_id).await {
|
|
return Err("Export cancelled".to_string());
|
|
}
|
|
|
|
on_progress(ExportProgress {
|
|
export_id: request.export_id.clone(),
|
|
current_object: func_name.clone(),
|
|
object_index,
|
|
total_objects,
|
|
rows_exported: 0,
|
|
total_rows: None,
|
|
status: ExportStatus::Running,
|
|
error: None,
|
|
});
|
|
|
|
match crate::schema::get_object_source_core(
|
|
state,
|
|
&request.connection_id,
|
|
&request.database,
|
|
&request.schema,
|
|
func_name,
|
|
crate::db::ObjectSourceKind::Function,
|
|
)
|
|
.await
|
|
{
|
|
Ok(obj_source) => {
|
|
if !obj_source.source.is_empty() {
|
|
writeln!(file, "{};\n", obj_source.source).map_err(|e| format!("Failed to write file: {e}"))?;
|
|
}
|
|
}
|
|
Err(e) => {
|
|
writeln!(file, "-- ERROR exporting function {func_name}: {e}")
|
|
.map_err(|e| format!("Failed to write file: {e}"))?;
|
|
}
|
|
}
|
|
|
|
object_index += 1;
|
|
}
|
|
}
|
|
|
|
// For MySQL: re-enable foreign key checks
|
|
if matches!(db_type, DatabaseType::Mysql) {
|
|
writeln!(file, "SET FOREIGN_KEY_CHECKS = 1;").map_err(|e| format!("Failed to write file: {e}"))?;
|
|
}
|
|
|
|
// Emit Done progress
|
|
on_progress(ExportProgress {
|
|
export_id: request.export_id.clone(),
|
|
current_object: String::new(),
|
|
object_index,
|
|
total_objects,
|
|
rows_exported: 0,
|
|
total_rows: None,
|
|
status: ExportStatus::Done,
|
|
error: None,
|
|
});
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn filter_selected_table_infos(
|
|
tables: Vec<crate::types::TableInfo>,
|
|
selected_tables: &[String],
|
|
) -> Vec<crate::types::TableInfo> {
|
|
if selected_tables.is_empty() {
|
|
return tables;
|
|
}
|
|
let selected: HashSet<&str> = selected_tables.iter().map(String::as_str).collect();
|
|
tables.into_iter().filter(|table| selected.contains(table.name.as_str())).collect()
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::filter_selected_table_infos;
|
|
use crate::types::TableInfo;
|
|
|
|
fn table(name: &str, table_type: &str) -> TableInfo {
|
|
TableInfo { name: name.to_string(), table_type: table_type.to_string(), comment: None }
|
|
}
|
|
|
|
#[test]
|
|
fn filters_export_tables_by_selected_names() {
|
|
let tables = vec![table("users", "TABLE"), table("orders", "TABLE"), table("active_users", "VIEW")];
|
|
|
|
let filtered = filter_selected_table_infos(tables, &["active_users".to_string(), "users".to_string()]);
|
|
|
|
assert_eq!(filtered.iter().map(|table| table.name.as_str()).collect::<Vec<_>>(), vec!["users", "active_users"]);
|
|
}
|
|
|
|
#[test]
|
|
fn keeps_all_export_tables_when_selection_is_empty() {
|
|
let tables = vec![table("users", "TABLE"), table("orders", "TABLE")];
|
|
|
|
let filtered = filter_selected_table_infos(tables.clone(), &[]);
|
|
|
|
assert_eq!(filtered.iter().map(|table| table.name.as_str()).collect::<Vec<_>>(), vec!["users", "orders"]);
|
|
}
|
|
}
|