Skip to main content

sz_orm_core/
migration.rs

1//! Migration system
2//!
3//! Provides database schema migration management
4
5use crate::db_type::DbType;
6use crate::error::DbError;
7use std::path::PathBuf;
8
9/// 数据库迁移定义
10#[derive(Clone)]
11pub struct Migration {
12    /// 版本号
13    pub version: String,
14    /// 迁移名称
15    pub name: String,
16    /// 正向 SQL(升级)
17    pub sql_up: String,
18    /// 反向 SQL(回滚)
19    pub sql_down: String,
20    /// 批次号(0 表示未执行)
21    pub batch: i32,
22    /// 执行时间(None 表示未执行)
23    pub executed_at: Option<chrono::DateTime<chrono::Utc>>,
24}
25
26impl Migration {
27    /// 创建迁移,指定版本号、名称、正向和反向 SQL
28    pub fn new(version: &str, name: &str, sql_up: &str, sql_down: &str) -> Self {
29        Self {
30            version: version.to_string(),
31            name: name.to_string(),
32            sql_up: sql_up.to_string(),
33            sql_down: sql_down.to_string(),
34            batch: 0,
35            executed_at: None,
36        }
37    }
38
39    /// 设置批次号
40    pub fn with_batch(mut self, batch: i32) -> Self {
41        self.batch = batch;
42        self
43    }
44
45    /// 设置执行时间
46    pub fn with_executed_at(mut self, time: chrono::DateTime<chrono::Utc>) -> Self {
47        self.executed_at = Some(time);
48        self
49    }
50}
51
52impl std::fmt::Debug for Migration {
53    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54        f.debug_struct("Migration")
55            .field("version", &self.version)
56            .field("name", &self.name)
57            .field("batch", &self.batch)
58            .finish()
59    }
60}
61
62/// 迁移解析器 trait
63pub trait MigrationResolver: Send + Sync {
64    /// 解析指定数据库类型的迁移列表
65    fn resolve(&self, db_type: DbType) -> Result<Vec<Migration>, DbError>;
66}
67
68/// 文件迁移解析器
69pub struct FileMigrationResolver {
70    /// 迁移文件目录
71    pub path: PathBuf,
72}
73
74impl FileMigrationResolver {
75    /// 创建文件迁移解析器,指定迁移文件目录
76    pub fn new(path: PathBuf) -> Self {
77        Self { path }
78    }
79}
80
81impl MigrationResolver for FileMigrationResolver {
82    fn resolve(&self, db_type: DbType) -> Result<Vec<Migration>, DbError> {
83        let mut migrations = Vec::new();
84
85        // 读取迁移目录
86        let entries = match std::fs::read_dir(&self.path) {
87            Ok(entries) => entries,
88            Err(e) => {
89                return Err(DbError::MigrationError(format!(
90                    "Cannot read migration directory {}: {}",
91                    self.path.display(),
92                    e
93                )));
94            }
95        };
96
97        let _ = db_type; // 当前实现不区分数据库类型
98
99        // 收集所有 .sql 文件
100        let mut sql_files: Vec<std::path::PathBuf> = Vec::new();
101        for entry in entries {
102            let entry = entry.map_err(|e| {
103                DbError::MigrationError(format!("Cannot read directory entry: {}", e))
104            })?;
105            let path = entry.path();
106            if path.extension().and_then(|s| s.to_str()) == Some("sql") {
107                sql_files.push(path);
108            }
109        }
110
111        // 按文件名排序
112        sql_files.sort();
113
114        // 解析文件名格式:<version>_<name>_up.sql 或 <version>_<name>_down.sql
115        // 也支持简单的 <name>.sql(不区分 up/down)
116        let mut version_map: std::collections::HashMap<
117            String,
118            (Option<String>, Option<String>, String),
119        > = std::collections::HashMap::new();
120
121        for path in sql_files {
122            let filename = match path.file_stem().and_then(|s| s.to_str()) {
123                Some(name) => name.to_string(),
124                None => continue,
125            };
126
127            let content = std::fs::read_to_string(&path).map_err(|e| {
128                DbError::MigrationError(format!(
129                    "Cannot read migration file {}: {}",
130                    path.display(),
131                    e
132                ))
133            })?;
134
135            // 尝试解析文件名
136            if filename.ends_with("_up") {
137                let base = &filename[..filename.len() - 3];
138                let (version, name) = parse_migration_filename(base);
139                let entry = version_map
140                    .entry(version.clone())
141                    .or_insert((None, None, name));
142                entry.0 = Some(content);
143            } else if filename.ends_with("_down") {
144                let base = &filename[..filename.len() - 5];
145                let (version, name) = parse_migration_filename(base);
146                let entry = version_map
147                    .entry(version.clone())
148                    .or_insert((None, None, name));
149                entry.1 = Some(content);
150            } else {
151                // 简单格式:整个文件作为 up SQL,down 为空
152                let (version, name) = parse_migration_filename(&filename);
153                let entry = version_map
154                    .entry(version.clone())
155                    .or_insert((None, None, name));
156                if entry.0.is_none() {
157                    entry.0 = Some(content);
158                }
159            }
160        }
161
162        // 转换为 Migration 列表并按 version 排序
163        type VersionEntry = (Option<String>, Option<String>, String);
164        let mut sorted_versions: Vec<(String, VersionEntry)> = version_map.into_iter().collect();
165        sorted_versions.sort_by(|a, b| a.0.cmp(&b.0));
166
167        for (version, (sql_up, sql_down, name)) in sorted_versions {
168            let migration = Migration::new(
169                &version,
170                &name,
171                sql_up.unwrap_or_default().as_str(),
172                sql_down.unwrap_or_default().as_str(),
173            );
174            migrations.push(migration);
175        }
176
177        Ok(migrations)
178    }
179}
180
181/// 解析迁移文件名:格式 <version>_<name>,如 "001_create_users"
182fn parse_migration_filename(filename: &str) -> (String, String) {
183    if let Some(underscore_pos) = filename.find('_') {
184        let version = filename[..underscore_pos].to_string();
185        let name = filename[underscore_pos + 1..].to_string();
186        (version, name)
187    } else {
188        // 没有下划线,整个作为 version,name 为空
189        (filename.to_string(), filename.to_string())
190    }
191}
192
193/// 迁移上下文
194pub struct MigrationContext {
195    /// 迁移记录表名
196    pub table_name: String,
197    /// 数据库连接
198    pub connection: Option<Box<dyn crate::pool::Connection>>,
199    /// 数据库类型(用于判断是否支持 DDL 事务包裹)
200    pub db_type: Option<DbType>,
201}
202
203impl Default for MigrationContext {
204    fn default() -> Self {
205        Self {
206            table_name: "__migrations".to_string(),
207            connection: None,
208            db_type: None,
209        }
210    }
211}
212
213impl MigrationContext {
214    /// 设置数据库类型
215    pub fn with_db_type(mut self, db_type: DbType) -> Self {
216        self.db_type = Some(db_type);
217        self
218    }
219}
220
221/// P1-4:校验迁移版本号安全性
222///
223/// 允许的格式:字母、数字、下划线、连字符、点(支持 "001"、"20240101"、"v1.0"、"create_users" 等常见格式)。
224/// 拒绝单引号、分号、注释、空格等可能用于 SQL 注入的字符。
225///
226/// 与 `validate_identifier` 的区别:版本号允许以数字开头(如 "001"),
227/// 而标识符通常要求以字母或下划线开头。
228fn validate_migration_version(version: &str) -> Result<(), DbError> {
229    if version.is_empty() || version.len() > 255 {
230        return Err(DbError::InvalidInput(format!(
231            "invalid migration version: empty or too long (max 255 chars): {:?}",
232            version
233        )));
234    }
235    // 只允许字母、数字、下划线、连字符、点
236    let valid = version
237        .chars()
238        .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-' || c == '.');
239    if !valid {
240        return Err(DbError::InvalidInput(format!(
241            "invalid migration version: only ASCII alphanumeric, underscore, hyphen, dot allowed, got {:?}",
242            version
243        )));
244    }
245    // 额外拒绝 "--" 注释序列
246    if version.contains("--") {
247        return Err(DbError::InvalidInput(format!(
248            "invalid migration version: SQL comment sequence '--' not allowed: {:?}",
249            version
250        )));
251    }
252    Ok(())
253}
254
255/// 判断指定数据库方言是否支持 DDL 事务
256///
257/// - PostgreSQL:✅ 支持 DDL 事务(CREATE/ALTER/DROP 可回滚)
258/// - SQLite:✅ 支持 DDL 事务
259/// - MySQL:❌ DDL 语句隐式提交,无法回滚
260/// - Oracle:❌ DDL 语句前后隐式 COMMIT
261/// - SQL Server:❌ 部分 DDL 不支持事务内执行(保守处理)
262/// - 其他:❌ 默认不支持
263fn supports_ddl_transactions(db_type: DbType) -> bool {
264    matches!(db_type, DbType::PostgreSQL | DbType::Sqlite)
265}
266
267/// 迁移方向
268#[derive(Debug, Clone, PartialEq)]
269pub enum MigrationDirection {
270    /// 正向迁移(升级)
271    Up,
272    /// 反向迁移(回滚)
273    Down,
274}
275
276/// 迁移执行器
277pub struct Migrator {
278    context: MigrationContext,
279    migrations: Vec<Migration>,
280}
281
282impl Migrator {
283    /// 创建迁移执行器,指定上下文
284    pub fn new(context: MigrationContext) -> Self {
285        Self {
286            context,
287            migrations: Vec::new(),
288        }
289    }
290
291    /// 添加单个迁移
292    pub fn add_migration(mut self, migration: Migration) -> Self {
293        self.migrations.push(migration);
294        self
295    }
296
297    /// 添加多个迁移
298    pub fn add_migrations(mut self, migrations: Vec<Migration>) -> Self {
299        self.migrations.extend(migrations);
300        self
301    }
302
303    /// 返回所有迁移的引用
304    pub fn get_migrations(&self) -> &Vec<Migration> {
305        &self.migrations
306    }
307
308    /// 返回待执行的迁移(batch == 0)
309    pub fn get_pending_migrations(&self) -> Vec<&Migration> {
310        self.migrations.iter().filter(|m| m.batch == 0).collect()
311    }
312
313    /// 返回已执行的迁移(batch > 0)
314    pub fn get_applied_migrations(&self) -> Vec<&Migration> {
315        self.migrations.iter().filter(|m| m.batch > 0).collect()
316    }
317
318    /// 返回最新版本号
319    pub fn latest_version(&self) -> Option<&str> {
320        self.migrations.last().map(|m| m.version.as_str())
321    }
322
323    /// 按版本号查找迁移
324    pub fn find_migration(&self, version: &str) -> Option<&Migration> {
325        self.migrations.iter().find(|m| m.version == version)
326    }
327
328    /// 检测迁移版本冲突(重复版本号)
329    ///
330    /// 返回第一个冲突的版本号(如有)。
331    /// 在 `migrate`/`up`/`down` 等执行方法入口处调用,确保迁移列表无重复版本。
332    pub fn check_version_conflicts(&self) -> Result<(), DbError> {
333        let mut seen = std::collections::HashSet::new();
334        for m in &self.migrations {
335            if !seen.insert(&m.version) {
336                return Err(DbError::MigrationError(format!(
337                    "迁移版本冲突:版本号 '{}' 重复定义",
338                    m.version
339                )));
340            }
341        }
342        Ok(())
343    }
344
345    // ==================== P1-4:__migrations 持久化表自动创建 ====================
346
347    /// P1-4:生成 `__migrations` 表的 CREATE TABLE IF NOT EXISTS SQL
348    ///
349    /// 表结构(跨方言兼容):
350    /// ```sql
351    /// CREATE TABLE IF NOT EXISTS __migrations (
352    ///     version VARCHAR(255) NOT NULL PRIMARY KEY,
353    ///     name VARCHAR(255),
354    ///     batch INTEGER NOT NULL,
355    ///     executed_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
356    /// );
357    /// ```
358    ///
359    /// 不同方言的差异:
360    /// - MySQL/PG/SQLite:`TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP`
361    /// - Oracle:`TIMESTAMP DEFAULT CURRENT_TIMESTAMP`
362    /// - SQL Server:`DATETIME DEFAULT GETDATE()`
363    pub fn build_create_migrations_table_sql(&self) -> String {
364        let table = &self.context.table_name;
365        let timestamp_default = match self.context.db_type {
366            Some(DbType::SqlServer) => "DATETIME DEFAULT GETDATE()",
367            Some(DbType::Oracle) => "TIMESTAMP DEFAULT CURRENT_TIMESTAMP",
368            _ => "TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP",
369        };
370        format!(
371            "CREATE TABLE IF NOT EXISTS {} (\
372                version VARCHAR(255) NOT NULL PRIMARY KEY, \
373                name VARCHAR(255), \
374                batch INTEGER NOT NULL, \
375                executed_at {ts}\
376            )",
377            table,
378            ts = timestamp_default
379        )
380    }
381
382    /// P1-4:确保 `__migrations` 表存在(若连接可用)
383    ///
384    /// 在 `migrate`/`up`/`down`/`rollback` 等方法入口处调用,
385    /// 自动创建持久化表(若不存在)。
386    ///
387    /// 若 `context.connection` 为 None,跳过(纯内存模式)。
388    async fn ensure_migrations_table(&mut self) -> Result<(), DbError> {
389        let sql = self.build_create_migrations_table_sql();
390        if let Some(ref mut conn) = self.context.connection {
391            conn.execute(&sql).await?;
392        }
393        Ok(())
394    }
395
396    /// P1-4:从 `__migrations` 表加载已执行的迁移记录
397    ///
398    /// 返回 `HashMap<version, batch>`,表示每个已执行迁移的版本号和批次号。
399    /// 调用方应根据返回值更新 `self.migrations` 中对应迁移的 `batch` 字段。
400    ///
401    /// 若 `context.connection` 为 None,返回空 map(纯内存模式)。
402    async fn load_applied_migrations(
403        &mut self,
404    ) -> Result<std::collections::HashMap<String, i32>, DbError> {
405        let mut applied = std::collections::HashMap::new();
406        if let Some(ref mut conn) = self.context.connection {
407            let sql = format!("SELECT version, batch FROM {}", self.context.table_name);
408            let rows = conn.query(&sql).await?;
409            for row in rows {
410                if let Some(crate::Value::String(version)) = row.get("version") {
411                    let batch = match row.get("batch") {
412                        Some(crate::Value::I32(b)) => *b,
413                        Some(crate::Value::I64(b)) => *b as i32,
414                        _ => 0,
415                    };
416                    applied.insert(version.clone(), batch);
417                }
418            }
419        }
420        Ok(applied)
421    }
422
423    /// P1-4:根据 `__migrations` 表的记录同步内存中的迁移状态
424    ///
425    /// 在 `migrate`/`up`/`down` 等方法入口处调用,确保内存状态与数据库持久化状态一致。
426    async fn sync_state_from_db(&mut self) -> Result<(), DbError> {
427        let applied = self.load_applied_migrations().await?;
428        for migration in &mut self.migrations {
429            if let Some(batch) = applied.get(&migration.version) {
430                migration.batch = *batch;
431                migration.executed_at = Some(chrono::Utc::now());
432            } else {
433                migration.batch = 0;
434                migration.executed_at = None;
435            }
436        }
437        Ok(())
438    }
439
440    /// P1-4:向 `__migrations` 表插入一条执行记录
441    ///
442    /// 在每个迁移成功执行 up SQL 后调用。
443    async fn record_migration(
444        &mut self,
445        version: &str,
446        name: &str,
447        batch: i32,
448    ) -> Result<(), DbError> {
449        if let Some(ref mut conn) = self.context.connection {
450            // 安全校验:version 允许字母/数字/下划线/连字符/点(支持 "001"、"20240101"、"v1.0" 等格式),
451            // 但拒绝单引号、分号、注释等危险字符
452            validate_migration_version(version)?;
453            // name 可能为空或含下划线,放宽校验:仅拒绝单引号和分号
454            if name.contains('\'') || name.contains(';') {
455                return Err(DbError::MigrationError(format!("非法迁移名称: {}", name)));
456            }
457            let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S");
458            let sql = format!(
459                "INSERT INTO {} (version, name, batch, executed_at) VALUES ('{}', '{}', {}, '{}')",
460                self.context.table_name, version, name, batch, now
461            );
462            conn.execute(&sql).await?;
463        }
464        Ok(())
465    }
466
467    /// P1-4:从 `__migrations` 表删除一条执行记录
468    ///
469    /// 在每个迁移成功执行 down SQL 后调用。
470    async fn remove_migration(&mut self, version: &str) -> Result<(), DbError> {
471        if let Some(ref mut conn) = self.context.connection {
472            validate_migration_version(version)?;
473            let sql = format!(
474                "DELETE FROM {} WHERE version = '{}'",
475                self.context.table_name, version
476            );
477            conn.execute(&sql).await?;
478        }
479        Ok(())
480    }
481
482    /// 执行所有待迁移(batch=0)的 up SQL
483    ///
484    /// 若数据库方言支持 DDL 事务(PostgreSQL/SQLite),则用事务包裹所有待执行迁移,
485    /// 任一迁移失败时回滚全部变更,避免部分迁移导致的状态不一致。
486    /// 不支持 DDL 事务的方言(MySQL/Oracle/SQL Server)逐条执行,失败时保留已执行的变更。
487    ///
488    /// P1-4:若连接可用,会自动创建 `__migrations` 持久化表,并在执行前从表中
489    /// 同步已执行记录、执行后向表插入新记录。
490    pub async fn migrate(&mut self) -> Result<Vec<String>, DbError> {
491        // 版本冲突检测
492        self.check_version_conflicts()?;
493
494        // P1-4:确保 __migrations 表存在
495        self.ensure_migrations_table().await?;
496
497        // P1-4:从 __migrations 表同步已执行状态
498        self.sync_state_from_db().await?;
499
500        let mut applied = Vec::new();
501        let current_batch = self.migrations.iter().map(|m| m.batch).max().unwrap_or(0) + 1;
502
503        // 收集待迁移的索引(避免在循环中再次 position(),消除 O(n²) 复杂度)
504        let pending_indices: Vec<usize> = self
505            .migrations
506            .iter()
507            .enumerate()
508            .filter(|(_, m)| m.batch == 0)
509            .map(|(idx, _)| idx)
510            .collect();
511
512        if pending_indices.is_empty() {
513            return Ok(applied);
514        }
515
516        // 判断是否需要事务包裹
517        let use_transaction = self
518            .context
519            .db_type
520            .map(supports_ddl_transactions)
521            .unwrap_or(false);
522
523        // 开启事务(若方言支持)
524        if use_transaction {
525            if let Some(ref mut conn) = self.context.connection {
526                conn.begin_transaction().await?;
527            }
528        }
529
530        // 逐条执行迁移
531        for migration_idx in &pending_indices {
532            let sql_up = self.migrations[*migration_idx].sql_up.clone();
533            let version = self.migrations[*migration_idx].version.clone();
534            let name = self.migrations[*migration_idx].name.clone();
535
536            let exec_result = async {
537                if let Some(ref mut conn) = self.context.connection {
538                    if !sql_up.is_empty() {
539                        conn.execute(&sql_up).await?;
540                    }
541                }
542                Ok::<(), DbError>(())
543            }
544            .await;
545
546            if let Err(e) = exec_result {
547                // 事务包裹下回滚
548                if use_transaction {
549                    if let Some(ref mut conn) = self.context.connection {
550                        let _ = conn.rollback().await;
551                    }
552                }
553                return Err(e);
554            }
555
556            // 标记为已执行
557            let now = chrono::Utc::now();
558            self.migrations[*migration_idx].batch = current_batch;
559            self.migrations[*migration_idx].executed_at = Some(now);
560
561            // P1-4:向 __migrations 表插入记录
562            if let Err(e) = self.record_migration(&version, &name, current_batch).await {
563                // 事务包裹下回滚
564                if use_transaction {
565                    if let Some(ref mut conn) = self.context.connection {
566                        let _ = conn.rollback().await;
567                    }
568                }
569                return Err(e);
570            }
571
572            applied.push(version);
573        }
574
575        // 提交事务(若方言支持)
576        if use_transaction {
577            if let Some(ref mut conn) = self.context.connection {
578                conn.commit().await?;
579            }
580        }
581
582        Ok(applied)
583    }
584
585    /// 回滚指定版本(执行 down SQL)
586    ///
587    /// P1-4:若连接可用,会自动创建 `__migrations` 表,并在回滚后从表中删除记录。
588    pub async fn rollback(&mut self, version: &str) -> Result<(), DbError> {
589        // P1-4:确保 __migrations 表存在并同步状态
590        self.ensure_migrations_table().await?;
591        self.sync_state_from_db().await?;
592
593        let migration_idx = self
594            .migrations
595            .iter()
596            .position(|m| m.version == version)
597            .ok_or_else(|| DbError::MigrationError(format!("Migration {} not found", version)))?;
598
599        if self.migrations[migration_idx].batch == 0 {
600            return Err(DbError::MigrationError(format!(
601                "Migration {} not applied",
602                version
603            )));
604        }
605
606        let sql_down = self.migrations[migration_idx].sql_down.clone();
607
608        if let Some(ref mut conn) = self.context.connection {
609            if !sql_down.is_empty() {
610                conn.execute(&sql_down).await?;
611            }
612        }
613
614        self.migrations[migration_idx].batch = 0;
615        self.migrations[migration_idx].executed_at = None;
616
617        // P1-4:从 __migrations 表删除记录
618        self.remove_migration(version).await?;
619
620        Ok(())
621    }
622
623    /// 执行到指定版本(包括该版本)
624    ///
625    /// P1-4:若连接可用,会自动创建 `__migrations` 表,并在执行前从表中
626    /// 同步已执行记录、执行后向表插入新记录。
627    pub async fn up(&mut self, target_version: Option<&str>) -> Result<Vec<String>, DbError> {
628        // 版本冲突检测
629        self.check_version_conflicts()?;
630
631        // P1-4:确保 __migrations 表存在并同步状态
632        self.ensure_migrations_table().await?;
633        self.sync_state_from_db().await?;
634
635        let mut applied = Vec::new();
636        let current_batch = self.migrations.iter().map(|m| m.batch).max().unwrap_or(0) + 1;
637
638        // 收集待执行迁移的版本和索引(避免在循环中借用冲突)
639        let pending: Vec<(usize, String, String)> = self
640            .migrations
641            .iter()
642            .enumerate()
643            .filter(|(_, m)| m.batch == 0)
644            .take_while(|(_, m)| {
645                if let Some(target) = target_version {
646                    m.version.as_str() <= target
647                } else {
648                    true
649                }
650            })
651            .map(|(idx, m)| (idx, m.version.clone(), m.name.clone()))
652            .collect();
653
654        for (idx, version, name) in pending {
655            let sql_up = self.migrations[idx].sql_up.clone();
656            if let Some(ref mut conn) = self.context.connection {
657                if !sql_up.is_empty() {
658                    conn.execute(&sql_up).await?;
659                }
660            }
661
662            self.migrations[idx].batch = current_batch;
663            self.migrations[idx].executed_at = Some(chrono::Utc::now());
664
665            // P1-4:向 __migrations 表插入记录
666            self.record_migration(&version, &name, current_batch)
667                .await?;
668
669            applied.push(version);
670        }
671
672        Ok(applied)
673    }
674
675    /// 回滚到指定版本(执行该版本之后所有迁移的 down SQL)
676    ///
677    /// P1-4:若连接可用,会自动创建 `__migrations` 表,并在回滚后从表中删除记录。
678    pub async fn down(&mut self, target_version: Option<&str>) -> Result<Vec<String>, DbError> {
679        // 版本冲突检测
680        self.check_version_conflicts()?;
681
682        // P1-4:确保 __migrations 表存在并同步状态
683        self.ensure_migrations_table().await?;
684        self.sync_state_from_db().await?;
685
686        let mut rolled_back = Vec::new();
687
688        // 从后往前回滚,收集待回滚的索引和版本
689        let mut indices: Vec<usize> = (0..self.migrations.len()).collect();
690        indices.reverse();
691
692        let pending_rollback: Vec<(usize, String)> = indices
693            .iter()
694            .filter(|&&idx| self.migrations[idx].batch > 0)
695            .take_while(|&&idx| {
696                if let Some(target) = target_version {
697                    self.migrations[idx].version.as_str() > target
698                } else {
699                    true
700                }
701            })
702            .map(|&idx| (idx, self.migrations[idx].version.clone()))
703            .collect();
704
705        for (idx, version) in pending_rollback {
706            let sql_down = self.migrations[idx].sql_down.clone();
707            if let Some(ref mut conn) = self.context.connection {
708                if !sql_down.is_empty() {
709                    conn.execute(&sql_down).await?;
710                }
711            }
712
713            self.migrations[idx].batch = 0;
714            self.migrations[idx].executed_at = None;
715
716            // P1-4:从 __migrations 表删除记录
717            self.remove_migration(&version).await?;
718
719            rolled_back.push(version);
720        }
721
722        Ok(rolled_back)
723    }
724
725    /// 重置:回滚所有已执行的迁移,然后重新执行
726    pub async fn reset(&mut self) -> Result<Vec<String>, DbError> {
727        // 先全部回滚
728        self.down(None).await?;
729        // 再全部执行
730        self.migrate().await
731    }
732
733    /// 刷新:回滚所有已执行的迁移,然后重新执行
734    pub async fn refresh(&mut self) -> Result<Vec<String>, DbError> {
735        self.reset().await
736    }
737
738    /// 获取迁移进度
739    pub fn progress(&self) -> MigrationProgress {
740        let total = self.migrations.len();
741        let applied = self.migrations.iter().filter(|m| m.batch > 0).count();
742        MigrationProgress::new(total, applied)
743    }
744}
745
746/// 迁移进度
747#[derive(Debug, Clone)]
748pub struct MigrationProgress {
749    /// 总迁移数
750    pub total: usize,
751    /// 已执行数
752    pub applied: usize,
753    /// 待执行数
754    pub pending: usize,
755    /// 当前批次号
756    pub current_batch: i32,
757}
758
759impl MigrationProgress {
760    /// 创建迁移进度,指定总数和已执行数
761    pub fn new(total: usize, applied: usize) -> Self {
762        Self {
763            total,
764            applied,
765            pending: total - applied,
766            current_batch: 0,
767        }
768    }
769
770    /// 返回完成百分比
771    pub fn percent_complete(&self) -> f64 {
772        if self.total == 0 {
773            return 100.0;
774        }
775        (self.applied as f64 / self.total as f64) * 100.0
776    }
777}
778
779/// Schema 构建器
780pub struct SchemaBuilder {
781    table_name: String,
782    columns: Vec<ColumnDef>,
783    indexes: Vec<IndexDef>,
784    foreign_keys: Vec<ForeignKeyDef>,
785    if_not_exists: bool,
786}
787
788impl SchemaBuilder {
789    /// 创建 Schema 构建器,指定表名
790    pub fn new(table_name: &str) -> Self {
791        Self {
792            table_name: table_name.to_string(),
793            columns: Vec::new(),
794            indexes: Vec::new(),
795            foreign_keys: Vec::new(),
796            if_not_exists: true,
797        }
798    }
799
800    /// 添加列定义
801    pub fn add_column(mut self, column: ColumnDef) -> Self {
802        self.columns.push(column);
803        self
804    }
805
806    /// 添加索引定义
807    pub fn add_index(mut self, index: IndexDef) -> Self {
808        self.indexes.push(index);
809        self
810    }
811
812    /// 添加外键定义
813    pub fn add_foreign_key(mut self, fk: ForeignKeyDef) -> Self {
814        self.foreign_keys.push(fk);
815        self
816    }
817
818    /// 设置是否添加 `IF NOT EXISTS`
819    pub fn if_not_exists(mut self, value: bool) -> Self {
820        self.if_not_exists = value;
821        self
822    }
823
824    /// 构建 CREATE TABLE SQL
825    pub fn build(&self, db_type: DbType) -> Result<String, DbError> {
826        let mut sql = String::new();
827        sql.push_str("CREATE TABLE ");
828        if self.if_not_exists {
829            sql.push_str("IF NOT EXISTS ");
830        }
831        sql.push_str(&self.table_name);
832        sql.push_str(" (");
833
834        let col_defs: Vec<String> = self.columns.iter().map(|c| c.build(db_type)).collect();
835        sql.push_str(&col_defs.join(", "));
836
837        for index in &self.indexes {
838            sql.push_str(", ");
839            sql.push_str(&index.build(db_type));
840        }
841
842        for fk in &self.foreign_keys {
843            sql.push_str(", ");
844            sql.push_str(&fk.build(db_type)?);
845        }
846
847        sql.push(')');
848        Ok(sql)
849    }
850}
851
852/// 列定义
853#[derive(Debug, Clone)]
854pub struct ColumnDef {
855    /// 列名
856    pub name: String,
857    /// 列类型
858    pub col_type: String,
859    /// 长度(如 VARCHAR(255))
860    pub length: Option<usize>,
861    /// 精度与小数位(如 DECIMAL(10,2))
862    pub precision: Option<(u32, u32)>,
863    /// 是否允许 NULL
864    pub nullable: bool,
865    /// 默认值表达式
866    pub default: Option<String>,
867    /// 是否自增
868    pub auto_increment: bool,
869    /// 是否唯一
870    pub unique: bool,
871    /// 列注释
872    pub comment: Option<String>,
873}
874
875impl ColumnDef {
876    /// 创建列定义,指定列名和类型
877    pub fn new(name: &str, col_type: &str) -> Self {
878        Self {
879            name: name.to_string(),
880            col_type: col_type.to_string(),
881            length: None,
882            precision: None,
883            nullable: true,
884            default: None,
885            auto_increment: false,
886            unique: false,
887            comment: None,
888        }
889    }
890
891    /// 设置 NOT NULL
892    pub fn not_null(mut self) -> Self {
893        self.nullable = false;
894        self
895    }
896
897    /// 设置默认值
898    pub fn default(mut self, value: &str) -> Self {
899        self.default = Some(value.to_string());
900        self
901    }
902
903    /// 设置自增
904    pub fn auto_increment(mut self) -> Self {
905        self.auto_increment = true;
906        self
907    }
908
909    /// 设置唯一约束
910    pub fn unique(mut self) -> Self {
911        self.unique = true;
912        self
913    }
914
915    /// 设置列注释
916    pub fn comment(mut self, comment: &str) -> Self {
917        self.comment = Some(comment.to_string());
918        self
919    }
920
921    /// 设置列长度
922    pub fn length(mut self, len: usize) -> Self {
923        self.length = Some(len);
924        self
925    }
926
927    fn build(&self, db_type: DbType) -> String {
928        let mut sql = format!("{} {}", self.name, self.col_type);
929        if let Some(len) = self.length {
930            if matches!(db_type, DbType::MySQL) {
931                sql.push_str(&format!("({})", len));
932            }
933        }
934        if self.auto_increment {
935            match db_type {
936                DbType::MySQL => sql.push_str(" AUTO_INCREMENT"),
937                DbType::PostgreSQL => sql.push_str(" GENERATED BY DEFAULT AS IDENTITY"),
938                DbType::Sqlite => sql.push_str(" AUTOINCREMENT"),
939                _ => {}
940            }
941        }
942        if !self.nullable {
943            sql.push_str(" NOT NULL");
944        }
945        if let Some(ref def) = self.default {
946            sql.push_str(&format!(" DEFAULT {}", def));
947        }
948        if self.unique {
949            sql.push_str(" UNIQUE");
950        }
951        sql
952    }
953}
954
955/// 索引定义
956#[derive(Debug, Clone)]
957pub struct IndexDef {
958    /// 索引名
959    pub name: String,
960    /// 索引列列表
961    pub columns: Vec<String>,
962    /// 是否唯一索引
963    pub unique: bool,
964    /// 索引类型(如 BTREE、HASH)
965    pub index_type: Option<String>,
966}
967
968impl IndexDef {
969    /// 创建索引定义,指定索引名和列列表
970    pub fn new(name: &str, columns: Vec<&str>) -> Self {
971        Self {
972            name: name.to_string(),
973            columns: columns.into_iter().map(|s| s.to_string()).collect(),
974            unique: false,
975            index_type: None,
976        }
977    }
978
979    /// 设置为唯一索引
980    pub fn unique(mut self) -> Self {
981        self.unique = true;
982        self
983    }
984
985    fn build(&self, _db_type: DbType) -> String {
986        let unique_str = if self.unique { "UNIQUE " } else { "" };
987        format!(
988            "{}KEY {} ({})",
989            unique_str,
990            self.name,
991            self.columns.join(", ")
992        )
993    }
994}
995
996/// 外键定义
997#[derive(Debug, Clone)]
998pub struct ForeignKeyDef {
999    /// 约束名
1000    pub name: String,
1001    /// 本表列名
1002    pub column: String,
1003    /// 引用表名
1004    pub referenced_table: String,
1005    /// 引用列名
1006    pub referenced_column: String,
1007    /// ON DELETE 动作
1008    pub on_delete: Option<String>,
1009    /// ON UPDATE 动作
1010    pub on_update: Option<String>,
1011}
1012
1013impl ForeignKeyDef {
1014    /// 创建外键定义,指定约束名、本表列、引用表和引用列
1015    pub fn new(name: &str, column: &str, referenced_table: &str, referenced_column: &str) -> Self {
1016        Self {
1017            name: name.to_string(),
1018            column: column.to_string(),
1019            referenced_table: referenced_table.to_string(),
1020            referenced_column: referenced_column.to_string(),
1021            on_delete: None,
1022            on_update: None,
1023        }
1024    }
1025
1026    /// 设置 ON DELETE 动作
1027    pub fn on_delete(mut self, action: &str) -> Self {
1028        self.on_delete = Some(action.to_string());
1029        self
1030    }
1031
1032    /// 设置 ON UPDATE 动作
1033    pub fn on_update(mut self, action: &str) -> Self {
1034        self.on_update = Some(action.to_string());
1035        self
1036    }
1037
1038    fn build(&self, _db_type: DbType) -> Result<String, DbError> {
1039        // v0.2.2 修复 C-3:FOREIGN KEY 标识符与 ON DELETE/ON UPDATE 动作严格校验
1040        crate::sql_safety::validate_identifier(&self.name, "foreign key constraint name")?;
1041        crate::sql_safety::validate_identifier(&self.column, "foreign key column")?;
1042        crate::sql_safety::validate_identifier(
1043            &self.referenced_table,
1044            "foreign key referenced table",
1045        )?;
1046        crate::sql_safety::validate_identifier(
1047            &self.referenced_column,
1048            "foreign key referenced column",
1049        )?;
1050        if let Some(ref on_delete) = self.on_delete {
1051            crate::sql_safety::validate_fk_action(on_delete)?;
1052        }
1053        if let Some(ref on_update) = self.on_update {
1054            crate::sql_safety::validate_fk_action(on_update)?;
1055        }
1056        let mut sql = format!(
1057            "CONSTRAINT {} FOREIGN KEY ({}) REFERENCES {} ({})",
1058            self.name, self.column, self.referenced_table, self.referenced_column
1059        );
1060        if let Some(ref on_delete) = self.on_delete {
1061            sql.push_str(&format!(" ON DELETE {}", on_delete.trim().to_uppercase()));
1062        }
1063        if let Some(ref on_update) = self.on_update {
1064            sql.push_str(&format!(" ON UPDATE {}", on_update.trim().to_uppercase()));
1065        }
1066        Ok(sql)
1067    }
1068}
1069
1070#[cfg(test)]
1071mod tests {
1072
1073    use super::*;
1074
1075    #[test]
1076    fn test_migration_new() {
1077        let m = Migration::new("001", "create_users", "CREATE TABLE...", "DROP TABLE...");
1078        assert_eq!(m.version, "001");
1079        assert_eq!(m.name, "create_users");
1080    }
1081
1082    #[test]
1083    fn test_migration_with_batch() {
1084        let m = Migration::new("001", "create_users", "UP", "DOWN").with_batch(1);
1085        assert_eq!(m.batch, 1);
1086    }
1087
1088    #[test]
1089    fn test_migrator_latest_version() {
1090        let ctx = MigrationContext::default();
1091        let migrator = Migrator::new(ctx)
1092            .add_migration(Migration::new("001", "v1", "UP", "DOWN"))
1093            .add_migration(Migration::new("002", "v2", "UP", "DOWN"));
1094
1095        assert_eq!(migrator.latest_version(), Some("002"));
1096    }
1097
1098    #[test]
1099    fn test_migrator_find_migration() {
1100        let ctx = MigrationContext::default();
1101        let migrator =
1102            Migrator::new(ctx).add_migration(Migration::new("001", "create_users", "UP", "DOWN"));
1103
1104        assert!(migrator.find_migration("001").is_some());
1105        assert!(migrator.find_migration("999").is_none());
1106    }
1107
1108    #[test]
1109    fn test_column_def() {
1110        let col = ColumnDef::new("id", "INT").not_null().auto_increment();
1111        assert_eq!(col.name, "id");
1112        assert!(!col.nullable);
1113        assert!(col.auto_increment);
1114    }
1115
1116    #[test]
1117    fn test_column_build_mysql() {
1118        let col = ColumnDef::new("id", "INT").not_null();
1119        let sql = col.build(DbType::MySQL);
1120        assert!(sql.contains("NOT NULL"));
1121    }
1122
1123    #[test]
1124    fn test_index_build() {
1125        let idx = IndexDef::new("idx_name", vec!["name"]).unique();
1126        let sql = idx.build(DbType::MySQL);
1127        assert!(sql.contains("UNIQUE KEY"));
1128    }
1129
1130    #[test]
1131    fn test_foreign_key_build() {
1132        let fk = ForeignKeyDef::new("fk_user", "user_id", "users", "id").on_delete("CASCADE");
1133        let sql = fk.build(DbType::MySQL).unwrap();
1134        assert!(sql.contains("FOREIGN KEY"));
1135        assert!(sql.contains("ON DELETE CASCADE"));
1136    }
1137
1138    #[test]
1139    fn test_foreign_key_build_normalizes_action_case() {
1140        // v0.2.2 修复 C-3:动作大小写不敏感,输出统一为大写
1141        let fk = ForeignKeyDef::new("fk_user", "user_id", "users", "id").on_delete("cascade");
1142        let sql = fk.build(DbType::MySQL).unwrap();
1143        assert!(sql.contains("ON DELETE CASCADE"));
1144    }
1145
1146    #[test]
1147    fn test_foreign_key_rejects_sql_injection_in_column() {
1148        let fk = ForeignKeyDef::new("fk_user", "user_id; DROP TABLE users", "users", "id");
1149        let result = fk.build(DbType::MySQL);
1150        assert!(result.is_err());
1151    }
1152
1153    #[test]
1154    fn test_foreign_key_rejects_sql_injection_in_ref_table() {
1155        let fk = ForeignKeyDef::new("fk_user", "user_id", "users; DROP TABLE users", "id");
1156        let result = fk.build(DbType::MySQL);
1157        assert!(result.is_err());
1158    }
1159
1160    #[test]
1161    fn test_foreign_key_rejects_sql_injection_in_on_delete() {
1162        let fk = ForeignKeyDef::new("fk_user", "user_id", "users", "id")
1163            .on_delete("CASCADE; DROP TABLE users");
1164        let result = fk.build(DbType::MySQL);
1165        assert!(result.is_err());
1166    }
1167
1168    #[test]
1169    fn test_foreign_key_rejects_invalid_on_update_action() {
1170        let fk = ForeignKeyDef::new("fk_user", "user_id", "users", "id").on_update("EVIL_ACTION");
1171        let result = fk.build(DbType::MySQL);
1172        assert!(result.is_err());
1173    }
1174
1175    #[test]
1176    fn test_schema_builder() {
1177        let schema = SchemaBuilder::new("users")
1178            .add_column(ColumnDef::new("id", "INT").not_null().auto_increment())
1179            .add_column(ColumnDef::new("name", "VARCHAR").length(255));
1180
1181        let sql = schema.build(DbType::MySQL).unwrap();
1182        assert!(sql.contains("CREATE TABLE"));
1183        assert!(sql.contains("users"));
1184    }
1185
1186    #[test]
1187    fn test_migration_progress() {
1188        let progress = MigrationProgress::new(10, 4);
1189        assert_eq!(progress.pending, 6);
1190        assert!((progress.percent_complete() - 40.0).abs() < 0.01);
1191    }
1192
1193    #[test]
1194    fn test_migrate_macro() {
1195        let m = crate::migrate!(
1196            "001",
1197            "create_users",
1198            "CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT)",
1199            "DROP TABLE users"
1200        );
1201        assert_eq!(m.version, "001");
1202        assert_eq!(m.name, "create_users");
1203        assert!(m.sql_up.contains("CREATE TABLE"));
1204        assert!(m.sql_down.contains("DROP TABLE"));
1205    }
1206}