From df94a7c1bf9c44c50ab9b92c55abfdbb7f33e1fe Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Mon, 22 Jun 2026 16:47:10 +0800 Subject: [PATCH] fix(postgres): export sequences with database dump --- .../export/DatabaseExportDialog.vue | 2 +- apps/desktop/src/i18n/locales/en.ts | 3 +- apps/desktop/src/i18n/locales/es.ts | 3 +- apps/desktop/src/i18n/locales/it.ts | 3 +- apps/desktop/src/i18n/locales/ja.ts | 3 +- apps/desktop/src/i18n/locales/pt-BR.ts | 3 +- apps/desktop/src/i18n/locales/zh-CN.ts | 3 +- apps/desktop/src/i18n/locales/zh-TW.ts | 3 +- crates/dbx-core/src/database_export.rs | 299 +++++++++++++++++- 9 files changed, 310 insertions(+), 12 deletions(-) diff --git a/apps/desktop/src/components/export/DatabaseExportDialog.vue b/apps/desktop/src/components/export/DatabaseExportDialog.vue index fc7f73fb9..e4de55cc9 100644 --- a/apps/desktop/src/components/export/DatabaseExportDialog.vue +++ b/apps/desktop/src/components/export/DatabaseExportDialog.vue @@ -418,7 +418,7 @@ watch(
- {{ t("settings.options") }} + {{ t("databaseExport.options") }}
diff --git a/apps/desktop/src/i18n/locales/en.ts b/apps/desktop/src/i18n/locales/en.ts index 6677b8ffe..2b0f3d4f8 100644 --- a/apps/desktop/src/i18n/locales/en.ts +++ b/apps/desktop/src/i18n/locales/en.ts @@ -2672,10 +2672,11 @@ export default { }, databaseExport: { title: "Export Database", + options: "Options", includeStructure: "Table structure (DDL)", dropTableIfExists: "Add DROP TABLE IF EXISTS before DDL", includeData: "Table data (INSERT)", - includeObjects: "Views / Procedures / Functions", + includeObjects: "Views / Procedures / Functions / Sequences", tableSelection: "Tables", selectedTables: "{selected}/{total} selected", selectAllTables: "Select all", diff --git a/apps/desktop/src/i18n/locales/es.ts b/apps/desktop/src/i18n/locales/es.ts index 658ad3641..13eeb7de2 100644 --- a/apps/desktop/src/i18n/locales/es.ts +++ b/apps/desktop/src/i18n/locales/es.ts @@ -2320,10 +2320,11 @@ export default { }, databaseExport: { title: "Exportar base de datos", + options: "Opciones", includeStructure: "Estructura de tablas (DDL)", dropTableIfExists: "Agregar DROP TABLE IF EXISTS antes del DDL", includeData: "Datos de tablas (INSERT)", - includeObjects: "Vistas / Procedimientos / Funciones", + includeObjects: "Vistas / Procedimientos / Funciones / Secuencias", tableSelection: "Tablas", selectedTables: "{selected}/{total} seleccionadas", selectAllTables: "Seleccionar todo", diff --git a/apps/desktop/src/i18n/locales/it.ts b/apps/desktop/src/i18n/locales/it.ts index 289062932..2be29d0a7 100644 --- a/apps/desktop/src/i18n/locales/it.ts +++ b/apps/desktop/src/i18n/locales/it.ts @@ -2395,10 +2395,11 @@ export default { }, databaseExport: { title: "Esporta Database", + options: "Opzioni", includeStructure: "Struttura tabella (DDL)", dropTableIfExists: "Aggiungi DROP TABLE IF EXISTS prima del DDL", includeData: "Dati tabella (INSERT)", - includeObjects: "Viste / Procedure / Funzioni", + includeObjects: "Viste / Procedure / Funzioni / Sequenze", tableSelection: "Tabelle", selectedTables: "{selected}/{total} selezionate", selectAllTables: "Seleziona tutte", diff --git a/apps/desktop/src/i18n/locales/ja.ts b/apps/desktop/src/i18n/locales/ja.ts index 49e4f2225..255244b65 100644 --- a/apps/desktop/src/i18n/locales/ja.ts +++ b/apps/desktop/src/i18n/locales/ja.ts @@ -2626,10 +2626,11 @@ export default { }, databaseExport: { title: "データベースをエクスポート", + options: "オプション", includeStructure: "テーブル構造 (DDL)", dropTableIfExists: "DDLの前にDROP TABLE IF EXISTSを追加", includeData: "テーブルデータ (INSERT)", - includeObjects: "ビュー / プロシージャ / 関数", + includeObjects: "ビュー / プロシージャ / 関数 / シーケンス", tableSelection: "テーブル", selectedTables: "{selected}/{total}件選択中", selectAllTables: "すべて選択", diff --git a/apps/desktop/src/i18n/locales/pt-BR.ts b/apps/desktop/src/i18n/locales/pt-BR.ts index bced4d0d2..39771e65a 100644 --- a/apps/desktop/src/i18n/locales/pt-BR.ts +++ b/apps/desktop/src/i18n/locales/pt-BR.ts @@ -2406,10 +2406,11 @@ export default { }, databaseExport: { title: "Exportar banco de dados", + options: "Opções", includeStructure: "Estrutura da tabela (DDL)", dropTableIfExists: "Adicionar DROP TABLE IF EXISTS antes do DDL", includeData: "Dados da tabela (INSERT)", - includeObjects: "Views / Procedures / Funções", + includeObjects: "Views / Procedures / Funções / Sequências", tableSelection: "Tabelas", selectedTables: "{selected}/{total} selecionadas", selectAllTables: "Selecionar todas", diff --git a/apps/desktop/src/i18n/locales/zh-CN.ts b/apps/desktop/src/i18n/locales/zh-CN.ts index ae2f50e6e..30afd9185 100644 --- a/apps/desktop/src/i18n/locales/zh-CN.ts +++ b/apps/desktop/src/i18n/locales/zh-CN.ts @@ -2696,10 +2696,11 @@ export default { }, databaseExport: { title: "导出数据库", + options: "选项", includeStructure: "表结构 (DDL)", dropTableIfExists: "导出前添加 DROP TABLE IF EXISTS", includeData: "表数据 (INSERT)", - includeObjects: "视图 / 存储过程 / 函数", + includeObjects: "视图 / 存储过程 / 函数 / 序列", tableSelection: "选择表", selectedTables: "已选择 {selected}/{total}", selectAllTables: "全选", diff --git a/apps/desktop/src/i18n/locales/zh-TW.ts b/apps/desktop/src/i18n/locales/zh-TW.ts index e9f438f98..e73c98803 100644 --- a/apps/desktop/src/i18n/locales/zh-TW.ts +++ b/apps/desktop/src/i18n/locales/zh-TW.ts @@ -2405,10 +2405,11 @@ export default { }, databaseExport: { title: "匯出資料庫", + options: "選項", includeStructure: "資料表結構 (DDL)", dropTableIfExists: "匯出前新增 DROP TABLE IF EXISTS", includeData: "資料表資料 (INSERT)", - includeObjects: "檢視 / 預存程序 / 函式", + includeObjects: "檢視 / 預存程序 / 函式 / 序列", tableSelection: "資料表", selectedTables: "已選擇 {selected}/{total}", selectAllTables: "全選", diff --git a/crates/dbx-core/src/database_export.rs b/crates/dbx-core/src/database_export.rs index cade6416c..f1a495071 100644 --- a/crates/dbx-core/src/database_export.rs +++ b/crates/dbx-core/src/database_export.rs @@ -6,7 +6,7 @@ use tokio::sync::RwLock; use crate::models::connection::DatabaseType; use crate::sql_dialect::{qualified_table_name, quote_table_identifier}; -use crate::transfer::{format_ch_array_sql_literal, format_pg_array_sql_literal}; +use crate::transfer::{format_ch_array_sql_literal, format_pg_array_sql_literal, quote_identifier}; static EXPORT_CANCELLED: std::sync::LazyLock>> = std::sync::LazyLock::new(|| RwLock::new(HashSet::new())); @@ -55,6 +55,21 @@ pub const DATABASE_EXPORT_ROW_LIMIT: usize = 10_000; pub const DATABASE_EXPORT_PAGE_SIZE: usize = 500; pub const DATABASE_EXPORT_INSERT_BATCH_SIZE: usize = 100; +#[derive(Debug, Clone, PartialEq, Eq)] +struct PostgresExportSequence { + name: String, + data_type: String, + start_value: String, + min_value: String, + max_value: String, + increment: String, + cycle: bool, + cache_value: String, + last_value: Option, + owner_table: Option, + owner_column: Option, +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct ExportedTableSql { @@ -354,6 +369,156 @@ fn normalize_export_table_ddl(ddl: &str, database_type: Option) -> LEGACY_MYSQL_ROW_FORMAT_RE.replace_all(ddl, "ROW_FORMAT=DYNAMIC").into_owned() } +fn postgres_sequence_qualified_name(schema: &str, sequence_name: &str) -> String { + let db_type = DatabaseType::Postgres; + if schema.trim().is_empty() { + quote_identifier(sequence_name, &db_type) + } else { + format!("{}.{}", quote_identifier(schema, &db_type), quote_identifier(sequence_name, &db_type)) + } +} + +fn postgres_string_literal(value: &str) -> String { + format!("'{}'", value.replace('\'', "''")) +} + +fn generate_postgres_sequence_create_ddl(sequence: &PostgresExportSequence, schema: &str) -> String { + let qualified_name = postgres_sequence_qualified_name(schema, &sequence.name); + let cycle = if sequence.cycle { "CYCLE" } else { "NO CYCLE" }; + format!( + "CREATE SEQUENCE IF NOT EXISTS {qualified_name}\n AS {data_type}\n START WITH {start_value}\n INCREMENT BY {increment}\n MINVALUE {min_value}\n MAXVALUE {max_value}\n CACHE {cache_value}\n {cycle}", + data_type = sequence.data_type, + start_value = sequence.start_value, + increment = sequence.increment, + min_value = sequence.min_value, + max_value = sequence.max_value, + cache_value = sequence.cache_value, + ) +} + +fn generate_postgres_sequence_owner_ddl(sequence: &PostgresExportSequence, schema: &str) -> Option { + let owner_table = sequence.owner_table.as_deref()?; + let owner_column = sequence.owner_column.as_deref()?; + Some(format!( + "ALTER SEQUENCE {} OWNED BY {}.{}", + postgres_sequence_qualified_name(schema, &sequence.name), + crate::transfer::qualified_table(owner_table, schema, &DatabaseType::Postgres), + quote_identifier(owner_column, &DatabaseType::Postgres) + )) +} + +fn generate_postgres_sequence_setval_sql(sequence: &PostgresExportSequence, schema: &str) -> Option { + let last_value = sequence.last_value.as_deref()?.trim(); + if last_value.is_empty() { + return None; + } + + let sequence_literal = postgres_string_literal(&postgres_sequence_qualified_name(schema, &sequence.name)); + match (sequence.owner_table.as_deref(), sequence.owner_column.as_deref()) { + (Some(owner_table), Some(owner_column)) => { + let owner_table = crate::transfer::qualified_table(owner_table, schema, &DatabaseType::Postgres); + let owner_column = quote_identifier(owner_column, &DatabaseType::Postgres); + Some(format!( + "SELECT setval({sequence_literal}, GREATEST(COALESCE(MAX({owner_column}), {last_value}), {last_value}), true) FROM {owner_table}" + )) + } + _ => Some(format!("SELECT setval({sequence_literal}, {last_value}, true)")), + } +} + +async fn list_postgres_export_sequences( + state: &crate::connection::AppState, + pool_key: &str, + schema: &str, + selected_tables: &[String], + include_objects: bool, +) -> Result, String> { + let pool = { + let connections = state.connections.read().await; + match connections.get(pool_key) { + Some(crate::connection::PoolKind::Postgres(pool)) => pool.clone(), + _ => return Ok(Vec::new()), + } + }; + let client = pool.get().await.map_err(|e| e.to_string())?; + let rows = client + .query( + "SELECT c.relname, \ + COALESCE(format_type(s.seqtypid, NULL), 'bigint'), \ + COALESCE(s.seqstart::text, '1'), \ + COALESCE(s.seqmin::text, '1'), \ + COALESCE(s.seqmax::text, '9223372036854775807'), \ + COALESCE(s.seqincrement::text, '1'), \ + COALESCE(s.seqcycle, false), \ + COALESCE(s.seqcache::text, '1'), \ + t.relname, \ + a.attname \ + FROM pg_class c \ + JOIN pg_namespace n ON n.oid = c.relnamespace \ + LEFT JOIN pg_sequence s ON s.seqrelid = c.oid \ + LEFT JOIN pg_depend d ON d.classid = 'pg_class'::regclass \ + AND d.objid = c.oid \ + AND d.refclassid = 'pg_class'::regclass \ + AND d.deptype IN ('a', 'i') \ + LEFT JOIN pg_class t ON t.oid = d.refobjid \ + LEFT JOIN pg_namespace tn ON tn.oid = t.relnamespace AND tn.nspname = n.nspname \ + LEFT JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = d.refobjsubid \ + WHERE c.relkind = 'S' AND n.nspname = $1 \ + ORDER BY c.relname", + &[&schema], + ) + .await + .map_err(|e| e.to_string())?; + + let selected: HashSet<&str> = selected_tables.iter().map(String::as_str).collect(); + let mut sequences = rows + .iter() + .map(|row| PostgresExportSequence { + name: row.get::<_, String>(0), + data_type: row.get::<_, String>(1), + start_value: row.get::<_, String>(2), + min_value: row.get::<_, String>(3), + max_value: row.get::<_, String>(4), + increment: row.get::<_, String>(5), + cycle: row.get::<_, bool>(6), + cache_value: row.get::<_, String>(7), + last_value: None, + owner_table: row.get::<_, Option>(8), + owner_column: row.get::<_, Option>(9), + }) + .filter(|sequence| { + selected.is_empty() + || sequence.owner_table.as_deref().map(|owner_table| selected.contains(owner_table)).unwrap_or(false) + }) + .filter(|sequence| sequence.owner_table.is_some() || (include_objects && selected.is_empty())) + .collect::>(); + + if sequences.is_empty() { + return Ok(sequences); + } + + if let Ok(rows) = client + .query( + "SELECT c.relname, pg_sequence_last_value(c.oid)::text \ + FROM pg_class c \ + JOIN pg_namespace n ON n.oid = c.relnamespace \ + WHERE c.relkind = 'S' AND n.nspname = $1", + &[&schema], + ) + .await + { + for row in rows { + let name: String = row.get(0); + let last_value: Option = row.get(1); + if let Some(sequence) = sequences.iter_mut().find(|sequence| sequence.name == name) { + sequence.last_value = last_value; + } + } + } + + Ok(sequences) +} + pub async fn is_export_cancelled(export_id: &str) -> bool { EXPORT_CANCELLED.read().await.contains(export_id) } @@ -415,6 +580,26 @@ pub async fn export_database_sql_core( // 7. Separate tables and views let mut 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(); + let postgres_sequences = if request.include_structure && matches!(db_type, DatabaseType::Postgres) { + match list_postgres_export_sequences( + state, + &pool_key, + &request.schema, + &request.selected_tables, + request.include_objects, + ) + .await + { + Ok(sequences) => sequences, + Err(e) => { + writeln!(file, "-- ERROR exporting sequences: {e}") + .map_err(|e| format!("Failed to write file: {e}"))?; + Vec::new() + } + } + } else { + Vec::new() + }; // Sort tables by foreign key dependency so referenced (parent) tables are // exported before referencing (child) tables. @@ -435,7 +620,7 @@ pub async fn export_database_sql_core( } // 8. Calculate total objects - let mut total_objects = tables.len() + views.len(); + let mut total_objects = tables.len() + views.len() + postgres_sequences.len(); // We'll add procedures/functions count later if include_objects let mut procedures: Vec = Vec::new(); @@ -462,6 +647,27 @@ pub async fn export_database_sql_core( // Export tables let batch_size = if request.batch_size == 0 { 1000 } else { request.batch_size }; + for sequence in postgres_sequences.iter().filter(|sequence| sequence.owner_table.is_none()) { + 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: sequence.name.clone(), + object_index, + total_objects, + rows_exported: 0, + total_rows: None, + status: ExportStatus::Running, + error: None, + }); + + writeln!(file, "{};\n", generate_postgres_sequence_create_ddl(sequence, &request.schema)) + .map_err(|e| format!("Failed to write file: {e}"))?; + object_index += 1; + } + for table_info in &tables { // Check cancellation if is_export_cancelled(&request.export_id).await { @@ -498,6 +704,25 @@ pub async fn export_database_sql_core( writeln!(file, "{}\n", drop_table_if_exists_sql(table_name, &request.schema, &db_type)) .map_err(|e| format!("Failed to write file: {e}"))?; } + for sequence in postgres_sequences + .iter() + .filter(|sequence| sequence.owner_table.as_deref() == Some(table_name.as_str())) + { + on_progress(ExportProgress { + export_id: request.export_id.clone(), + current_object: sequence.name.clone(), + object_index, + total_objects, + rows_exported: 0, + total_rows: None, + status: ExportStatus::Running, + error: None, + }); + + writeln!(file, "{};\n", generate_postgres_sequence_create_ddl(sequence, &request.schema)) + .map_err(|e| format!("Failed to write file: {e}"))?; + object_index += 1; + } match crate::schema::get_table_ddl_core( state, &request.connection_id, @@ -634,6 +859,19 @@ pub async fn export_database_sql_core( object_index += 1; } + if request.include_structure && !postgres_sequences.is_empty() { + for sequence in &postgres_sequences { + if let Some(sql) = generate_postgres_sequence_owner_ddl(sequence, &request.schema) { + writeln!(file, "{};\n", sql).map_err(|e| format!("Failed to write file: {e}"))?; + } + } + for sequence in &postgres_sequences { + if let Some(sql) = generate_postgres_sequence_setval_sql(sequence, &request.schema) { + writeln!(file, "{};\n", sql).map_err(|e| format!("Failed to write file: {e}"))?; + } + } + } + // Export views (if include_objects) if request.include_objects { for view_info in &views { @@ -800,8 +1038,9 @@ fn drop_table_if_exists_sql(table_name: &str, schema: &str, db_type: &DatabaseTy mod tests { use super::{ build_database_sql_export, build_export_insert_statements, drop_table_if_exists_sql, - filter_selected_table_infos, format_export_sql_literal, normalize_export_table_ddl, - BuildDatabaseSqlExportOptions, BuildExportInsertStatementsOptions, ExportedTableSql, + filter_selected_table_infos, format_export_sql_literal, generate_postgres_sequence_create_ddl, + generate_postgres_sequence_owner_ddl, generate_postgres_sequence_setval_sql, normalize_export_table_ddl, + BuildDatabaseSqlExportOptions, BuildExportInsertStatementsOptions, ExportedTableSql, PostgresExportSequence, DATABASE_EXPORT_INSERT_BATCH_SIZE, DATABASE_EXPORT_ROW_LIMIT, }; use crate::models::connection::DatabaseType; @@ -992,4 +1231,56 @@ mod tests { assert_eq!(normalize_export_table_ddl(mysql_ddl, Some(DatabaseType::Mysql)), mysql_ddl); assert_eq!(normalize_export_table_ddl(postgres_ddl, Some(DatabaseType::Postgres)), postgres_ddl); } + + fn postgres_sequence(name: &str) -> PostgresExportSequence { + PostgresExportSequence { + name: name.to_string(), + data_type: "integer".to_string(), + start_value: "1".to_string(), + min_value: "1".to_string(), + max_value: "2147483647".to_string(), + increment: "1".to_string(), + cycle: false, + cache_value: "1".to_string(), + last_value: Some("42".to_string()), + owner_table: Some("permissions".to_string()), + owner_column: Some("id".to_string()), + } + } + + #[test] + fn postgres_sequence_create_ddl_is_importable_before_table_ddl() { + let ddl = generate_postgres_sequence_create_ddl(&postgres_sequence("permissions_id_seq"), "public"); + + assert_eq!( + ddl, + [ + "CREATE SEQUENCE IF NOT EXISTS \"public\".\"permissions_id_seq\"", + " AS integer", + " START WITH 1", + " INCREMENT BY 1", + " MINVALUE 1", + " MAXVALUE 2147483647", + " CACHE 1", + " NO CYCLE", + ] + .join("\n") + ); + } + + #[test] + fn postgres_sequence_owner_and_setval_sql_are_qualified() { + let sequence = postgres_sequence("permissions_id_seq"); + + assert_eq!( + generate_postgres_sequence_owner_ddl(&sequence, "public").as_deref(), + Some("ALTER SEQUENCE \"public\".\"permissions_id_seq\" OWNED BY \"public\".\"permissions\".\"id\"") + ); + assert_eq!( + generate_postgres_sequence_setval_sql(&sequence, "public").as_deref(), + Some( + "SELECT setval('\"public\".\"permissions_id_seq\"', GREATEST(COALESCE(MAX(\"id\"), 42), 42), true) FROM \"public\".\"permissions\"" + ) + ); + } }