github/dbx
4
0 Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-11-15 14:52:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(export): support database setup and batch SQL scripts

Browse Source
This commit is contained in:
二丫讲梵 authored and GitHub committed 2026-07-24 23:52:10 +08:00
1 parent 3d1601c807
commit 9cfee7f497
21 files changed
+622 -235
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/dbx-core/src/database_export.rs
+102 -20
View File
@@ -38,6 +38,11 @@ pub struct DatabaseExportRequest {
pub include_structure: bool,
pub include_data: bool,
pub include_objects: bool,
/// Include a MySQL `CREATE DATABASE IF NOT EXISTS` and `USE` preamble.
/// The option is intentionally opt-in because it needs CREATE privileges
/// when importing into a new server.
#[serde(default)]
pub include_create_database: bool,
#[serde(default)]
pub drop_table_if_exists: bool,
/// Drop the table-level `AUTO_INCREMENT=N` clause from exported MySQL DDL,
@@ -89,6 +94,59 @@ fn database_export_total_objects(request: &DatabaseExportRequest, counts: &Datab
total
}
fn mysql_sql_string_literal(value: &str) -> String {
format!("'{}'", value.replace('\\', "\\\\").replace('\'', "''"))
}
fn mysql_database_export_preamble(database: &str, charset: Option<&str>, collation: Option<&str>) -> String {
let database = quote_identifier(database, &DatabaseType::Mysql);
let charset = charset.map(str::trim).filter(|value| !value.is_empty());
let collation = collation.map(str::trim).filter(|value| !value.is_empty());
let options = match charset {
Some(charset) => match collation {
Some(collation) => format!(" CHARACTER SET {charset} COLLATE {collation}"),
None => format!(" CHARACTER SET {charset}"),
},
None => String::new(),
};
format!("CREATE DATABASE IF NOT EXISTS {database}{options};\nUSE {database};\n")
}
async fn mysql_database_export_preamble_for_request(
state: &crate::connection::AppState,
request: &DatabaseExportRequest,
) -> String {
let metadata_sql = format!(
"SELECT DEFAULT_CHARACTER_SET_NAME, DEFAULT_COLLATION_NAME FROM information_schema.SCHEMATA WHERE SCHEMA_NAME = {}",
mysql_sql_string_literal(&request.database)
);
let metadata = crate::query::execute_sql_statement(
state,
&request.connection_id,
&request.database,
&metadata_sql,
None,
None,
)
.await
.ok()
.and_then(|result| {
result.rows.first().map(|row| {
(
row.first().and_then(Value::as_str).map(str::to_string),
row.get(1).and_then(Value::as_str).map(str::to_string),
)
})
});
match metadata {
Some((charset, collation)) => {
mysql_database_export_preamble(&request.database, charset.as_deref(), collation.as_deref())
}
None => mysql_database_export_preamble(&request.database, None, None),
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct DatabaseBackupSnapshot {
@@ -1274,6 +1332,12 @@ pub async fn export_database_sql_core(
// 4. Create file
let mut file = std::fs::File::create(&request.file_path).map_err(|e| format!("Failed to write file: {e}"))?;
let create_database_preamble = if request.include_create_database && matches!(db_type, DatabaseType::Mysql) {
Some(mysql_database_export_preamble_for_request(state, request).await)
} else {
None
};
// 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}"))?;
@@ -1281,6 +1345,11 @@ pub async fn export_database_sql_core(
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}"))?;
if let Some(preamble) = create_database_preamble {
writeln!(file, "-- Database setup").map_err(|e| format!("Failed to write file: {e}"))?;
writeln!(file, "{preamble}").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}"))?;
@@ -1425,6 +1494,7 @@ pub async fn export_database_sql_core(
);
let mut object_index: usize = 0;
let mut total_rows_exported = 0_u64;
// Export tables
let batch_size = if request.batch_size == 0 { 1000 } else { request.batch_size };
@@ -1438,7 +1508,7 @@ pub async fn export_database_sql_core(
current_object: extension.name.clone(),
object_index,
total_objects,
rows_exported: 0,
rows_exported: total_rows_exported,
total_rows: None,
status: ExportStatus::Running,
error: None,
@@ -1458,7 +1528,7 @@ pub async fn export_database_sql_core(
current_object: sequence.name.clone(),
object_index,
total_objects,
rows_exported: 0,
rows_exported: total_rows_exported,
total_rows: None,
status: ExportStatus::Running,
error: None,
@@ -1559,7 +1629,7 @@ pub async fn export_database_sql_core(
current_object: table_info.name.clone(),
object_index,
total_objects,
rows_exported: 0,
rows_exported: total_rows_exported,
total_rows: None,
status: ExportStatus::Cancelled,
error: None,
@@ -1575,7 +1645,7 @@ pub async fn export_database_sql_core(
current_object: table_name.clone(),
object_index,
total_objects,
rows_exported: 0,
rows_exported: total_rows_exported,
total_rows: None,
status: ExportStatus::Running,
error: None,
@@ -1596,7 +1666,7 @@ pub async fn export_database_sql_core(
current_object: sequence.name.clone(),
object_index,
total_objects,
rows_exported: 0,
rows_exported: total_rows_exported,
total_rows: None,
status: ExportStatus::Running,
error: None,
@@ -1682,7 +1752,6 @@ pub async fn export_database_sql_core(
if !col_names.is_empty() {
if let Some(snapshot_session_id) = request.snapshot_session_id.as_deref() {
let sql = database_export_select_sql(&col_names, table_name, &request.schema, &db_type);
let mut rows_exported = 0_u64;
crate::query::stream_rows_in_manual_transaction(
state,
snapshot_session_id,
@@ -1702,13 +1771,13 @@ pub async fn export_database_sql_core(
&request.schema,
&db_type,
)?;
rows_exported += rows.len() as u64;
total_rows_exported += rows.len() as u64;
on_progress(ExportProgress {
export_id: request.export_id.clone(),
current_object: table_name.clone(),
object_index,
total_objects,
rows_exported,
rows_exported: total_rows_exported,
total_rows: None,
status: ExportStatus::Running,
error: None,
@@ -1737,7 +1806,6 @@ pub async fn export_database_sql_core(
Err(_) => None,
};
let mut offset = 0_u64;
let mut rows_exported = 0_u64;
loop {
if is_export_cancelled(&request.export_id).await {
@@ -1746,7 +1814,7 @@ pub async fn export_database_sql_core(
current_object: table_name.clone(),
object_index,
total_objects,
rows_exported,
rows_exported: total_rows_exported,
total_rows,
status: ExportStatus::Cancelled,
error: None,
@@ -1787,14 +1855,14 @@ pub async fn export_database_sql_core(
&request.schema,
&db_type,
)?;
rows_exported += row_count as u64;
total_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,
rows_exported: total_rows_exported,
total_rows,
status: ExportStatus::Running,
error: None,
@@ -1837,7 +1905,7 @@ pub async fn export_database_sql_core(
current_object: view_name.clone(),
object_index,
total_objects,
rows_exported: 0,
rows_exported: total_rows_exported,
total_rows: None,
status: ExportStatus::Running,
error: None,
@@ -1888,7 +1956,7 @@ pub async fn export_database_sql_core(
current_object: proc_name.clone(),
object_index,
total_objects,
rows_exported: 0,
rows_exported: total_rows_exported,
total_rows: None,
status: ExportStatus::Running,
error: None,
@@ -1943,7 +2011,7 @@ pub async fn export_database_sql_core(
current_object: func_name.clone(),
object_index,
total_objects,
rows_exported: 0,
rows_exported: total_rows_exported,
total_rows: None,
status: ExportStatus::Running,
error: None,
@@ -1997,7 +2065,7 @@ pub async fn export_database_sql_core(
current_object: String::new(),
object_index,
total_objects,
rows_exported: 0,
rows_exported: total_rows_exported,
total_rows: None,
status: ExportStatus::Done,
error: None,
@@ -2054,10 +2122,11 @@ mod tests {
database_export_total_objects, drop_table_if_exists_sql, filter_export_table_infos, format_export_sql_literal,
format_export_table_ddl, generate_postgres_extension_ddl, generate_postgres_sequence_create_ddl,
generate_postgres_sequence_owner_ddl, generate_postgres_sequence_setval_sql,
is_postgres_extension_member_routine, normalize_export_table_ddl, record_export_error,
BuildDatabaseSqlExportOptions, BuildExportInsertStatementsOptions, DatabaseExportObjectCounts,
DatabaseExportRequest, DdlNormalizeOptions, ExportedTableSql, PostgresExportExtension, PostgresExportSequence,
PostgresExtensionMembers, DATABASE_EXPORT_INSERT_BATCH_SIZE, DATABASE_EXPORT_ROW_LIMIT,
is_postgres_extension_member_routine, mysql_database_export_preamble, normalize_export_table_ddl,
record_export_error, BuildDatabaseSqlExportOptions, BuildExportInsertStatementsOptions,
DatabaseExportObjectCounts, DatabaseExportRequest, DdlNormalizeOptions, ExportedTableSql,
PostgresExportExtension, PostgresExportSequence, PostgresExtensionMembers, DATABASE_EXPORT_INSERT_BATCH_SIZE,
DATABASE_EXPORT_ROW_LIMIT,
};
use crate::models::connection::DatabaseType;
use crate::types::{ObjectInfo, ObjectSourceKind, TableInfo};
@@ -2105,6 +2174,7 @@ mod tests {
include_structure,
include_data,
include_objects,
include_create_database: false,
drop_table_if_exists: false,
omit_auto_increment: false,
fail_on_error: false,
@@ -2137,6 +2207,18 @@ mod tests {
}
}
#[test]
fn mysql_database_export_preamble_preserves_database_settings() {
assert_eq!(
mysql_database_export_preamble("app`db", Some("utf8mb4"), Some("utf8mb4_unicode_ci")),
"CREATE DATABASE IF NOT EXISTS `app``db` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;\nUSE `app``db`;\n"
);
assert_eq!(
mysql_database_export_preamble("app", None, None),
"CREATE DATABASE IF NOT EXISTS `app`;\nUSE `app`;\n"
);
}
#[test]
fn export_progress_total_excludes_schema_routines_for_selected_tables() {
// Counts are already filtered to the selected table/view set before