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