Skip to main content

shared_framework/data/seed/
mod.rs

1//! Database seeding with execution tracking.
2//!
3//! Provides the [`DatabaseSeeder`] trait, the [`DatabaseSeederRunner`] that
4//! runs seeders in dependency order inside a single transaction, the
5//! [`EntitySeeder`] helper for entity-specific seeders, and the
6//! `__database_seeders` tracking table (see [`ensure_table`] and [`entity`]).
7//!
8//! Each seeder runs once: completed names are recorded in
9//! `__database_seeders` and skipped on later runs. Deletion runs in reverse
10//! order and removes the tracking rows.
11//!
12//! ```ignore
13//! use shared_framework::data::seed::DatabaseSeederRunner;
14//!
15//! let runner = DatabaseSeederRunner::new().add(MySeeder);
16//! runner.seed(&db).await?;
17//! ```
18
19pub mod entity;
20
21pub use entity::Entity as SeederEntryEntity;
22pub use entity::ModelEx as SeederEntryModel;
23
24use sea_orm::{
25    ActiveValue::Set, ColumnTrait, ConnectionTrait, DatabaseConnection, DatabaseTransaction,
26    EntityTrait, QueryFilter, TransactionTrait,
27};
28use std::collections::HashSet;
29use std::sync::Arc;
30
31// Re-export entity for external use
32pub use entity::{ActiveModel as SeederEntryActiveModel, ModelEx as SeederEntry};
33
34/// Creates the `__database_seeders` tracking table when missing.
35///
36/// Idempotent `CREATE TABLE IF NOT EXISTS`, callable from the library without
37/// a separate migration run. Returns an error when the DDL fails.
38pub async fn ensure_table(db: &DatabaseConnection) -> anyhow::Result<()> {
39    db.execute_unprepared(
40        r#"CREATE TABLE IF NOT EXISTS "__database_seeders" (
41            "id" BIGSERIAL PRIMARY KEY,
42            "uid" UUID NOT NULL UNIQUE,
43            "seederName" VARCHAR(255) NOT NULL UNIQUE,
44            "createdAt" TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
45            "updatedAt" TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
46        )"#,
47    )
48    .await?;
49    Ok(())
50}
51
52/// Trait for database seeders.
53///
54/// Implementations receive an `Arc<DatabaseTransaction>` so all work runs
55/// inside the runner's single transaction while the handle stays freely
56/// shareable (e.g., for nested repository calls via `RepositoryOptions`).
57/// Do not retain clones beyond the call: committing requires sole ownership.
58#[async_trait::async_trait]
59pub trait DatabaseSeeder: Send + Sync {
60    /// Returns the unique tracking name; defaults to the type name.
61    fn name(&self) -> &str {
62        std::any::type_name::<Self>()
63    }
64    /// Returns the execution order (lower runs first, default 0).
65    fn order(&self) -> i32 {
66        0
67    }
68    /// Inserts seeded data inside the given transaction.
69    async fn seed(&self, txn: Arc<DatabaseTransaction>) -> anyhow::Result<()>;
70    /// Removes seeded data inside the given transaction.
71    async fn delete(&self, txn: Arc<DatabaseTransaction>) -> anyhow::Result<()>;
72}
73
74/// Alias for [`DatabaseSeeder`] kept for existing imports.
75pub use DatabaseSeeder as RepositorySeeder;
76
77/// Helper marking an entity-specific seeder.
78///
79/// `E` is the SeaORM entity being seeded.
80pub struct EntitySeeder<E>
81where
82    E: EntityTrait,
83    E::ModelEx: crate::data::BaseEntity + Send + Sync,
84{
85    _marker: std::marker::PhantomData<E>,
86}
87
88impl<E> EntitySeeder<E>
89where
90    E: EntityTrait,
91    E::ModelEx: crate::data::BaseEntity + Send + Sync,
92{
93    /// Creates an empty helper for the entity type.
94    pub fn new() -> Self {
95        Self { _marker: std::marker::PhantomData }
96    }
97}
98
99/// Runs registered seeders in order inside a single transaction.
100///
101/// Ensures the `__database_seeders` table exists before querying it.
102pub struct DatabaseSeederRunner {
103    seeders: Vec<Box<dyn DatabaseSeeder>>,
104}
105
106impl DatabaseSeederRunner {
107    /// Creates a runner with no seeders.
108    pub fn new() -> Self {
109        Self { seeders: Vec::new() }
110    }
111
112    /// Registers a seeder and returns the runner for chaining.
113    pub fn add<S: DatabaseSeeder + 'static>(mut self, seeder: S) -> Self {
114        self.seeders.push(Box::new(seeder));
115        self
116    }
117
118    /// Registers an already-boxed seeder.
119    pub fn add_boxed(&mut self, seeder: Box<dyn DatabaseSeeder>) {
120        self.seeders.push(seeder);
121    }
122
123    /// Runs all seeders in ascending `order` inside one transaction, recording each name.
124    ///
125    /// Already-recorded seeders are skipped. Does nothing when no seeders are registered.
126    pub async fn seed(&self, db: &DatabaseConnection) -> anyhow::Result<()> {
127        if self.seeders.is_empty() {
128            tracing::info!("No seeders registered; nothing to run");
129            return Ok(());
130        }
131        // Ensure table exists — clean bundling without external migration step
132        ensure_table(db).await?;
133
134        let mut ordered: Vec<&Box<dyn DatabaseSeeder>> = self.seeders.iter().collect();
135        ordered.sort_by_key(|s| s.order());
136
137        let existing = entity::Entity::load().all(db).await?;
138        let executed: HashSet<String> = existing.into_iter().map(|e| e.seeder_name).collect();
139
140        // Single transaction across ALL seeder runs, shared by handle so
141        // seeders can fan the same transaction out to repository calls.
142        let txn = Arc::new(db.begin().await?);
143        let mut in_txn_executed = executed.clone();
144
145        for seeder in ordered {
146            let name = seeder.name().to_string();
147            if in_txn_executed.contains(&name) {
148                tracing::info!(seeder = %name, action = "seed", "Skipping already executed seeder");
149                continue;
150            }
151            tracing::info!(seeder = %name, action = "seed", "Running seeder");
152            seeder.seed(txn.clone()).await.map_err(|e| anyhow::anyhow!("seeder {} failed: {}", name, e))?;
153            let active = entity::ActiveModel {
154                seeder_name: Set(name.clone()),
155                uid: Set(uuid::Uuid::new_v4()),
156                created_at: Set(chrono::Utc::now().into()),
157                updated_at: Set(chrono::Utc::now().into()),
158                ..Default::default()
159            };
160            entity::Entity::insert(active).exec(txn.as_ref()).await.map_err(|e| anyhow::anyhow!("failed to track seeder {}: {}", name, e))?;
161            in_txn_executed.insert(name);
162        }
163
164        let owned = Arc::try_unwrap(txn).map_err(|_| anyhow::anyhow!("seeder transaction handle still shared; refusing to commit"))?;
165        owned.commit().await?;
166        tracing::info!(action = "seed", seeder_count = in_txn_executed.len(), "Seeder transaction committed");
167        Ok(())
168    }
169
170    /// Deletes seeded data in descending `order` inside one transaction, removing each tracking row.
171    ///
172    /// Seeders without a tracking row are skipped. Does nothing when no seeders are registered.
173    pub async fn delete(&self, db: &DatabaseConnection) -> anyhow::Result<()> {
174        if self.seeders.is_empty() {
175            tracing::info!("No seeders registered; nothing to delete");
176            return Ok(());
177        }
178        ensure_table(db).await?;
179
180        let mut ordered: Vec<&Box<dyn DatabaseSeeder>> = self.seeders.iter().collect();
181        ordered.sort_by_key(|s| std::cmp::Reverse(s.order()));
182
183        let existing = entity::Entity::load().all(db).await?;
184        let executed: HashSet<String> = existing.into_iter().map(|e| e.seeder_name).collect();
185
186        let txn = Arc::new(db.begin().await?);
187        let mut in_txn_executed = executed.clone();
188
189        for seeder in ordered {
190            let name = seeder.name().to_string();
191            if !in_txn_executed.contains(&name) {
192                tracing::info!(seeder = %name, action = "delete", "Skipping seeder that has not been executed");
193                continue;
194            }
195            tracing::info!(seeder = %name, action = "delete", "Deleting seeded data");
196            seeder.delete(txn.clone()).await.map_err(|e| anyhow::anyhow!("seeder delete {} failed: {}", name, e))?;
197            entity::Entity::delete_many()
198                .filter(entity::Column::SeederName.eq(name.clone()))
199                .exec(txn.as_ref())
200                .await?;
201            in_txn_executed.remove(&name);
202        }
203
204        let owned = Arc::try_unwrap(txn).map_err(|_| anyhow::anyhow!("seeder transaction handle still shared; refusing to commit"))?;
205        owned.commit().await?;
206        tracing::info!(action = "delete", "Seeder deletion transaction committed");
207        Ok(())
208    }
209}
210
211impl Default for DatabaseSeederRunner {
212    fn default() -> Self { Self::new() }
213}
214
215/// Alias for [`DatabaseSeederRunner`].
216pub type DatabaseSeederHelper = DatabaseSeederRunner;
217
218#[deprecated(note = "Use DatabaseSeederRunner or DatabaseSeederHelper")]
219/// Deprecated alias for [`DatabaseSeederRunner`].
220pub type SeederRunner = DatabaseSeederRunner;