drizzle_cli/commands/
seed.rs

1//! `drizzle seed`: fills a live database with deterministic test data.
2//!
3//! The tables are read from the database itself, so no Rust schema is
4//! needed. Needs a database driver feature enabled at build time.
5
6use crate::commands::overrides::{self, ConnectionOverrides, FilterArgs};
7use crate::config::{Config, Dialect, Driver};
8use crate::error::CliError;
9use crate::output;
10use std::path::PathBuf;
11
12#[derive(clap::Args, Debug, Clone)]
13pub struct SeedOptions {
14    /// RNG seed: the same seed and drizzle version give the same rows
15    #[arg(long, default_value_t = 0)]
16    pub seed: u64,
17
18    /// Rows per table: `--count 50` for every table, `--count users=100`
19    /// for one (repeatable). A table without a count that references a
20    /// seeded table gets one row per parent row, or `--relation` rows.
21    #[arg(long, value_name = "N|TABLE=N")]
22    pub count: Vec<String>,
23
24    /// Rows of CHILD per row of PARENT: `--relation users:posts=3`
25    /// (repeatable). CHILD needs a foreign key to PARENT.
26    #[arg(long, value_name = "PARENT:CHILD=N")]
27    pub relation: Vec<String>,
28
29    /// Leave a table out (repeatable)
30    #[arg(long, value_name = "TABLE")]
31    pub skip: Vec<String>,
32
33    /// Empty the seeded tables first (children before parents)
34    #[arg(long)]
35    pub reset: bool,
36
37    /// Write the SQL to FILE instead of running it (`-` for stdout)
38    #[arg(long, value_name = "FILE")]
39    pub out: Option<PathBuf>,
40
41    /// Run `--reset` without asking for confirmation
42    #[arg(long)]
43    pub force: bool,
44
45    /// Override dialect from config
46    #[arg(long)]
47    pub dialect: Option<Dialect>,
48
49    #[command(flatten)]
50    pub filters: FilterArgs,
51
52    #[command(flatten)]
53    pub connection: ConnectionOverrides,
54}
55
56/// The parsed `--count`, `--relation` and `--skip` arguments.
57#[derive(Debug, Default, PartialEq, Eq)]
58pub struct SeedPlan {
59    /// `--seed`.
60    pub seed: u64,
61    /// A bare `--count N`.
62    pub default_count: Option<usize>,
63    /// `--count TABLE=N`.
64    pub counts: Vec<(String, usize)>,
65    /// `--relation PARENT:CHILD=N`.
66    pub relations: Vec<(String, String, usize)>,
67    /// `--skip TABLE`.
68    pub skips: Vec<String>,
69    /// `--reset`.
70    pub reset: bool,
71}
72
73impl SeedPlan {
74    /// Parses the seed options.
75    ///
76    /// # Errors
77    ///
78    /// Returns [`CliError::Other`] for a malformed `--count` or `--relation`
79    /// value, or two bare `--count`s.
80    pub fn from_options(opts: &SeedOptions) -> Result<Self, CliError> {
81        let number = |text: &str, flag: &str, value: &str| {
82            text.trim().parse::<usize>().map_err(|_| {
83                CliError::Other(format!(
84                    "--{flag} {value}: `{}` is not a row count",
85                    text.trim()
86                ))
87            })
88        };
89        let mut plan = Self {
90            seed: opts.seed,
91            reset: opts.reset,
92            skips: opts.skip.clone(),
93            ..Self::default()
94        };
95        for value in &opts.count {
96            match value.rsplit_once('=') {
97                Some((table, count)) if !table.trim().is_empty() => {
98                    plan.counts
99                        .push((table.trim().to_owned(), number(count, "count", value)?));
100                }
101                Some(_) => {
102                    return Err(CliError::Other(format!(
103                        "--count {value}: expected N or TABLE=N"
104                    )));
105                }
106                None if plan.default_count.is_some() => {
107                    return Err(CliError::Other(
108                        "--count without a table was given twice".to_owned(),
109                    ));
110                }
111                None => plan.default_count = Some(number(value, "count", value)?),
112            }
113        }
114        for value in &opts.relation {
115            let parsed = value.rsplit_once('=').and_then(|(tables, count)| {
116                let (parent, child) = tables.split_once(':')?;
117                (!parent.trim().is_empty() && !child.trim().is_empty())
118                    .then(|| (parent.trim().to_owned(), child.trim().to_owned(), count))
119            });
120            let Some((parent, child, count)) = parsed else {
121                return Err(CliError::Other(format!(
122                    "--relation {value}: expected PARENT:CHILD=N"
123                )));
124            };
125            plan.relations
126                .push((parent, child, number(count, "relation", value)?));
127        }
128        Ok(plan)
129    }
130}
131
132/// The SQL a seed runs, with every value written inline.
133#[derive(Debug, Default)]
134pub struct SeedScript {
135    /// `DELETE`s (and resets) that empty the tables, children first.
136    pub reset: Vec<String>,
137    /// The INSERTs (and `PostgreSQL` sequence updates), parents first.
138    pub inserts: Vec<String>,
139    /// Rows inserted per table, in insert order.
140    pub rows: Vec<(String, usize)>,
141}
142
143impl SeedScript {
144    #[allow(dead_code)] // unused in a build without drivers
145    fn add_rows(&mut self, table: &str, rows: usize) {
146        match self.rows.last_mut() {
147            Some((last, count)) if last == table => *count += rows,
148            _ if rows > 0 => self.rows.push((table.to_owned(), rows)),
149            _ => {}
150        }
151    }
152
153    /// The whole script: one statement per line, each ending in `;`.
154    #[must_use]
155    pub fn to_sql(&self) -> String {
156        let mut sql = String::new();
157        for statement in self.reset.iter().chain(&self.inserts) {
158            sql.push_str(statement);
159            sql.push_str(";\n");
160        }
161        sql
162    }
163}
164
165fn seed_error(error: drizzle_seed::SeedError) -> CliError {
166    CliError::Other(format!("Seeding failed: {error}"))
167}
168
169/// Builds the script for `schema` with the dialect's `SeedConfig`
170/// constructor.
171#[allow(unused_macros)] // unused in a build without drivers
172macro_rules! build_script {
173    ($constructor:ident, $schema:expr, $plan:expr, $max_params:expr) => {{
174        let plan: &SeedPlan = $plan;
175        let mut config = drizzle_seed::SeedConfig::$constructor($schema).seed(plan.seed);
176        if let Some(count) = plan.default_count {
177            config = config.default_count(count);
178        }
179        if let Some(limit) = $max_params {
180            config = config.max_params(limit);
181        }
182        for (table, count) in &plan.counts {
183            config = config.count_by_name(table, *count);
184        }
185        for (parent, child, count) in &plan.relations {
186            config = config.relation_by_name(parent, child, *count);
187        }
188        for table in &plan.skips {
189            config = config.skip_by_name(table);
190        }
191        let mut script = SeedScript::default();
192        if plan.reset {
193            for statement in config.reset_plan().map_err(seed_error)? {
194                script
195                    .reset
196                    .push(statement.inline_sql().map_err(seed_error)?);
197            }
198        }
199        for statement in config.try_generate() {
200            let statement = statement.map_err(seed_error)?;
201            script.add_rows(statement.table(), statement.rows());
202            script
203                .inserts
204                .push(statement.inline_sql().map_err(seed_error)?);
205        }
206        Ok(script)
207    }};
208}
209
210/// Builds the seed script for `schema` in `dialect`.
211///
212/// `max_params` caps the values per INSERT (D1's HTTP API limits statement
213/// length, so its driver passes a lower cap).
214///
215/// # Errors
216///
217/// Returns [`CliError`] when this build has no driver for `dialect`, a
218/// table or column name in `plan` is unknown, or seeding fails (for example
219/// a foreign-key cycle).
220#[allow(unused_variables)] // every dialect arm can be compiled out
221pub fn build_script(
222    dialect: drizzle_types::Dialect,
223    schema: &drizzle_seed::schema::Schema,
224    plan: &SeedPlan,
225    max_params: Option<usize>,
226) -> Result<SeedScript, CliError> {
227    match dialect {
228        #[cfg(any(
229            feature = "rusqlite",
230            feature = "libsql",
231            feature = "turso",
232            feature = "d1-http"
233        ))]
234        drizzle_types::Dialect::SQLite => build_script!(sqlite, schema, plan, max_params),
235        #[cfg(any(feature = "postgres-sync", feature = "tokio-postgres"))]
236        drizzle_types::Dialect::PostgreSQL => build_script!(postgres, schema, plan, max_params),
237        #[cfg(any(feature = "mysql-sync", feature = "mysql-async"))]
238        drizzle_types::Dialect::MySQL => build_script!(mysql, schema, plan, max_params),
239        #[allow(unreachable_patterns)]
240        other => Err(CliError::Other(format!(
241            "this drizzle build has no {other:?} driver; build it with a driver feature to seed"
242        ))),
243    }
244}
245
246/// Runs `drizzle seed`: reads the tables from the live database (after the
247/// configured table and schema filters, without the migrations tracking
248/// table), generates rows for them, and inserts them in one transaction
249/// where the driver supports it. With `--out`, writes the SQL instead.
250///
251/// # Errors
252///
253/// Returns [`CliError`] if `db_name` does not match the config, an option is
254/// malformed, there are no credentials or no driver for them, reading the
255/// database fails, seeding fails, writing `--out` fails, a statement fails,
256/// or the user declines `--reset` ([`CliError::Aborted`]).
257pub fn run(config: &Config, db_name: Option<&str>, opts: &SeedOptions) -> Result<(), CliError> {
258    let db = config.database(db_name)?;
259    let plan = SeedPlan::from_options(opts)?;
260    let effective_dialect = overrides::resolve_dialect(db, opts.dialect);
261    let to_stdout = opts.out.as_deref() == Some(std::path::Path::new("-"));
262    // With `--out -`, stdout carries only the SQL.
263    let say = |line: String| {
264        if to_stdout {
265            eprintln!("{line}");
266        } else {
267            println!("{line}");
268        }
269    };
270
271    if !to_stdout {
272        crate::commands::harness::print_db_header(config, db_name);
273    }
274    say(output::heading("Seeding database..."));
275    say(format!(
276        "  {}: {}",
277        output::label("Dialect"),
278        effective_dialect.as_str()
279    ));
280
281    let connection = overrides::resolve_connection(db, effective_dialect, &opts.connection)?;
282    let Some(connection) = connection else {
283        return Err(CliError::MissingCredentials("seed"));
284    };
285    say(format!(
286        "  {}: {}",
287        output::label("Driver"),
288        connection.driver
289    ));
290
291    let filters = crate::db::SnapshotFilters {
292        tables: overrides::resolve_filter_list(
293            opts.filters.tables_filter.as_deref(),
294            db.tables_filter.as_ref(),
295        ),
296        schemas: overrides::resolve_schema_filters(
297            effective_dialect,
298            opts.filters.schema_filters.as_deref(),
299            db.schema_filter.as_ref(),
300        ),
301        extensions: overrides::resolve_extensions_filter(
302            effective_dialect,
303            opts.filters.extensions_filters.as_deref(),
304            db.extensions_filters.as_deref(),
305        ),
306        roles: None,
307    };
308    let snapshot = crate::db::introspect_for_seed(&connection, &filters, db.migrations_table())?;
309    let schema = drizzle_seed::schema::Schema::from_snapshot(&snapshot).map_err(seed_error)?;
310    if schema.tables().is_empty() {
311        say(output::warning("No tables found to seed."));
312        return Ok(());
313    }
314
315    let max_params = (connection.driver == Driver::D1Http).then_some(250);
316    let script = build_script(effective_dialect.to_base(), &schema, &plan, max_params)?;
317
318    if let Some(out) = &opts.out {
319        let sql = script.to_sql();
320        if to_stdout {
321            print!("{sql}");
322        } else {
323            std::fs::write(out, sql).map_err(|error| {
324                CliError::IoError(format!("Failed to write {}: {error}", out.display()))
325            })?;
326            say(format!(
327                "{} Wrote {} statement(s) to {}",
328                output::success("Done!"),
329                script.reset.len() + script.inserts.len(),
330                out.display()
331            ));
332        }
333        return Ok(());
334    }
335
336    if plan.reset && !opts.force && !confirm_reset(&script)? {
337        return Err(CliError::Aborted(
338            "Seed aborted: --reset was not confirmed, so nothing was changed".into(),
339        ));
340    }
341
342    let statements: Vec<String> = script
343        .reset
344        .iter()
345        .chain(&script.inserts)
346        .cloned()
347        .collect();
348    crate::db::execute_seed(&connection, &statements)?;
349
350    println!();
351    for (table, rows) in &script.rows {
352        println!("  {} {table}: {rows} row(s)", output::success("+"));
353    }
354    println!();
355    println!("{}", output::success("Seed complete!"));
356    Ok(())
357}
358
359/// Asks before `--reset` deletes rows; refuses in a non-interactive shell.
360fn confirm_reset(script: &SeedScript) -> Result<bool, CliError> {
361    use std::io::{self, IsTerminal, Write};
362
363    if !io::stdin().is_terminal() {
364        return Err(CliError::Other(
365            "Refusing to delete rows for --reset in non-interactive mode. Use --force.".into(),
366        ));
367    }
368    println!(
369        "{}",
370        output::warning(&format!(
371            "--reset deletes every row of the {} seeded table(s) first.",
372            script.rows.len()
373        ))
374    );
375    print!("Continue? [y/N]: ");
376    io::stdout()
377        .flush()
378        .map_err(|error| CliError::IoError(error.to_string()))?;
379    let mut answer = String::new();
380    io::stdin()
381        .read_line(&mut answer)
382        .map_err(|error| CliError::IoError(error.to_string()))?;
383    Ok(matches!(answer.trim(), "y" | "Y" | "yes" | "YES" | "Yes"))
384}
385
386#[cfg(test)]
387mod tests {
388    use super::*;
389    use clap::Parser;
390
391    #[derive(Parser)]
392    struct Cli {
393        #[command(flatten)]
394        seed: SeedOptions,
395    }
396
397    fn plan(args: &[&str]) -> Result<SeedPlan, CliError> {
398        let mut argv = vec!["drizzle"];
399        argv.extend_from_slice(args);
400        SeedPlan::from_options(&Cli::parse_from(argv).seed)
401    }
402
403    #[test]
404    fn parses_counts_relations_and_skips() {
405        let plan = plan(&[
406            "--seed",
407            "7",
408            "--count",
409            "20",
410            "--count",
411            "users=100",
412            "--count",
413            "auth.sessions = 5",
414            "--relation",
415            "users:posts=3",
416            "--skip",
417            "audit_log",
418            "--reset",
419        ])
420        .unwrap();
421        assert_eq!(
422            plan,
423            SeedPlan {
424                seed: 7,
425                default_count: Some(20),
426                counts: vec![("users".into(), 100), ("auth.sessions".into(), 5)],
427                relations: vec![("users".into(), "posts".into(), 3)],
428                skips: vec!["audit_log".into()],
429                reset: true,
430            }
431        );
432    }
433
434    #[test]
435    fn rejects_malformed_values() {
436        for args in [
437            &["--count", "lots"][..],
438            &["--count", "users=lots"],
439            &["--count", "=5"],
440            &["--count", "1", "--count", "2"],
441            &["--relation", "users=3"],
442            &["--relation", "users:=3"],
443            &["--relation", "users:posts=x"],
444        ] {
445            assert!(plan(args).is_err(), "{args:?}");
446        }
447    }
448}
drizzle_cli/commands/seed.rs

drizzle_cli/commands/
seed.rs
