sz_rust_cli/cmd/
seed.rs

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

sz_rust_cli/cmd/
seed.rs
