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().map_err(seed_error)? {
200            script.add_rows(statement.table(), statement.rows());
201            script
202                .inserts
203                .push(statement.inline_sql().map_err(seed_error)?);
204        }
205        Ok(script)
206    }};
207}
208
209/// Builds the seed script for `schema` in `dialect`.
210///
211/// `max_params` caps the values per INSERT (D1's HTTP API limits statement
212/// length, so its driver passes a lower cap).
213///
214/// # Errors
215///
216/// Returns [`CliError`] when this build has no driver for `dialect`, a
217/// table or column name in `plan` is unknown, or seeding fails (for example
218/// a foreign-key cycle).
219#[allow(unused_variables)] // every dialect arm can be compiled out
220pub fn build_script(
221    dialect: drizzle_types::Dialect,
222    schema: &drizzle_seed::schema::Schema,
223    plan: &SeedPlan,
224    max_params: Option<usize>,
225) -> Result<SeedScript, CliError> {
226    match dialect {
227        #[cfg(any(
228            feature = "rusqlite",
229            feature = "libsql",
230            feature = "turso",
231            feature = "d1-http"
232        ))]
233        drizzle_types::Dialect::SQLite => build_script!(sqlite, schema, plan, max_params),
234        #[cfg(any(feature = "postgres-sync", feature = "tokio-postgres"))]
235        drizzle_types::Dialect::PostgreSQL => build_script!(postgres, schema, plan, max_params),
236        #[cfg(any(feature = "mysql-sync", feature = "mysql-async"))]
237        drizzle_types::Dialect::MySQL => build_script!(mysql, schema, plan, max_params),
238        #[allow(unreachable_patterns)]
239        other => Err(CliError::Other(format!(
240            "this drizzle build has no {other:?} driver; build it with a driver feature to seed"
241        ))),
242    }
243}
244
245/// Runs `drizzle seed`: reads the tables from the live database (after the
246/// configured table and schema filters, without the migrations tracking
247/// table), generates rows for them, and inserts them in one transaction
248/// where the driver supports it. With `--out`, writes the SQL instead.
249///
250/// # Errors
251///
252/// Returns [`CliError`] if `db_name` does not match the config, an option is
253/// malformed, there are no credentials or no driver for them, reading the
254/// database fails, seeding fails, writing `--out` fails, a statement fails,
255/// or the user declines `--reset` ([`CliError::Aborted`]).
256pub fn run(config: &Config, db_name: Option<&str>, opts: &SeedOptions) -> Result<(), CliError> {
257    let db = config.database(db_name)?;
258    let plan = SeedPlan::from_options(opts)?;
259    let effective_dialect = overrides::resolve_dialect(db, opts.dialect);
260    let to_stdout = opts.out.as_deref() == Some(std::path::Path::new("-"));
261    // With `--out -`, stdout carries only the SQL.
262    let say = |line: String| {
263        if to_stdout {
264            eprintln!("{line}");
265        } else {
266            println!("{line}");
267        }
268    };
269
270    if !to_stdout {
271        crate::commands::harness::print_db_header(config, db_name);
272    }
273    say(output::heading("Seeding database..."));
274    say(format!(
275        "  {}: {}",
276        output::label("Dialect"),
277        effective_dialect.as_str()
278    ));
279
280    let connection = overrides::resolve_connection(db, effective_dialect, &opts.connection)?;
281    let Some(connection) = connection else {
282        return Err(CliError::MissingCredentials("seed"));
283    };
284    say(format!(
285        "  {}: {}",
286        output::label("Driver"),
287        connection.driver
288    ));
289
290    let filters = crate::db::SnapshotFilters {
291        tables: overrides::resolve_filter_list(
292            opts.filters.tables_filter.as_deref(),
293            db.tables_filter.as_ref(),
294        ),
295        schemas: overrides::resolve_schema_filters(
296            effective_dialect,
297            opts.filters.schema_filters.as_deref(),
298            db.schema_filter.as_ref(),
299        ),
300        extensions: overrides::resolve_extensions_filter(
301            effective_dialect,
302            opts.filters.extensions_filters.as_deref(),
303            db.extensions_filters.as_deref(),
304        ),
305        roles: None,
306    };
307    let snapshot = crate::db::introspect_for_seed(&connection, &filters, db.migrations_table())?;
308    let schema = drizzle_seed::schema::Schema::from_snapshot(&snapshot).map_err(seed_error)?;
309    if schema.tables().is_empty() {
310        say(output::warning("No tables found to seed."));
311        return Ok(());
312    }
313
314    let max_params = (connection.driver == Driver::D1Http).then_some(250);
315    let script = build_script(effective_dialect.to_base(), &schema, &plan, max_params)?;
316
317    if let Some(out) = &opts.out {
318        let sql = script.to_sql();
319        if to_stdout {
320            print!("{sql}");
321        } else {
322            std::fs::write(out, sql).map_err(|error| {
323                CliError::IoError(format!("Failed to write {}: {error}", out.display()))
324            })?;
325            say(format!(
326                "{} Wrote {} statement(s) to {}",
327                output::success("Done!"),
328                script.reset.len() + script.inserts.len(),
329                out.display()
330            ));
331        }
332        return Ok(());
333    }
334
335    if plan.reset && !opts.force && !confirm_reset(&script)? {
336        return Err(CliError::Aborted(
337            "Seed aborted: --reset was not confirmed, so nothing was changed".into(),
338        ));
339    }
340
341    let statements: Vec<String> = script
342        .reset
343        .iter()
344        .chain(&script.inserts)
345        .cloned()
346        .collect();
347    crate::db::execute_seed(&connection, &statements)?;
348
349    println!();
350    for (table, rows) in &script.rows {
351        println!("  {} {table}: {rows} row(s)", output::success("+"));
352    }
353    println!();
354    println!("{}", output::success("Seed complete!"));
355    Ok(())
356}
357
358/// Asks before `--reset` deletes rows; refuses in a non-interactive shell.
359fn confirm_reset(script: &SeedScript) -> Result<bool, CliError> {
360    use std::io::{self, IsTerminal, Write};
361
362    if !io::stdin().is_terminal() {
363        return Err(CliError::Other(
364            "Refusing to delete rows for --reset in non-interactive mode. Use --force.".into(),
365        ));
366    }
367    println!(
368        "{}",
369        output::warning(&format!(
370            "--reset deletes every row of the {} seeded table(s) first.",
371            script.rows.len()
372        ))
373    );
374    print!("Continue? [y/N]: ");
375    io::stdout()
376        .flush()
377        .map_err(|error| CliError::IoError(error.to_string()))?;
378    let mut answer = String::new();
379    io::stdin()
380        .read_line(&mut answer)
381        .map_err(|error| CliError::IoError(error.to_string()))?;
382    Ok(matches!(answer.trim(), "y" | "Y" | "yes" | "YES" | "Yes"))
383}
384
385#[cfg(test)]
386mod tests {
387    use super::*;
388    use clap::Parser;
389
390    #[derive(Parser)]
391    struct Cli {
392        #[command(flatten)]
393        seed: SeedOptions,
394    }
395
396    fn plan(args: &[&str]) -> Result<SeedPlan, CliError> {
397        let mut argv = vec!["drizzle"];
398        argv.extend_from_slice(args);
399        SeedPlan::from_options(&Cli::parse_from(argv).seed)
400    }
401
402    #[test]
403    fn parses_counts_relations_and_skips() {
404        let plan = plan(&[
405            "--seed",
406            "7",
407            "--count",
408            "20",
409            "--count",
410            "users=100",
411            "--count",
412            "auth.sessions = 5",
413            "--relation",
414            "users:posts=3",
415            "--skip",
416            "audit_log",
417            "--reset",
418        ])
419        .unwrap();
420        assert_eq!(
421            plan,
422            SeedPlan {
423                seed: 7,
424                default_count: Some(20),
425                counts: vec![("users".into(), 100), ("auth.sessions".into(), 5)],
426                relations: vec![("users".into(), "posts".into(), 3)],
427                skips: vec!["audit_log".into()],
428                reset: true,
429            }
430        );
431    }
432
433    #[test]
434    fn rejects_malformed_values() {
435        for args in [
436            &["--count", "lots"][..],
437            &["--count", "users=lots"],
438            &["--count", "=5"],
439            &["--count", "1", "--count", "2"],
440            &["--relation", "users=3"],
441            &["--relation", "users:=3"],
442            &["--relation", "users:posts=x"],
443        ] {
444            assert!(plan(args).is_err(), "{args:?}");
445        }
446    }
447}
drizzle_cli/commands/seed.rs

drizzle_cli/commands/
seed.rs
