shared_framework/data/seed/
mod.rs1pub 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
31pub use entity::{ActiveModel as SeederEntryActiveModel, ModelEx as SeederEntry};
33
34pub 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#[async_trait::async_trait]
59pub trait DatabaseSeeder: Send + Sync {
60 fn name(&self) -> &str {
62 std::any::type_name::<Self>()
63 }
64 fn order(&self) -> i32 {
66 0
67 }
68 async fn seed(&self, txn: Arc<DatabaseTransaction>) -> anyhow::Result<()>;
70 async fn delete(&self, txn: Arc<DatabaseTransaction>) -> anyhow::Result<()>;
72}
73
74pub use DatabaseSeeder as RepositorySeeder;
76
77pub 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 pub fn new() -> Self {
95 Self { _marker: std::marker::PhantomData }
96 }
97}
98
99pub struct DatabaseSeederRunner {
103 seeders: Vec<Box<dyn DatabaseSeeder>>,
104}
105
106impl DatabaseSeederRunner {
107 pub fn new() -> Self {
109 Self { seeders: Vec::new() }
110 }
111
112 pub fn add<S: DatabaseSeeder + 'static>(mut self, seeder: S) -> Self {
114 self.seeders.push(Box::new(seeder));
115 self
116 }
117
118 pub fn add_boxed(&mut self, seeder: Box<dyn DatabaseSeeder>) {
120 self.seeders.push(seeder);
121 }
122
123 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(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 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 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
215pub type DatabaseSeederHelper = DatabaseSeederRunner;
217
218#[deprecated(note = "Use DatabaseSeederRunner or DatabaseSeederHelper")]
219pub type SeederRunner = DatabaseSeederRunner;