1use std::path::{Path, PathBuf};
33
34use clap::Args;
35use sz_rust_core::orm::migration::{
36 FileMigrationResolver, Migration, MigrationContext, MigrationResolver, Migrator,
37};
38use sz_rust_core::orm::{Connection, ConnectionFactory, DbType};
39
40use crate::error::CliError;
41
42#[derive(Args, Debug)]
46pub struct MigrateArgs {
47 #[arg(long)]
49 pub rollback: bool,
50
51 #[arg(short = 'p', long, default_value = "migrations")]
53 pub path: String,
54
55 #[arg(long, default_value = "postgres")]
61 pub db_type: String,
62
63 #[arg(long)]
65 pub show_sql: bool,
66
67 #[arg(long)]
75 pub url: Option<String>,
76}
77
78pub fn execute_migrate(args: &MigrateArgs) -> Result<(), CliError> {
88 let path = PathBuf::from(&args.path);
89
90 if !path.exists() {
91 return Err(CliError::Migration(format!(
92 "Migration directory not found: {}",
93 path.display()
94 )));
95 }
96
97 let db_type = DbType::from_str(&args.db_type)
98 .ok_or_else(|| CliError::Migration(format!("Unknown database type: {}", args.db_type)))?;
99
100 let migrations = resolve_migrations(&path, db_type)?;
101
102 if migrations.is_empty() {
103 println!("No migrations found in: {}", path.display());
104 return Ok(());
105 }
106
107 match &args.url {
108 None => execute_migrate_offline(args, &migrations),
109 Some(url) => execute_migrate_online(args, &migrations, url, db_type),
110 }
111}
112
113fn execute_migrate_offline(args: &MigrateArgs, migrations: &[Migration]) -> Result<(), CliError> {
115 if args.rollback {
116 println!("Rolling back last batch in: {}", args.path);
117 if let Some(last) = migrations.last() {
118 println!(" Would rollback: {} ({})", last.version, last.name);
119 if args.show_sql {
120 print_sql_block("SQL DOWN", &last.sql_down);
121 }
122 }
123 println!("Note: Actual rollback requires database connection (offline mode).");
124 } else {
125 println!("Running migrations in: {}", args.path);
126 for m in migrations {
127 println!(" Would apply: {} ({})", m.version, m.name);
128 if args.show_sql {
129 print_sql_block("SQL UP", &m.sql_up);
130 }
131 }
132 println!(
133 "Total: {} migration(s). Note: Actual execution requires database connection (offline mode).",
134 migrations.len()
135 );
136 }
137 Ok(())
138}
139
140fn execute_migrate_online(
142 args: &MigrateArgs,
143 migrations: &[Migration],
144 url: &str,
145 db_type: DbType,
146) -> Result<(), CliError> {
147 let rt = tokio::runtime::Builder::new_current_thread()
149 .enable_all()
150 .build()
151 .map_err(|e| CliError::Migration(format!("Failed to create tokio runtime: {}", e)))?;
152
153 rt.block_on(async move {
154 let mut conn = create_connection(url, db_type).await?;
155
156 if args.rollback {
157 let last = migrations
159 .last()
160 .ok_or_else(|| CliError::Migration("No migrations to rollback".to_string()))?;
161 println!("Rolling back: {} ({})", last.version, last.name);
162 if args.show_sql {
163 print_sql_block("SQL DOWN", &last.sql_down);
164 }
165 if !last.sql_down.is_empty() {
166 conn.execute(&last.sql_down)
167 .await
168 .map_err(|e| CliError::Migration(format!("Rollback failed: {}", e)))?;
169 }
170 delete_migration_record(&mut conn, &last.version, db_type).await?;
172 println!("Rollback completed: {} ({})", last.version, last.name);
173 } else {
174 ensure_migrations_table(&mut conn, db_type).await?;
176
177 let applied = fetch_applied_versions(&mut conn, db_type).await?;
179
180 let pending: Vec<&Migration> = migrations
182 .iter()
183 .filter(|m| !applied.contains(&m.version))
184 .collect();
185
186 if pending.is_empty() {
187 println!("No pending migrations. Database is up to date.");
188 return Ok(());
189 }
190
191 println!("Running {} pending migration(s):", pending.len());
192
193 let mut context = MigrationContext::default().with_db_type(db_type);
195 context.connection = Some(conn);
196
197 let mut migrator = Migrator::new(context);
198 for m in migrations {
199 let rebuilt = Migration::new(&m.version, &m.name, &m.sql_up, &m.sql_down);
201 if applied.contains(&m.version) {
202 migrator = migrator.add_migration(rebuilt.with_batch(1));
204 } else {
205 migrator = migrator.add_migration(rebuilt);
206 }
207 }
208
209 let applied_versions = migrator
210 .migrate()
211 .await
212 .map_err(|e| CliError::Migration(format!("Migration failed: {}", e)))?;
213
214 for v in &applied_versions {
215 println!(" Applied: {}", v);
216 }
217 println!("Migration completed: {} applied.", applied_versions.len());
218 }
219
220 Ok::<(), CliError>(())
221 })
222}
223
224pub fn execute_status(path: &str) -> Result<(), CliError> {
230 execute_status_full(path, "postgres", false, None)
231}
232
233pub fn execute_status_with(path: &str, db_type_str: &str, show_sql: bool) -> Result<(), CliError> {
241 execute_status_full(path, db_type_str, show_sql, None)
242}
243
244pub fn execute_status_full(
246 path: &str,
247 db_type_str: &str,
248 show_sql: bool,
249 url: Option<&str>,
250) -> Result<(), CliError> {
251 let path_buf = PathBuf::from(path);
252
253 if !path_buf.exists() {
254 return Err(CliError::Migration(format!(
255 "Migration directory not found: {}",
256 path_buf.display()
257 )));
258 }
259
260 let db_type = DbType::from_str(db_type_str)
261 .ok_or_else(|| CliError::Migration(format!("Unknown database type: {}", db_type_str)))?;
262
263 let migrations = resolve_migrations(&path_buf, db_type)?;
264
265 if migrations.is_empty() {
266 println!("No migrations found in: {}", path_buf.display());
267 return Ok(());
268 }
269
270 let applied_versions = if let Some(url) = url {
272 let rt = tokio::runtime::Builder::new_current_thread()
273 .enable_all()
274 .build()
275 .map_err(|e| CliError::Migration(format!("Failed to create tokio runtime: {}", e)))?;
276 rt.block_on(async move {
277 let mut conn = create_connection(url, db_type).await?;
278 ensure_migrations_table(&mut conn, db_type).await?;
279 fetch_applied_versions(&mut conn, db_type).await
280 })?
281 } else {
282 std::collections::HashSet::new()
283 };
284
285 println!(
287 "{:<15} {:<30} {:<20}",
288 "Version", "Migration Name", "Status"
289 );
290 println!("{}", "-".repeat(65));
291
292 for m in &migrations {
293 let status = if applied_versions.contains(&m.version) {
294 "Applied"
295 } else if url.is_some() {
296 "Pending"
297 } else {
298 "Pending*"
299 };
300 println!("{:<15} {:<30} {:<20}", m.version, m.name, status);
301 if show_sql {
302 print_sql_block("SQL UP", &m.sql_up);
303 print_sql_block("SQL DOWN", &m.sql_down);
304 }
305 }
306
307 println!();
308 if url.is_some() {
309 let applied = migrations
310 .iter()
311 .filter(|m| applied_versions.contains(&m.version))
312 .count();
313 println!(
314 "Total: {} migration(s), {} applied, {} pending.",
315 migrations.len(),
316 applied,
317 migrations.len() - applied
318 );
319 } else {
320 println!("* Status cannot be determined without database connection (offline mode).");
321 }
322
323 Ok(())
324}
325
326async fn create_connection(url: &str, db_type: DbType) -> Result<Box<dyn Connection>, CliError> {
328 use std::sync::Arc;
329 use sz_orm_sqlx::{
330 MySqlPoolHandle, PgPoolHandle, SqlitePoolHandle, SqlxMySqlConnectionFactory,
331 SqlxPgConnectionFactory, SqlxSqliteConnectionFactory,
332 };
333
334 match db_type {
335 DbType::PostgreSQL => {
336 let pool = PgPoolHandle::connect(url).await.map_err(|e| {
337 CliError::Migration(format!("PostgreSQL connect failed: {}", e))
338 })?;
339 let factory = SqlxPgConnectionFactory::new(Arc::new(pool));
340 let conn = factory.create().await.map_err(|e| {
341 CliError::Migration(format!("PostgreSQL acquire failed: {}", e))
342 })?;
343 Ok(conn)
344 }
345 DbType::MySQL => {
346 let pool = MySqlPoolHandle::connect(url).await.map_err(|e| {
347 CliError::Migration(format!("MySQL connect failed: {}", e))
348 })?;
349 let factory = SqlxMySqlConnectionFactory::new(Arc::new(pool));
350 let conn = factory.create().await.map_err(|e| {
351 CliError::Migration(format!("MySQL acquire failed: {}", e))
352 })?;
353 Ok(conn)
354 }
355 DbType::Sqlite => {
356 let pool = SqlitePoolHandle::connect(url).await.map_err(|e| {
357 CliError::Migration(format!("SQLite connect failed: {}", e))
358 })?;
359 let factory = SqlxSqliteConnectionFactory::new(Arc::new(pool));
360 let conn = factory.create().await.map_err(|e| {
361 CliError::Migration(format!("SQLite acquire failed: {}", e))
362 })?;
363 Ok(conn)
364 }
365 _ => Err(CliError::Migration(format!(
366 "Online migration not supported for db_type {:?}. Supported: PostgreSQL, MySQL, SQLite.",
367 db_type
368 ))),
369 }
370}
371
372async fn ensure_migrations_table(
374 conn: &mut Box<dyn Connection>,
375 db_type: DbType,
376) -> Result<(), CliError> {
377 let sql = match db_type {
378 DbType::PostgreSQL | DbType::Sqlite => {
379 "CREATE TABLE IF NOT EXISTS __migrations (
380 version VARCHAR(255) PRIMARY KEY,
381 name VARCHAR(255) NOT NULL,
382 batch INTEGER NOT NULL,
383 run_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
384 )"
385 }
386 DbType::MySQL => {
387 "CREATE TABLE IF NOT EXISTS __migrations (
388 version VARCHAR(255) PRIMARY KEY,
389 name VARCHAR(255) NOT NULL,
390 batch INT NOT NULL,
391 run_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
392 )"
393 }
394 _ => {
395 return Err(CliError::Migration(format!(
396 "Cannot ensure __migrations table for db_type {:?}",
397 db_type
398 )))
399 }
400 };
401 conn.execute(sql)
402 .await
403 .map_err(|e| CliError::Migration(format!("Failed to create __migrations table: {}", e)))?;
404 Ok(())
405}
406
407async fn fetch_applied_versions(
412 conn: &mut Box<dyn Connection>,
413 _db_type: DbType,
414) -> Result<std::collections::HashSet<String>, CliError> {
415 let rows = conn
418 .query("SELECT version FROM __migrations")
419 .await
420 .map_err(|e| CliError::Migration(format!("Failed to query __migrations: {}", e)))?;
421
422 let mut versions = std::collections::HashSet::new();
423 for row in &rows {
424 if let Some(val) = row.get("version") {
426 use sz_rust_core::orm::Value;
427 match val {
428 Value::String(s) => versions.insert(s.clone()),
429 Value::I64(i) => versions.insert(i.to_string()),
430 Value::I32(i) => versions.insert(i.to_string()),
431 _ => false,
432 };
433 }
434 }
435 Ok(versions)
436}
437
438async fn delete_migration_record(
440 conn: &mut Box<dyn Connection>,
441 version: &str,
442 db_type: DbType,
443) -> Result<(), CliError> {
444 let sql = match db_type {
445 DbType::PostgreSQL | DbType::Sqlite => {
446 format!("DELETE FROM __migrations WHERE version = '{}'", version)
447 }
448 DbType::MySQL => {
449 format!("DELETE FROM __migrations WHERE version = '{}'", version)
450 }
451 _ => return Ok(()),
452 };
453 conn.execute(&sql)
454 .await
455 .map_err(|e| CliError::Migration(format!("Failed to delete migration record: {}", e)))?;
456 Ok(())
457}
458
459fn resolve_migrations(path: &Path, db_type: DbType) -> Result<Vec<Migration>, CliError> {
467 let resolver = FileMigrationResolver::new(path.to_path_buf());
468 resolver
469 .resolve(db_type)
470 .map_err(|e| CliError::Migration(format!("Failed to resolve migrations: {}", e)))
471}
472
473fn print_sql_block(title: &str, sql: &str) {
482 if sql.is_empty() {
483 return;
484 }
485 println!(" --- {} ---", title);
486 for line in sql.lines() {
487 println!(" {}", line);
488 }
489 println!(" {}", "-".repeat(title.len() + 8));
490}
491
492#[cfg(test)]
493mod tests {
494 use super::*;
495 use std::fs;
496 use std::io::Write;
497
498 fn create_test_migration(dir: &Path, version: &str, name: &str) {
500 let up_name = format!("{}_{}_up.sql", version, name);
501 let down_name = format!("{}_{}_down.sql", version, name);
502
503 let up_path = dir.join(up_name);
504 let down_path = dir.join(down_name);
505
506 let mut up_file = fs::File::create(&up_path).unwrap();
507 writeln!(up_file, "-- {} up", name).unwrap();
508
509 let mut down_file = fs::File::create(&down_path).unwrap();
510 writeln!(down_file, "-- {} down", name).unwrap();
511 }
512
513 #[test]
514 fn test_resolve_migrations_empty() {
515 let temp = tempfile::tempdir().unwrap();
516 let path = temp.path().to_path_buf();
517 let result = resolve_migrations(&path, DbType::PostgreSQL).unwrap();
518 assert!(result.is_empty());
519 }
520
521 #[test]
522 fn test_resolve_migrations_with_files() {
523 let temp = tempfile::tempdir().unwrap();
524 let path = temp.path().to_path_buf();
525
526 create_test_migration(&path, "001", "create_users");
527 create_test_migration(&path, "002", "add_index");
528
529 let result = resolve_migrations(&path, DbType::PostgreSQL).unwrap();
530 assert_eq!(result.len(), 2);
531 assert_eq!(result[0].version, "001");
532 assert_eq!(result[0].name, "create_users");
533 assert_eq!(result[1].version, "002");
534 assert_eq!(result[1].name, "add_index");
535 }
536
537 #[test]
538 fn test_resolve_migrations_returns_sql_content() {
539 let temp = tempfile::tempdir().unwrap();
540 let path = temp.path().to_path_buf();
541
542 let up_path = path.join("001_init_up.sql");
543 let down_path = path.join("001_init_down.sql");
544 fs::write(&up_path, "CREATE TABLE users (id INT);").unwrap();
545 fs::write(&down_path, "DROP TABLE users;").unwrap();
546
547 let result = resolve_migrations(&path, DbType::PostgreSQL).unwrap();
548 assert_eq!(result.len(), 1);
549 assert!(result[0].sql_up.contains("CREATE TABLE users"));
550 assert!(result[0].sql_down.contains("DROP TABLE users"));
551 }
552
553 #[test]
554 fn test_resolve_migrations_supports_multiple_db_types() {
555 let temp = tempfile::tempdir().unwrap();
556 let path = temp.path().to_path_buf();
557 create_test_migration(&path, "001", "init");
558
559 let mysql_result = resolve_migrations(&path, DbType::MySQL).unwrap();
560 let pg_result = resolve_migrations(&path, DbType::PostgreSQL).unwrap();
561
562 assert_eq!(mysql_result.len(), 1);
563 assert_eq!(pg_result.len(), 1);
564 }
565
566 #[test]
567 fn test_execute_status_nonexistent_dir() {
568 let result = execute_status("/nonexistent/path/migrations");
569 assert!(matches!(result, Err(CliError::Migration(_))));
570 }
571
572 #[test]
573 fn test_execute_status_empty_dir() {
574 let temp = tempfile::tempdir().unwrap();
575 let path = temp.path().to_str().unwrap();
576 let result = execute_status(path);
577 assert!(result.is_ok());
578 }
579
580 #[test]
581 fn test_execute_status_with_migrations() {
582 let temp = tempfile::tempdir().unwrap();
583 let path = temp.path().to_path_buf();
584 create_test_migration(&path, "001", "create_users");
585
586 let path_str = temp.path().to_str().unwrap();
587 let result = execute_status(path_str);
588 assert!(result.is_ok());
589 }
590
591 #[test]
592 fn test_execute_status_with_invalid_db_type() {
593 let temp = tempfile::tempdir().unwrap();
594 let path = temp.path().to_str().unwrap();
595 let result = execute_status_with(path, "invalid_db_type", false);
596 assert!(matches!(result, Err(CliError::Migration(_))));
597 }
598
599 #[test]
600 fn test_execute_status_with_show_sql() {
601 let temp = tempfile::tempdir().unwrap();
602 let path = temp.path().to_path_buf();
603
604 let up_path = path.join("001_init_up.sql");
605 let down_path = path.join("001_init_down.sql");
606 fs::write(&up_path, "CREATE TABLE users (id INT);").unwrap();
607 fs::write(&down_path, "DROP TABLE users;").unwrap();
608
609 let path_str = temp.path().to_str().unwrap();
610 let result = execute_status_with(path_str, "postgres", true);
611 assert!(result.is_ok());
612 }
613
614 #[test]
615 fn test_execute_migrate_nonexistent_dir() {
616 let args = MigrateArgs {
617 rollback: false,
618 path: "/nonexistent/migrations".to_string(),
619 db_type: "postgres".to_string(),
620 show_sql: false,
621 url: None,
622 };
623 let result = execute_migrate(&args);
624 assert!(matches!(result, Err(CliError::Migration(_))));
625 }
626
627 #[test]
628 fn test_execute_migrate_empty_dir() {
629 let temp = tempfile::tempdir().unwrap();
630 let args = MigrateArgs {
631 rollback: false,
632 path: temp.path().to_str().unwrap().to_string(),
633 db_type: "postgres".to_string(),
634 show_sql: false,
635 url: None,
636 };
637 let result = execute_migrate(&args);
638 assert!(result.is_ok());
639 }
640
641 #[test]
642 fn test_execute_migrate_with_files_offline() {
643 let temp = tempfile::tempdir().unwrap();
644 let path = temp.path().to_path_buf();
645 create_test_migration(&path, "001", "create_users");
646
647 let args = MigrateArgs {
648 rollback: false,
649 path: temp.path().to_str().unwrap().to_string(),
650 db_type: "postgres".to_string(),
651 show_sql: false,
652 url: None,
653 };
654 let result = execute_migrate(&args);
655 assert!(result.is_ok());
656 }
657
658 #[test]
659 fn test_execute_migrate_with_show_sql_offline() {
660 let temp = tempfile::tempdir().unwrap();
661 let path = temp.path().to_path_buf();
662
663 let up_path = path.join("001_init_up.sql");
664 let down_path = path.join("001_init_down.sql");
665 fs::write(&up_path, "CREATE TABLE users (id INT);").unwrap();
666 fs::write(&down_path, "DROP TABLE users;").unwrap();
667
668 let args = MigrateArgs {
669 rollback: false,
670 path: temp.path().to_str().unwrap().to_string(),
671 db_type: "postgres".to_string(),
672 show_sql: true,
673 url: None,
674 };
675 let result = execute_migrate(&args);
676 assert!(result.is_ok());
677 }
678
679 #[test]
680 fn test_execute_migrate_with_invalid_db_type() {
681 let temp = tempfile::tempdir().unwrap();
682 let args = MigrateArgs {
683 rollback: false,
684 path: temp.path().to_str().unwrap().to_string(),
685 db_type: "invalid_db_type".to_string(),
686 show_sql: false,
687 url: None,
688 };
689 let result = execute_migrate(&args);
690 assert!(matches!(result, Err(CliError::Migration(_))));
691 }
692
693 #[test]
694 fn test_execute_migrate_rollback_offline() {
695 let temp = tempfile::tempdir().unwrap();
696 let path = temp.path().to_path_buf();
697 create_test_migration(&path, "001", "create_users");
698 create_test_migration(&path, "002", "add_index");
699
700 let args = MigrateArgs {
701 rollback: true,
702 path: temp.path().to_str().unwrap().to_string(),
703 db_type: "postgres".to_string(),
704 show_sql: false,
705 url: None,
706 };
707 let result = execute_migrate(&args);
708 assert!(result.is_ok());
709 }
710
711 #[test]
712 fn test_execute_migrate_rollback_with_show_sql_offline() {
713 let temp = tempfile::tempdir().unwrap();
714 let path = temp.path().to_path_buf();
715
716 let up_path = path.join("001_init_up.sql");
717 let down_path = path.join("001_init_down.sql");
718 fs::write(&up_path, "CREATE TABLE users (id INT);").unwrap();
719 fs::write(&down_path, "DROP TABLE users;").unwrap();
720
721 let args = MigrateArgs {
722 rollback: true,
723 path: temp.path().to_str().unwrap().to_string(),
724 db_type: "postgres".to_string(),
725 show_sql: true,
726 url: None,
727 };
728 let result = execute_migrate(&args);
729 assert!(result.is_ok());
730 }
731
732 #[test]
733 fn test_print_sql_block_empty_sql() {
734 print_sql_block("SQL UP", "");
735 }
736
737 #[test]
738 fn test_print_sql_block_with_content() {
739 print_sql_block("SQL UP", "CREATE TABLE users (id INT);");
740 }
741
742 #[test]
743 fn test_execute_status_full_offline_no_url() {
744 let temp = tempfile::tempdir().unwrap();
745 let path = temp.path().to_path_buf();
746 create_test_migration(&path, "001", "init");
747
748 let path_str = temp.path().to_str().unwrap();
749 let result = execute_status_full(path_str, "postgres", false, None);
750 assert!(result.is_ok());
751 }
752
753 #[test]
754 fn test_execute_status_full_offline_with_show_sql() {
755 let temp = tempfile::tempdir().unwrap();
756 let path = temp.path().to_path_buf();
757 let up_path = path.join("001_init_up.sql");
758 let down_path = path.join("001_init_down.sql");
759 fs::write(&up_path, "CREATE TABLE t (id INT);").unwrap();
760 fs::write(&down_path, "DROP TABLE t;").unwrap();
761
762 let path_str = temp.path().to_str().unwrap();
763 let result = execute_status_full(path_str, "postgres", true, None);
764 assert!(result.is_ok());
765 }
766
767 #[test]
768 fn test_execute_status_full_invalid_db_type() {
769 let temp = tempfile::tempdir().unwrap();
770 let path_str = temp.path().to_str().unwrap();
771 let result = execute_status_full(path_str, "invalid_db", false, None);
772 assert!(matches!(result, Err(CliError::Migration(_))));
773 }
774
775 #[test]
776 fn test_execute_migrate_online_with_invalid_url_returns_error() {
777 let temp = tempfile::tempdir().unwrap();
778 let path = temp.path().to_path_buf();
779 create_test_migration(&path, "001", "init");
780
781 let args = MigrateArgs {
782 rollback: false,
783 path: temp.path().to_str().unwrap().to_string(),
784 db_type: "postgres".to_string(),
785 show_sql: false,
786 url: Some("postgres://invalid:invalid@127.0.0.1:1/invalid".to_string()),
787 };
788 let result = execute_migrate(&args);
789 assert!(result.is_err());
791 }
792}