1use 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 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; 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 sql_files.sort();
97
98 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 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 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 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
165fn 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 (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}
181
182impl Default for MigrationContext {
183 fn default() -> Self {
184 Self {
185 table_name: "__migrations".to_string(),
186 connection: None,
187 }
188 }
189}
190
191#[derive(Debug, Clone, PartialEq)]
192pub enum MigrationDirection {
193 Up,
194 Down,
195}
196
197pub struct Migrator {
198 context: MigrationContext,
199 migrations: Vec<Migration>,
200}
201
202impl Migrator {
203 pub fn new(context: MigrationContext) -> Self {
204 Self {
205 context,
206 migrations: Vec::new(),
207 }
208 }
209
210 pub fn add_migration(mut self, migration: Migration) -> Self {
211 self.migrations.push(migration);
212 self
213 }
214
215 pub fn add_migrations(mut self, migrations: Vec<Migration>) -> Self {
216 self.migrations.extend(migrations);
217 self
218 }
219
220 pub fn get_migrations(&self) -> &Vec<Migration> {
221 &self.migrations
222 }
223
224 pub fn get_pending_migrations(&self) -> Vec<&Migration> {
225 self.migrations.iter().filter(|m| m.batch == 0).collect()
226 }
227
228 pub fn get_applied_migrations(&self) -> Vec<&Migration> {
229 self.migrations.iter().filter(|m| m.batch > 0).collect()
230 }
231
232 pub fn latest_version(&self) -> Option<&str> {
233 self.migrations.last().map(|m| m.version.as_str())
234 }
235
236 pub fn find_migration(&self, version: &str) -> Option<&Migration> {
237 self.migrations.iter().find(|m| m.version == version)
238 }
239
240 pub async fn migrate(&mut self) -> Result<Vec<String>, DbError> {
242 let mut applied = Vec::new();
243 let current_batch = self.migrations.iter().map(|m| m.batch).max().unwrap_or(0) + 1;
244
245 let pending_indices: Vec<usize> = self
247 .migrations
248 .iter()
249 .enumerate()
250 .filter(|(_, m)| m.batch == 0)
251 .map(|(idx, _)| idx)
252 .collect();
253
254 for migration_idx in pending_indices {
255 let sql_up = self.migrations[migration_idx].sql_up.clone();
256
257 if let Some(ref mut conn) = self.context.connection {
259 if !sql_up.is_empty() {
260 conn.execute(&sql_up).await?;
261 }
262 }
263
264 let now = chrono::Utc::now();
266 self.migrations[migration_idx].batch = current_batch;
267 self.migrations[migration_idx].executed_at = Some(now);
268
269 applied.push(self.migrations[migration_idx].version.clone());
270 }
271
272 Ok(applied)
273 }
274
275 pub async fn rollback(&mut self, version: &str) -> Result<(), DbError> {
277 let migration_idx = self
278 .migrations
279 .iter()
280 .position(|m| m.version == version)
281 .ok_or_else(|| DbError::MigrationError(format!("Migration {} not found", version)))?;
282
283 if self.migrations[migration_idx].batch == 0 {
284 return Err(DbError::MigrationError(format!(
285 "Migration {} not applied",
286 version
287 )));
288 }
289
290 let sql_down = self.migrations[migration_idx].sql_down.clone();
291
292 if let Some(ref mut conn) = self.context.connection {
293 if !sql_down.is_empty() {
294 conn.execute(&sql_down).await?;
295 }
296 }
297
298 self.migrations[migration_idx].batch = 0;
299 self.migrations[migration_idx].executed_at = None;
300 Ok(())
301 }
302
303 pub async fn up(&mut self, target_version: Option<&str>) -> Result<Vec<String>, DbError> {
305 let mut applied = Vec::new();
306 let current_batch = self.migrations.iter().map(|m| m.batch).max().unwrap_or(0) + 1;
307
308 for migration in &mut self.migrations {
309 if migration.batch > 0 {
310 continue; }
312
313 if let Some(target) = target_version {
314 if migration.version.as_str() > target {
315 break; }
317 }
318
319 let sql_up = migration.sql_up.clone();
320 if let Some(ref mut conn) = self.context.connection {
321 if !sql_up.is_empty() {
322 conn.execute(&sql_up).await?;
323 }
324 }
325
326 migration.batch = current_batch;
327 migration.executed_at = Some(chrono::Utc::now());
328 applied.push(migration.version.clone());
329 }
330
331 Ok(applied)
332 }
333
334 pub async fn down(&mut self, target_version: Option<&str>) -> Result<Vec<String>, DbError> {
336 let mut rolled_back = Vec::new();
337
338 let mut indices: Vec<usize> = (0..self.migrations.len()).collect();
340 indices.reverse();
341
342 for idx in indices {
343 let migration = &mut self.migrations[idx];
344 if migration.batch == 0 {
345 continue; }
347
348 if let Some(target) = target_version {
349 if migration.version.as_str() <= target {
350 break; }
352 }
353
354 let sql_down = migration.sql_down.clone();
355 if let Some(ref mut conn) = self.context.connection {
356 if !sql_down.is_empty() {
357 conn.execute(&sql_down).await?;
358 }
359 }
360
361 migration.batch = 0;
362 migration.executed_at = None;
363 rolled_back.push(migration.version.clone());
364 }
365
366 Ok(rolled_back)
367 }
368
369 pub async fn reset(&mut self) -> Result<Vec<String>, DbError> {
371 self.down(None).await?;
373 self.migrate().await
375 }
376
377 pub async fn refresh(&mut self) -> Result<Vec<String>, DbError> {
379 self.reset().await
380 }
381
382 pub fn progress(&self) -> MigrationProgress {
384 let total = self.migrations.len();
385 let applied = self.migrations.iter().filter(|m| m.batch > 0).count();
386 MigrationProgress::new(total, applied)
387 }
388}
389
390#[derive(Debug, Clone)]
391pub struct MigrationProgress {
392 pub total: usize,
393 pub applied: usize,
394 pub pending: usize,
395 pub current_batch: i32,
396}
397
398impl MigrationProgress {
399 pub fn new(total: usize, applied: usize) -> Self {
400 Self {
401 total,
402 applied,
403 pending: total - applied,
404 current_batch: 0,
405 }
406 }
407
408 pub fn percent_complete(&self) -> f64 {
409 if self.total == 0 {
410 return 100.0;
411 }
412 (self.applied as f64 / self.total as f64) * 100.0
413 }
414}
415
416pub struct SchemaBuilder {
417 table_name: String,
418 columns: Vec<ColumnDef>,
419 indexes: Vec<IndexDef>,
420 foreign_keys: Vec<ForeignKeyDef>,
421 if_not_exists: bool,
422}
423
424impl SchemaBuilder {
425 pub fn new(table_name: &str) -> Self {
426 Self {
427 table_name: table_name.to_string(),
428 columns: Vec::new(),
429 indexes: Vec::new(),
430 foreign_keys: Vec::new(),
431 if_not_exists: true,
432 }
433 }
434
435 pub fn add_column(mut self, column: ColumnDef) -> Self {
436 self.columns.push(column);
437 self
438 }
439
440 pub fn add_index(mut self, index: IndexDef) -> Self {
441 self.indexes.push(index);
442 self
443 }
444
445 pub fn add_foreign_key(mut self, fk: ForeignKeyDef) -> Self {
446 self.foreign_keys.push(fk);
447 self
448 }
449
450 pub fn if_not_exists(mut self, value: bool) -> Self {
451 self.if_not_exists = value;
452 self
453 }
454
455 pub fn build(&self, db_type: DbType) -> String {
456 let mut sql = String::new();
457 sql.push_str("CREATE TABLE ");
458 if self.if_not_exists {
459 sql.push_str("IF NOT EXISTS ");
460 }
461 sql.push_str(&self.table_name);
462 sql.push_str(" (");
463
464 let col_defs: Vec<String> = self.columns.iter().map(|c| c.build(db_type)).collect();
465 sql.push_str(&col_defs.join(", "));
466
467 for index in &self.indexes {
468 sql.push_str(", ");
469 sql.push_str(&index.build(db_type));
470 }
471
472 for fk in &self.foreign_keys {
473 sql.push_str(", ");
474 sql.push_str(&fk.build(db_type));
475 }
476
477 sql.push(')');
478 sql
479 }
480}
481
482#[derive(Debug, Clone)]
483pub struct ColumnDef {
484 pub name: String,
485 pub col_type: String,
486 pub length: Option<usize>,
487 pub precision: Option<(u32, u32)>,
488 pub nullable: bool,
489 pub default: Option<String>,
490 pub auto_increment: bool,
491 pub unique: bool,
492 pub comment: Option<String>,
493}
494
495impl ColumnDef {
496 pub fn new(name: &str, col_type: &str) -> Self {
497 Self {
498 name: name.to_string(),
499 col_type: col_type.to_string(),
500 length: None,
501 precision: None,
502 nullable: true,
503 default: None,
504 auto_increment: false,
505 unique: false,
506 comment: None,
507 }
508 }
509
510 pub fn not_null(mut self) -> Self {
511 self.nullable = false;
512 self
513 }
514
515 pub fn default(mut self, value: &str) -> Self {
516 self.default = Some(value.to_string());
517 self
518 }
519
520 pub fn auto_increment(mut self) -> Self {
521 self.auto_increment = true;
522 self
523 }
524
525 pub fn unique(mut self) -> Self {
526 self.unique = true;
527 self
528 }
529
530 pub fn comment(mut self, comment: &str) -> Self {
531 self.comment = Some(comment.to_string());
532 self
533 }
534
535 pub fn length(mut self, len: usize) -> Self {
536 self.length = Some(len);
537 self
538 }
539
540 fn build(&self, db_type: DbType) -> String {
541 let mut sql = format!("{} {}", self.name, self.col_type);
542 if let Some(len) = self.length {
543 if matches!(db_type, DbType::MySQL) {
544 sql.push_str(&format!("({})", len));
545 }
546 }
547 if self.auto_increment {
548 match db_type {
549 DbType::MySQL => sql.push_str(" AUTO_INCREMENT"),
550 DbType::PostgreSQL => sql.push_str(" GENERATED BY DEFAULT AS IDENTITY"),
551 DbType::Sqlite => sql.push_str(" AUTOINCREMENT"),
552 _ => {}
553 }
554 }
555 if !self.nullable {
556 sql.push_str(" NOT NULL");
557 }
558 if let Some(ref def) = self.default {
559 sql.push_str(&format!(" DEFAULT {}", def));
560 }
561 if self.unique {
562 sql.push_str(" UNIQUE");
563 }
564 sql
565 }
566}
567
568#[derive(Debug, Clone)]
569pub struct IndexDef {
570 pub name: String,
571 pub columns: Vec<String>,
572 pub unique: bool,
573 pub index_type: Option<String>,
574}
575
576impl IndexDef {
577 pub fn new(name: &str, columns: Vec<&str>) -> Self {
578 Self {
579 name: name.to_string(),
580 columns: columns.into_iter().map(|s| s.to_string()).collect(),
581 unique: false,
582 index_type: None,
583 }
584 }
585
586 pub fn unique(mut self) -> Self {
587 self.unique = true;
588 self
589 }
590
591 fn build(&self, _db_type: DbType) -> String {
592 let unique_str = if self.unique { "UNIQUE " } else { "" };
593 format!(
594 "{}KEY {} ({})",
595 unique_str,
596 self.name,
597 self.columns.join(", ")
598 )
599 }
600}
601
602#[derive(Debug, Clone)]
603pub struct ForeignKeyDef {
604 pub name: String,
605 pub column: String,
606 pub referenced_table: String,
607 pub referenced_column: String,
608 pub on_delete: Option<String>,
609 pub on_update: Option<String>,
610}
611
612impl ForeignKeyDef {
613 pub fn new(name: &str, column: &str, referenced_table: &str, referenced_column: &str) -> Self {
614 Self {
615 name: name.to_string(),
616 column: column.to_string(),
617 referenced_table: referenced_table.to_string(),
618 referenced_column: referenced_column.to_string(),
619 on_delete: None,
620 on_update: None,
621 }
622 }
623
624 pub fn on_delete(mut self, action: &str) -> Self {
625 self.on_delete = Some(action.to_string());
626 self
627 }
628
629 pub fn on_update(mut self, action: &str) -> Self {
630 self.on_update = Some(action.to_string());
631 self
632 }
633
634 fn build(&self, _db_type: DbType) -> String {
635 crate::sql_safety::validate_identifier(&self.name, "foreign key constraint name")
637 .expect("invalid foreign key constraint name");
638 crate::sql_safety::validate_identifier(&self.column, "foreign key column")
639 .expect("invalid foreign key column name");
640 crate::sql_safety::validate_identifier(
641 &self.referenced_table,
642 "foreign key referenced table",
643 )
644 .expect("invalid foreign key referenced table name");
645 crate::sql_safety::validate_identifier(
646 &self.referenced_column,
647 "foreign key referenced column",
648 )
649 .expect("invalid foreign key referenced column name");
650 if let Some(ref on_delete) = self.on_delete {
651 crate::sql_safety::validate_fk_action(on_delete).expect("invalid ON DELETE action");
652 }
653 if let Some(ref on_update) = self.on_update {
654 crate::sql_safety::validate_fk_action(on_update).expect("invalid ON UPDATE action");
655 }
656 let mut sql = format!(
657 "CONSTRAINT {} FOREIGN KEY ({}) REFERENCES {} ({})",
658 self.name, self.column, self.referenced_table, self.referenced_column
659 );
660 if let Some(ref on_delete) = self.on_delete {
661 sql.push_str(&format!(" ON DELETE {}", on_delete.trim().to_uppercase()));
662 }
663 if let Some(ref on_update) = self.on_update {
664 sql.push_str(&format!(" ON UPDATE {}", on_update.trim().to_uppercase()));
665 }
666 sql
667 }
668}
669
670#[cfg(test)]
671mod tests {
672 use super::*;
673
674 #[test]
675 fn test_migration_new() {
676 let m = Migration::new("001", "create_users", "CREATE TABLE...", "DROP TABLE...");
677 assert_eq!(m.version, "001");
678 assert_eq!(m.name, "create_users");
679 }
680
681 #[test]
682 fn test_migration_with_batch() {
683 let m = Migration::new("001", "create_users", "UP", "DOWN").with_batch(1);
684 assert_eq!(m.batch, 1);
685 }
686
687 #[test]
688 fn test_migrator_latest_version() {
689 let ctx = MigrationContext::default();
690 let migrator = Migrator::new(ctx)
691 .add_migration(Migration::new("001", "v1", "UP", "DOWN"))
692 .add_migration(Migration::new("002", "v2", "UP", "DOWN"));
693
694 assert_eq!(migrator.latest_version(), Some("002"));
695 }
696
697 #[test]
698 fn test_migrator_find_migration() {
699 let ctx = MigrationContext::default();
700 let migrator =
701 Migrator::new(ctx).add_migration(Migration::new("001", "create_users", "UP", "DOWN"));
702
703 assert!(migrator.find_migration("001").is_some());
704 assert!(migrator.find_migration("999").is_none());
705 }
706
707 #[test]
708 fn test_column_def() {
709 let col = ColumnDef::new("id", "INT").not_null().auto_increment();
710 assert_eq!(col.name, "id");
711 assert!(!col.nullable);
712 assert!(col.auto_increment);
713 }
714
715 #[test]
716 fn test_column_build_mysql() {
717 let col = ColumnDef::new("id", "INT").not_null();
718 let sql = col.build(DbType::MySQL);
719 assert!(sql.contains("NOT NULL"));
720 }
721
722 #[test]
723 fn test_index_build() {
724 let idx = IndexDef::new("idx_name", vec!["name"]).unique();
725 let sql = idx.build(DbType::MySQL);
726 assert!(sql.contains("UNIQUE KEY"));
727 }
728
729 #[test]
730 fn test_foreign_key_build() {
731 let fk = ForeignKeyDef::new("fk_user", "user_id", "users", "id").on_delete("CASCADE");
732 let sql = fk.build(DbType::MySQL);
733 assert!(sql.contains("FOREIGN KEY"));
734 assert!(sql.contains("ON DELETE CASCADE"));
735 }
736
737 #[test]
738 fn test_foreign_key_build_normalizes_action_case() {
739 let fk = ForeignKeyDef::new("fk_user", "user_id", "users", "id").on_delete("cascade");
741 let sql = fk.build(DbType::MySQL);
742 assert!(sql.contains("ON DELETE CASCADE"));
743 }
744
745 #[test]
746 #[should_panic(expected = "invalid foreign key column name")]
747 fn test_foreign_key_rejects_sql_injection_in_column() {
748 let fk = ForeignKeyDef::new("fk_user", "user_id; DROP TABLE users", "users", "id");
749 let _ = fk.build(DbType::MySQL);
750 }
751
752 #[test]
753 #[should_panic(expected = "invalid foreign key referenced table name")]
754 fn test_foreign_key_rejects_sql_injection_in_ref_table() {
755 let fk = ForeignKeyDef::new("fk_user", "user_id", "users; DROP TABLE users", "id");
756 let _ = fk.build(DbType::MySQL);
757 }
758
759 #[test]
760 #[should_panic(expected = "invalid ON DELETE action")]
761 fn test_foreign_key_rejects_sql_injection_in_on_delete() {
762 let fk = ForeignKeyDef::new("fk_user", "user_id", "users", "id")
763 .on_delete("CASCADE; DROP TABLE users");
764 let _ = fk.build(DbType::MySQL);
765 }
766
767 #[test]
768 #[should_panic(expected = "invalid ON UPDATE action")]
769 fn test_foreign_key_rejects_invalid_on_update_action() {
770 let fk = ForeignKeyDef::new("fk_user", "user_id", "users", "id").on_update("EVIL_ACTION");
771 let _ = fk.build(DbType::MySQL);
772 }
773
774 #[test]
775 fn test_schema_builder() {
776 let schema = SchemaBuilder::new("users")
777 .add_column(ColumnDef::new("id", "INT").not_null().auto_increment())
778 .add_column(ColumnDef::new("name", "VARCHAR").length(255));
779
780 let sql = schema.build(DbType::MySQL);
781 assert!(sql.contains("CREATE TABLE"));
782 assert!(sql.contains("users"));
783 }
784
785 #[test]
786 fn test_migration_progress() {
787 let progress = MigrationProgress::new(10, 4);
788 assert_eq!(progress.pending, 6);
789 assert!((progress.percent_complete() - 40.0).abs() < 0.01);
790 }
791}