Skip to main content

sz_rust_cli/cmd/
seed.rs

1//! `db:seed` 命令 — 数据填充(对齐 PHP `think db:seed`)
2//!
3//! ## PHP 对齐
4//!
5//! PHP `db:seed` 通过 `Seeder` 类执行数据填充:
6//! ```text
7//! php think db:seed              # 运行默认 DatabaseSeeder
8//! php think db:seed -s UserSeeder # 运行指定填充器
9//! ```
10//!
11//! Rust 端由于静态编译无法按类名动态加载,采用以下策略:
12//! - **SQL 文件模式**(默认):从 `seeds/` 目录加载 `.sql` 文件并执行
13//! - **程序化模式**:业务实现 `sz_rust_core::seed::Seeder` trait,通过 `SeedRunner` 注册执行
14//!
15//! ## 文件命名约定
16//!
17//! `seeds/` 目录下的 `.sql` 文件按文件名升序执行:
18//! ```text
19//! seeds/
20//! ├── 001_users_seed.sql
21//! ├── 002_roles_seed.sql
22//! └── 003_user_roles_seed.sql
23//! ```
24//!
25//! ## 模式
26//!
27//! - **离线模式**(默认,未提供 `--url`):仅列出待执行的 SQL 文件内容
28//! - **在线模式**(提供 `--url`):连接数据库执行真实填充
29
30use std::path::{Path, PathBuf};
31
32use sz_rust_core::orm::{Connection, ConnectionFactory, DbType};
33
34use crate::error::CliError;
35
36/// `db:seed` 命令参数
37///
38/// 对齐 PHP `php think db:seed` / `php think db:seed -s <class>`。
39#[derive(Debug, Clone)]
40pub struct SeedArgs {
41    /// 填充目录(默认 `seeds`)
42    pub path: String,
43
44    /// 数据库类型(默认 `postgres`,对齐 sz-orm `DbType`)
45    pub db_type: String,
46
47    /// 仅打印 SQL 内容(dry-run 模式,便于审查)
48    pub show_sql: bool,
49
50    /// 数据库连接 URL(启用在线模式)
51    ///
52    /// 提供时连接数据库执行真实填充;省略时为离线模式(仅列出待执行的 SQL)。
53    pub url: Option<String>,
54
55    /// 指定填充器文件名(不含扩展名,如 `001_users_seed`)
56    ///
57    /// 省略时执行目录下所有 `.sql` 文件。
58    pub class: Option<String>,
59}
60
61impl Default for SeedArgs {
62    fn default() -> Self {
63        Self {
64            path: "seeds".to_string(),
65            db_type: "postgres".to_string(),
66            show_sql: false,
67            url: None,
68            class: None,
69        }
70    }
71}
72
73/// 执行 db:seed 命令
74///
75/// # 模式
76///
77/// - **离线模式**(默认,未提供 `--url`):仅解析填充目录并打印待执行内容
78/// - **在线模式**(提供 `--url`):连接数据库执行真实填充
79pub fn execute_seed(args: &SeedArgs) -> Result<(), CliError> {
80    let path = PathBuf::from(&args.path);
81
82    if !path.exists() {
83        return Err(CliError::Generic(format!(
84            "Seed directory not found: {}",
85            path.display()
86        )));
87    }
88
89    let db_type = DbType::from_str(&args.db_type)
90        .ok_or_else(|| CliError::Generic(format!("Unknown database type: {}", args.db_type)))?;
91
92    let seed_files = resolve_seed_files(&path, args.class.as_deref())?;
93
94    if seed_files.is_empty() {
95        println!("No seed files found in: {}", path.display());
96        return Ok(());
97    }
98
99    match &args.url {
100        None => execute_seed_offline(args, &seed_files),
101        Some(url) => execute_seed_online(args, &seed_files, url, db_type),
102    }
103}
104
105/// 离线模式执行 seed(仅打印,不连库)
106fn execute_seed_offline(args: &SeedArgs, seed_files: &[SeedFile]) -> Result<(), CliError> {
107    println!("Seed files in: {}", args.path);
108    for sf in seed_files {
109        println!("  Would execute: {}", sf.name);
110        if args.show_sql {
111            print_sql_block("SQL", &sf.content);
112        }
113    }
114    println!(
115        "Total: {} seed file(s). Note: Actual execution requires database connection (offline mode).",
116        seed_files.len()
117    );
118    Ok(())
119}
120
121/// 在线模式执行 seed(连接数据库真实执行)
122fn execute_seed_online(
123    args: &SeedArgs,
124    seed_files: &[SeedFile],
125    url: &str,
126    db_type: DbType,
127) -> Result<(), CliError> {
128    let rt = tokio::runtime::Builder::new_current_thread()
129        .enable_all()
130        .build()
131        .map_err(|e| CliError::Generic(format!("Failed to create tokio runtime: {}", e)))?;
132
133    rt.block_on(async move {
134        let mut conn = create_connection(url, db_type).await?;
135
136        println!("Running {} seed file(s):", seed_files.len());
137        for sf in seed_files {
138            println!("  Seeding: {}", sf.name);
139            if args.show_sql {
140                print_sql_block("SQL", &sf.content);
141            }
142            conn.execute(&sf.content)
143                .await
144                .map_err(|e| CliError::Generic(format!("Seed failed ({}): {}", sf.name, e)))?;
145            println!("  Completed: {}", sf.name);
146        }
147        println!("Seed completed: {} file(s) applied.", seed_files.len());
148
149        Ok::<(), CliError>(())
150    })
151}
152
153/// 填充文件
154#[derive(Debug, Clone)]
155struct SeedFile {
156    /// 文件名(含扩展名)
157    name: String,
158    /// SQL 内容
159    content: String,
160}
161
162/// 解析填充目录,返回按文件名排序的填充文件列表
163///
164/// # 参数
165///
166/// - `path`:填充目录
167/// - `class_filter`:可选的文件名过滤(不含扩展名)
168///
169/// # 错误
170///
171/// - [`CliError::Generic`]:目录读取失败或文件读取失败
172fn resolve_seed_files(path: &Path, class_filter: Option<&str>) -> Result<Vec<SeedFile>, CliError> {
173    let entries = std::fs::read_dir(path).map_err(|e| {
174        CliError::Generic(format!(
175            "Failed to read seed directory {}: {}",
176            path.display(),
177            e
178        ))
179    })?;
180
181    let mut files: Vec<PathBuf> = Vec::new();
182    for entry in entries {
183        let entry = entry
184            .map_err(|e| CliError::Generic(format!("Failed to read directory entry: {}", e)))?;
185        let p = entry.path();
186        if p.is_file() && p.extension().and_then(|e| e.to_str()) == Some("sql") {
187            if let Some(filter) = class_filter {
188                // 按文件名(不含扩展名)匹配
189                if p.file_stem().and_then(|s| s.to_str()) != Some(filter) {
190                    continue;
191                }
192            }
193            files.push(p);
194        }
195    }
196
197    // 按文件名升序排序
198    files.sort();
199
200    let mut seed_files = Vec::with_capacity(files.len());
201    for f in files {
202        let name = f
203            .file_name()
204            .and_then(|s| s.to_str())
205            .unwrap_or("unknown")
206            .to_string();
207        let content = std::fs::read_to_string(&f).map_err(|e| {
208            CliError::Generic(format!("Failed to read seed file {}: {}", f.display(), e))
209        })?;
210        seed_files.push(SeedFile { name, content });
211    }
212
213    Ok(seed_files)
214}
215
216/// 创建数据库连接(按 DbType 选择驱动)
217///
218/// 复用 `migrate` 模块的连接创建逻辑,保持一致性。
219async fn create_connection(url: &str, db_type: DbType) -> Result<Box<dyn Connection>, CliError> {
220    use std::sync::Arc;
221    use sz_orm_sqlx::{
222        MySqlPoolHandle, PgPoolHandle, SqlitePoolHandle, SqlxMySqlConnectionFactory,
223        SqlxPgConnectionFactory, SqlxSqliteConnectionFactory,
224    };
225
226    match db_type {
227        DbType::PostgreSQL => {
228            let pool = PgPoolHandle::connect(url)
229                .await
230                .map_err(|e| CliError::Generic(format!("PostgreSQL connect failed: {}", e)))?;
231            let factory = SqlxPgConnectionFactory::new(Arc::new(pool));
232            let conn = factory
233                .create()
234                .await
235                .map_err(|e| CliError::Generic(format!("PostgreSQL acquire failed: {}", e)))?;
236            Ok(conn)
237        }
238        DbType::MySQL => {
239            let pool = MySqlPoolHandle::connect(url)
240                .await
241                .map_err(|e| CliError::Generic(format!("MySQL connect failed: {}", e)))?;
242            let factory = SqlxMySqlConnectionFactory::new(Arc::new(pool));
243            let conn = factory
244                .create()
245                .await
246                .map_err(|e| CliError::Generic(format!("MySQL acquire failed: {}", e)))?;
247            Ok(conn)
248        }
249        DbType::Sqlite => {
250            let pool = SqlitePoolHandle::connect(url)
251                .await
252                .map_err(|e| CliError::Generic(format!("SQLite connect failed: {}", e)))?;
253            let factory = SqlxSqliteConnectionFactory::new(Arc::new(pool));
254            let conn = factory
255                .create()
256                .await
257                .map_err(|e| CliError::Generic(format!("SQLite acquire failed: {}", e)))?;
258            Ok(conn)
259        }
260        _ => Err(CliError::Generic(format!(
261            "Online seed not supported for db_type {:?}. Supported: PostgreSQL, MySQL, SQLite.",
262            db_type
263        ))),
264    }
265}
266
267/// 打印 SQL 代码块(带标题分隔符)
268fn print_sql_block(title: &str, sql: &str) {
269    if sql.is_empty() {
270        return;
271    }
272    println!("  --- {} ---", title);
273    for line in sql.lines() {
274        println!("  {}", line);
275    }
276    println!("  {}", "-".repeat(title.len() + 8));
277}
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282    use std::fs;
283
284    /// 创建测试用填充文件
285    fn create_test_seed_file(dir: &Path, name: &str, content: &str) {
286        let path = dir.join(name);
287        fs::write(&path, content).expect("Failed to write test seed file");
288    }
289
290    #[test]
291    fn test_resolve_seed_files_sorted_by_name() {
292        let tmp = tempfile::tempdir().expect("Failed to create temp dir");
293        create_test_seed_file(tmp.path(), "003_third.sql", "INSERT INTO t VALUES (3);");
294        create_test_seed_file(tmp.path(), "001_first.sql", "INSERT INTO t VALUES (1);");
295        create_test_seed_file(tmp.path(), "002_second.sql", "INSERT INTO t VALUES (2);");
296
297        let files = resolve_seed_files(tmp.path(), None).expect("resolve failed");
298        assert_eq!(files.len(), 3);
299        assert_eq!(files[0].name, "001_first.sql");
300        assert_eq!(files[1].name, "002_second.sql");
301        assert_eq!(files[2].name, "003_third.sql");
302    }
303
304    #[test]
305    fn test_resolve_seed_files_ignores_non_sql() {
306        let tmp = tempfile::tempdir().expect("Failed to create temp dir");
307        create_test_seed_file(tmp.path(), "001_first.sql", "INSERT 1;");
308        // 非 SQL 文件应被忽略
309        let txt_path = tmp.path().join("readme.txt");
310        fs::write(&txt_path, "ignore me").expect("write txt");
311        let md_path = tmp.path().join("notes.md");
312        fs::write(&md_path, "ignore me").expect("write md");
313
314        let files = resolve_seed_files(tmp.path(), None).expect("resolve failed");
315        assert_eq!(files.len(), 1);
316        assert_eq!(files[0].name, "001_first.sql");
317    }
318
319    #[test]
320    fn test_resolve_seed_files_class_filter() {
321        let tmp = tempfile::tempdir().expect("Failed to create temp dir");
322        create_test_seed_file(tmp.path(), "001_first.sql", "INSERT 1;");
323        create_test_seed_file(tmp.path(), "002_second.sql", "INSERT 2;");
324
325        let files = resolve_seed_files(tmp.path(), Some("002_second")).expect("resolve failed");
326        assert_eq!(files.len(), 1);
327        assert_eq!(files[0].name, "002_second.sql");
328    }
329
330    #[test]
331    fn test_resolve_seed_files_empty_directory() {
332        let tmp = tempfile::tempdir().expect("Failed to create temp dir");
333        let files = resolve_seed_files(tmp.path(), None).expect("resolve failed");
334        assert!(files.is_empty());
335    }
336
337    #[test]
338    fn test_resolve_seed_files_nonexistent_directory() {
339        let result = resolve_seed_files(Path::new("/nonexistent/path/to/seeds"), None);
340        assert!(result.is_err());
341    }
342
343    #[test]
344    fn test_seed_args_default() {
345        let args = SeedArgs::default();
346        assert_eq!(args.path, "seeds");
347        assert_eq!(args.db_type, "postgres");
348        assert!(!args.show_sql);
349        assert!(args.url.is_none());
350        assert!(args.class.is_none());
351    }
352
353    #[test]
354    fn test_execute_seed_offline_with_show_sql() {
355        let tmp = tempfile::tempdir().expect("Failed to create temp dir");
356        create_test_seed_file(
357            tmp.path(),
358            "001_users.sql",
359            "INSERT INTO users (name) VALUES ('admin');",
360        );
361
362        let args = SeedArgs {
363            path: tmp.path().to_string_lossy().to_string(),
364            show_sql: true,
365            ..SeedArgs::default()
366        };
367
368        let result = execute_seed(&args);
369        assert!(result.is_ok());
370    }
371
372    #[test]
373    fn test_execute_seed_offline_without_show_sql() {
374        let tmp = tempfile::tempdir().expect("Failed to create temp dir");
375        create_test_seed_file(
376            tmp.path(),
377            "001_users.sql",
378            "INSERT INTO users (name) VALUES ('admin');",
379        );
380
381        let args = SeedArgs {
382            path: tmp.path().to_string_lossy().to_string(),
383            ..SeedArgs::default()
384        };
385
386        let result = execute_seed(&args);
387        assert!(result.is_ok());
388    }
389
390    #[test]
391    fn test_execute_seed_directory_not_found() {
392        let args = SeedArgs {
393            path: "/nonexistent/path/to/seeds".to_string(),
394            ..SeedArgs::default()
395        };
396        let result = execute_seed(&args);
397        assert!(result.is_err());
398        let err = result.unwrap_err().to_string();
399        assert!(err.contains("Seed directory not found"));
400    }
401
402    #[test]
403    fn test_execute_seed_empty_directory() {
404        let tmp = tempfile::tempdir().expect("Failed to create temp dir");
405        let args = SeedArgs {
406            path: tmp.path().to_string_lossy().to_string(),
407            ..SeedArgs::default()
408        };
409        let result = execute_seed(&args);
410        assert!(result.is_ok());
411    }
412
413    #[test]
414    fn test_execute_seed_invalid_db_type() {
415        let tmp = tempfile::tempdir().expect("Failed to create temp dir");
416        create_test_seed_file(tmp.path(), "001.sql", "INSERT 1;");
417        let args = SeedArgs {
418            path: tmp.path().to_string_lossy().to_string(),
419            db_type: "invalid_db".to_string(),
420            ..SeedArgs::default()
421        };
422        let result = execute_seed(&args);
423        assert!(result.is_err());
424        assert!(result
425            .unwrap_err()
426            .to_string()
427            .contains("Unknown database type"));
428    }
429
430    #[test]
431    fn test_execute_seed_class_filter_offline() {
432        let tmp = tempfile::tempdir().expect("Failed to create temp dir");
433        create_test_seed_file(tmp.path(), "001_first.sql", "INSERT 1;");
434        create_test_seed_file(tmp.path(), "002_second.sql", "INSERT 2;");
435
436        let args = SeedArgs {
437            path: tmp.path().to_string_lossy().to_string(),
438            class: Some("002_second".to_string()),
439            ..SeedArgs::default()
440        };
441
442        let result = execute_seed(&args);
443        assert!(result.is_ok());
444    }
445
446    #[test]
447    fn test_execute_seed_class_filter_not_found() {
448        let tmp = tempfile::tempdir().expect("Failed to create temp dir");
449        create_test_seed_file(tmp.path(), "001_first.sql", "INSERT 1;");
450
451        let args = SeedArgs {
452            path: tmp.path().to_string_lossy().to_string(),
453            class: Some("nonexistent".to_string()),
454            ..SeedArgs::default()
455        };
456
457        let result = execute_seed(&args);
458        assert!(result.is_ok()); // 空列表不算错误
459    }
460
461    #[test]
462    fn test_print_sql_block_empty() {
463        // 空内容不应输出任何东西
464        print_sql_block("TITLE", "");
465    }
466
467    #[test]
468    fn test_print_sql_block_non_empty() {
469        print_sql_block("TITLE", "SELECT 1;\nSELECT 2;");
470    }
471}