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