drizzle_seed/
config.rs

1//! [`SeedConfig`], the builder for generating seed data.
2
3use crate::SeedError;
4use crate::generator::{Generator, GeneratorKind};
5use crate::identity::{ColumnId, TableId};
6use drizzle_core::{Relation, SQLSchemaImpl, SQLTableInfo, SchemaHasTable, TableRef};
7#[cfg(any(feature = "sqlite", feature = "postgres", feature = "mysql"))]
8use drizzle_core::{SQLColumn, SQLColumnInfo};
9use std::collections::{HashMap, HashSet};
10use std::marker::PhantomData;
11use std::sync::Arc;
12
13#[cfg(feature = "sqlite")]
14use crate::Sqlite;
15#[cfg(feature = "sqlite")]
16use drizzle_sqlite::traits::SQLiteColumn;
17#[cfg(feature = "sqlite")]
18use drizzle_sqlite::values::SQLiteValue;
19
20#[cfg(feature = "postgres")]
21use crate::Postgres;
22#[cfg(feature = "postgres")]
23use drizzle_postgres::traits::PostgresColumn;
24#[cfg(feature = "postgres")]
25use drizzle_postgres::values::PostgresValue;
26
27#[cfg(feature = "mysql")]
28use crate::MySql;
29#[cfg(feature = "mysql")]
30use drizzle_mysql::traits::MySQLColumn;
31#[cfg(feature = "mysql")]
32use drizzle_mysql::values::MySQLValue;
33
34/// Builder that generates seed INSERT statements for a schema.
35///
36/// Create one with `SeedConfig::sqlite`, `SeedConfig::postgres`, or
37/// `SeedConfig::mysql` (each behind its feature), chain the settings, then
38/// call `generate` / `try_generate`. Tables and columns passed to the
39/// builder must belong to the schema; anything else is a compile error.
40///
41/// Defaults: seed `0`, 10 rows per table, no skipped tables. See the
42/// [crate docs](crate) for how values and row counts are chosen.
43pub struct SeedConfig<'a, D, S> {
44    /// Source schema.
45    pub(crate) schema: &'a S,
46    /// Explicitly skipped tables.
47    pub(crate) skipped_tables: HashSet<TableId>,
48    /// User-provided seed for deterministic RNG.
49    pub(crate) seed: u64,
50    /// Default number of rows per table if not overridden.
51    pub(crate) default_count: usize,
52    /// Per-table row count overrides.
53    pub(crate) table_counts: HashMap<TableId, usize>,
54    /// Per-column generator overrides.
55    pub(crate) column_generators: HashMap<ColumnId, Arc<dyn Generator>>,
56    /// Per-column generator kind overrides.
57    pub(crate) column_kinds: HashMap<ColumnId, GeneratorKind>,
58    /// Relation cardinality overrides. Key: (`parent_table`, `child_table`).
59    pub(crate) relation_counts: HashMap<(TableId, TableId), usize>,
60    /// Optional override for maximum parameters per INSERT statement batch.
61    pub(crate) max_params_per_batch: Option<usize>,
62    /// Unknown names passed to the `*_by_name` methods, reported by
63    /// `try_generate` and `reset_plan`.
64    pub(crate) name_errors: Vec<SeedError>,
65    _dialect: PhantomData<D>,
66    _schema: PhantomData<&'a S>,
67}
68
69impl<'a, D, S> SeedConfig<'a, D, S> {
70    fn with_defaults(schema: &'a S) -> Self {
71        Self {
72            schema,
73            skipped_tables: HashSet::new(),
74            seed: 0,
75            default_count: 10,
76            table_counts: HashMap::new(),
77            column_generators: HashMap::new(),
78            column_kinds: HashMap::new(),
79            relation_counts: HashMap::new(),
80            max_params_per_batch: None,
81            name_errors: Vec::new(),
82            _dialect: PhantomData,
83            _schema: PhantomData,
84        }
85    }
86
87    /// Sets the RNG seed (default `0`). The same seed gives the same rows.
88    #[must_use]
89    pub const fn seed(mut self, seed: u64) -> Self {
90        self.seed = seed;
91        self
92    }
93
94    /// Sets the row count for tables with no explicit count and no seeded
95    /// parent (default 10).
96    #[must_use]
97    pub const fn default_count(mut self, count: usize) -> Self {
98        self.default_count = count;
99        self
100    }
101
102    /// Caps the bind parameters per INSERT; rows are split across more
103    /// statements to stay under it.
104    ///
105    /// Without this, a dialect-specific default is used.
106    ///
107    /// # Panics
108    ///
109    /// Panics if `limit` is zero.
110    #[must_use]
111    pub fn max_params(mut self, limit: usize) -> Self {
112        assert!(limit > 0, "max_params must be > 0");
113        self.max_params_per_batch = Some(limit);
114        self
115    }
116
117    pub(crate) fn count_for(&self, table: TableId) -> usize {
118        self.table_counts
119            .get(&table)
120            .copied()
121            .unwrap_or(self.default_count)
122    }
123}
124
125impl<D, S> SeedConfig<'_, D, S>
126where
127    S: SQLSchemaImpl,
128{
129    pub(crate) fn active_tables(&self) -> Vec<&'static TableRef> {
130        self.schema
131            .table_refs()
132            .iter()
133            .copied()
134            .filter(|t| !self.skipped_tables.contains(&TableId::from_ref(t)))
135            .collect()
136    }
137
138    /// Sets the row count for `table`.
139    ///
140    /// This wins over [`relation`](Self::relation) and
141    /// [`default_count`](Self::default_count()).
142    #[must_use]
143    pub fn count<T>(mut self, table: &T, count: usize) -> Self
144    where
145        T: SQLTableInfo,
146        S: SchemaHasTable<T>,
147    {
148        self.table_counts.insert(TableId::from_info(table), count);
149        self
150    }
151
152    /// Generates `count` rows of `child` for each row of `parent`; each
153    /// parent row gets `count` consecutive child rows.
154    ///
155    /// Without this, a child table with no explicit [`count`](Self::count)
156    /// gets one row per parent row. `child` must have a foreign key to
157    /// `parent` (checked at compile time).
158    #[must_use]
159    pub fn relation<P, C>(mut self, parent: &P, child: &C, count: usize) -> Self
160    where
161        P: SQLTableInfo,
162        C: SQLTableInfo + Relation<P>,
163        S: SchemaHasTable<P> + SchemaHasTable<C>,
164    {
165        self.relation_counts.insert(
166            (TableId::from_info(parent), TableId::from_info(child)),
167            count,
168        );
169        self
170    }
171}
172
173impl<D, S> SeedConfig<'_, D, S> {
174    /// Leaves `table` out: no INSERTs and no reset statements for it.
175    ///
176    /// Foreign keys that point at a skipped table keep their generated (or
177    /// custom) values, for example to reference rows already in the
178    /// database.
179    #[must_use]
180    pub fn skip<T>(mut self, table: &T) -> Self
181    where
182        T: SQLTableInfo,
183        S: SchemaHasTable<T>,
184    {
185        self.skipped_tables.insert(TableId::from_info(table));
186        self
187    }
188}
189
190/// Name-based settings, for a [`Schema`](crate::schema::Schema) built at
191/// runtime or names that come from configuration.
192///
193/// They take the same effect as the typed methods. A table name may be
194/// qualified as `schema.table`, and must be when two namespaces have a
195/// table with that name. An unknown name is reported as a
196/// [`SeedError`](crate::SeedError) by `try_generate`, `try_generate_rows`
197/// and `reset_plan` (and makes `generate` panic).
198impl<D, S> SeedConfig<'_, D, S>
199where
200    S: SQLSchemaImpl,
201{
202    /// Sets the row count for the table named `table`, like
203    /// [`count`](Self::count).
204    #[must_use]
205    pub fn count_by_name(mut self, table: &str, count: usize) -> Self {
206        if let Some(table) = self.resolve_table(table) {
207            self.table_counts.insert(TableId::from_ref(table), count);
208        }
209        self
210    }
211
212    /// Generates `count` rows of `child` for each row of `parent`, like
213    /// [`relation`](Self::relation). `child` must have a foreign key to
214    /// `parent`.
215    #[must_use]
216    pub fn relation_by_name(mut self, parent: &str, child: &str, count: usize) -> Self {
217        let (Some(parent_ref), Some(child_ref)) =
218            (self.resolve_table(parent), self.resolve_table(child))
219        else {
220            return self;
221        };
222        let parent_id = TableId::from_ref(parent_ref);
223        let related = child_ref
224            .foreign_keys
225            .iter()
226            .any(|fk| TableId::foreign_target(child_ref, fk) == parent_id);
227        if related {
228            self.relation_counts
229                .insert((parent_id, TableId::from_ref(child_ref)), count);
230        } else {
231            self.name_errors.push(SeedError::NotRelated {
232                parent: parent_id.to_string(),
233                child: TableId::from_ref(child_ref).to_string(),
234            });
235        }
236        self
237    }
238
239    /// Leaves the table named `table` out, like [`skip`](Self::skip).
240    #[must_use]
241    pub fn skip_by_name(mut self, table: &str) -> Self {
242        if let Some(table) = self.resolve_table(table) {
243            self.skipped_tables.insert(TableId::from_ref(table));
244        }
245        self
246    }
247
248    /// Uses a built-in [`GeneratorKind`] for `column` of `table`, like
249    /// `kind`.
250    #[must_use]
251    pub fn kind_by_name(mut self, table: &str, column: &str, kind: GeneratorKind) -> Self {
252        if let Some(column) = self.resolve_column(table, column) {
253            self.column_kinds.insert(column, kind);
254        }
255        self
256    }
257
258    /// Uses `generator` for `column` of `table`, like `generator`.
259    #[must_use]
260    pub fn generator_by_name(
261        mut self,
262        table: &str,
263        column: &str,
264        generator: impl Generator + 'static,
265    ) -> Self {
266        if let Some(column) = self.resolve_column(table, column) {
267            self.column_generators.insert(column, Arc::new(generator));
268        }
269        self
270    }
271
272    /// Finds the table named `name` (or `schema.name`), recording an error
273    /// when there is no such table or the name is ambiguous.
274    fn resolve_table(&mut self, name: &str) -> Option<&'static TableRef> {
275        let tables = self.schema.table_refs();
276        let exact: Vec<&'static TableRef> = tables
277            .iter()
278            .copied()
279            .filter(|table| table.name == name)
280            .collect();
281        let candidates = if exact.is_empty() {
282            name.split_once('.')
283                .map_or_else(Vec::new, |(schema, table_name)| {
284                    tables
285                        .iter()
286                        .copied()
287                        .filter(|table| {
288                            table.name == table_name
289                                && (table.schema == Some(schema)
290                                    || (schema == "public" && table.schema.is_none()))
291                        })
292                        .collect()
293                })
294        } else {
295            exact
296        };
297        match candidates.as_slice() {
298            [table] => Some(table),
299            [] => {
300                self.name_errors.push(SeedError::UnknownTable {
301                    table: name.to_string(),
302                    known: tables
303                        .iter()
304                        .map(|table| TableId::from_ref(table).to_string())
305                        .collect(),
306                });
307                None
308            }
309            _ => {
310                self.name_errors.push(SeedError::AmbiguousTable {
311                    table: name.to_string(),
312                    candidates: candidates
313                        .iter()
314                        .map(|table| TableId::from_ref(table).to_string())
315                        .collect(),
316                });
317                None
318            }
319        }
320    }
321
322    fn resolve_column(&mut self, table: &str, column: &str) -> Option<ColumnId> {
323        let table = self.resolve_table(table)?;
324        let table_id = TableId::from_ref(table);
325        if let Some(found) = table
326            .columns
327            .iter()
328            .find(|candidate| candidate.name == column)
329        {
330            return Some(ColumnId::new(table_id, found.name));
331        }
332        self.name_errors.push(SeedError::UnknownColumn {
333            table: table_id.to_string(),
334            column: column.to_string(),
335            known: table
336                .columns
337                .iter()
338                .map(|column| column.name.to_string())
339                .collect(),
340        });
341        None
342    }
343
344    /// The first unknown or ambiguous name given to a `*_by_name` method.
345    pub(crate) fn check_names(&self) -> Result<(), SeedError> {
346        self.name_errors.first().cloned().map_or(Ok(()), Err)
347    }
348}
349
350macro_rules! dialect_seed_config {
351    (
352        feature = $feature:literal,
353        dialect = $dialect:literal,
354        marker = $marker:ty,
355        constructor = $constructor:ident,
356        column = $column_trait:path,
357        value = $value:ty,
358        seed_statement = $seed_statement:path,
359        reset_statement = $reset_statement:path,
360        generate = $generate:ident,
361        reset = $reset:ident,
362        reset_note = $reset_note:literal
363    ) => {
364        #[cfg(feature = $feature)]
365        impl<'a> SeedConfig<'a, $marker, ()> {
366            #[doc = concat!("Creates a ", $dialect, " seed config for `schema` (usually a `#[derive(...Schema)]` struct).")]
367            pub fn $constructor<Schema>(schema: &'a Schema) -> SeedConfig<'a, $marker, Schema>
368            where
369                Schema: SQLSchemaImpl,
370            {
371                SeedConfig::<'a, $marker, Schema>::with_defaults(schema)
372            }
373        }
374
375        #[cfg(feature = $feature)]
376        impl<S> SeedConfig<'_, $marker, S>
377        where
378            S: SQLSchemaImpl,
379        {
380            /// Uses a built-in [`GeneratorKind`] for `column` instead of the
381            /// inferred one.
382            #[must_use]
383            pub fn kind<C>(mut self, column: &C, kind: GeneratorKind) -> Self
384            where
385                C: SQLColumnInfo + $column_trait,
386                S: SchemaHasTable<<C as SQLColumn<'static, $value>>::Table>,
387            {
388                self.column_kinds.insert(ColumnId::from_info(column), kind);
389                self
390            }
391
392            /// Uses `generator` for `column`: one from
393            /// [`generators`](crate::generators), a
394            /// [`GeneratorKind`], or your own [`Generator`]. Wins over
395            /// [`kind`](Self::kind).
396            ///
397            /// A foreign key column that references a seeded table is still
398            /// overwritten with the parent row's value.
399            #[must_use]
400            pub fn generator<C>(mut self, column: &C, generator: impl Generator + 'static) -> Self
401            where
402                C: SQLColumnInfo + $column_trait,
403                S: SchemaHasTable<<C as SQLColumn<'static, $value>>::Table>,
404            {
405                self.column_generators
406                    .insert(ColumnId::from_info(column), Arc::new(generator));
407                self
408            }
409
410            /// Generates the INSERT statements for every table not skipped,
411            /// parents first. Large tables are split over several statements
412            /// to stay under the bind-parameter limit. On `PostgreSQL`, a
413            /// table's rows can be followed by a `SELECT setval(...)` that
414            /// moves its `SERIAL`/`IDENTITY` sequences past the seeded ids.
415            ///
416            /// Execute every statement, in order.
417            ///
418            /// # Panics
419            ///
420            /// Panics when [`try_generate`](Self::try_generate) would return
421            /// an error.
422            #[must_use]
423            pub fn generate(&self) -> Vec<$seed_statement> {
424                self.try_generate()
425                    .unwrap_or_else(|error| panic!("invalid seed plan: {error}"))
426            }
427
428            /// Like [`generate`](Self::generate), but returns an error instead
429            /// of panicking.
430            ///
431            /// # Errors
432            ///
433            /// Returns a [`SeedError`](crate::SeedError) when the tables have a
434            /// foreign-key cycle, a `NOT NULL` foreign key has no parent row
435            /// to point at, a value does not fit its column, or one row
436            /// exceeds the parameter limit.
437            pub fn try_generate(&self) -> Result<Vec<$seed_statement>, crate::SeedError> {
438                crate::Seeder::new(self).$generate()
439            }
440
441            /// Generates the statements of [`generate`](Self::generate) as
442            /// one SQL script with every value written inline (see
443            /// `inline_sql` on the statement type): one statement per line,
444            /// each ending in `;`. Save it as a fixture, or run it with any
445            /// client.
446            ///
447            /// # Errors
448            ///
449            /// The same as [`try_generate`](Self::try_generate), plus
450            /// [`SeedError::NoLiteral`](crate::SeedError::NoLiteral) for a
451            /// value with no literal form.
452            pub fn try_generate_script(&self) -> Result<String, crate::SeedError> {
453                let mut script = String::new();
454                for statement in self.try_generate()? {
455                    script.push_str(&statement.inline_sql()?);
456                    script.push_str(";\n");
457                }
458                Ok(script)
459            }
460
461            /// Generates the rows without rendering SQL: per table, the
462            /// column names and one `Vec` of [`SeedValue`](crate::SeedValue)s
463            /// per row, parents first and foreign keys resolved.
464            ///
465            /// Use this to insert with any database driver, or to write
466            /// fixtures. Values follow the same rules as
467            /// [`generate`](Self::generate) for this dialect.
468            ///
469            /// # Errors
470            ///
471            /// The same as [`try_generate`](Self::try_generate), except that
472            /// no parameter limit applies.
473            pub fn try_generate_rows(&self) -> Result<Vec<crate::SeedRows>, crate::SeedError> {
474                crate::Seeder::new(self).generate_rows()
475            }
476
477            /// Returns statements that empty the non-skipped tables, children
478            /// first: `DELETE FROM` each table, preceded by an `UPDATE` that
479            /// clears nullable self-references.
480            ///
481            /// Execute the returned statements in order on one connection.
482            #[doc = $reset_note]
483            ///
484            /// # Errors
485            ///
486            /// Returns a [`SeedError`](crate::SeedError) when the tables have a
487            /// foreign-key cycle, or a table to be emptied is referenced by a
488            /// skipped table.
489            pub fn reset_plan(&self) -> Result<Vec<$reset_statement>, crate::SeedError> {
490                crate::Seeder::new(self).$reset()
491            }
492        }
493    };
494}
495
496dialect_seed_config!(
497    feature = "sqlite",
498    dialect = "SQLite",
499    marker = Sqlite,
500    constructor = sqlite,
501    column = SQLiteColumn<'static>,
502    value = SQLiteValue<'static>,
503    seed_statement = crate::SQLiteSeedStatement,
504    reset_statement = crate::SQLiteResetStatement,
505    generate = generate_sqlite,
506    reset = reset_sqlite,
507    reset_note = "Wrap execution in a transaction when the connection API supports one."
508);
509
510dialect_seed_config!(
511    feature = "postgres",
512    dialect = "PostgreSQL",
513    marker = Postgres,
514    constructor = postgres,
515    column = PostgresColumn<'static>,
516    value = PostgresValue<'static>,
517    seed_statement = crate::PostgresSeedStatement,
518    reset_statement = crate::PostgresResetStatement,
519    generate = generate_postgres,
520    reset = reset_postgres,
521    reset_note = "Wrap execution in a transaction when the connection API supports one."
522);
523
524dialect_seed_config!(
525    feature = "mysql",
526    dialect = "MySQL",
527    marker = MySql,
528    constructor = mysql,
529    column = MySQLColumn<'static>,
530    value = MySQLValue<'static>,
531    seed_statement = crate::MySQLSeedStatement,
532    reset_statement = crate::MySQLResetStatement,
533    generate = generate_mysql,
534    reset = reset_mysql,
535    reset_note = "The plan appends `ALTER TABLE ... AUTO_INCREMENT = 1` for auto-increment tables. Those statements implicitly commit in MySQL, so callers must not assume the whole reset is transactional."
536);
drizzle_seed/config.rs

drizzle_seed/
config.rs
