1use anyhow::Result;
4use sqlx::{PgPool, Row};
5use std::fs;
6use std::path::Path;
7use colored::*;
8use serde::{Deserialize, Serialize};
9use chrono::{DateTime, Utc};
10
11pub struct SeedManager {
13 pool: PgPool,
14}
15
16impl SeedManager {
17 pub fn new(pool: PgPool) -> Self {
18 Self { pool }
19 }
20
21 pub fn pool(&self) -> &PgPool {
23 &self.pool
24 }
25
26 pub async fn seed(&self) -> Result<()> {
28 println!("đą Running database seeds...");
29
30 self.create_seeds_table().await?;
32
33 let seeds = self.load_seed_files()?;
35
36 let applied_seeds = self.get_applied_seeds().await?;
38
39 for seed in seeds {
41 if !applied_seeds.contains(&seed.name) {
42 println!(" âŗ Applying seed: {}", seed.name.bright_green());
43 self.apply_seed(&seed).await?;
44 } else {
45 println!(" â Skipping already applied: {}", seed.name);
46 }
47 }
48
49 println!("â
Seeding completed successfully");
50 Ok(())
51 }
52
53 pub async fn seed_by_name(&self, seed_name: &str) -> Result<()> {
55 println!("đą Running specific seed: {}", seed_name.bright_cyan());
56
57 self.create_seeds_table().await?;
59
60 let seeds = self.load_seed_files()?;
62 let seed = seeds.iter().find(|s| s.name == seed_name);
63
64 match seed {
65 Some(seed) => {
66 println!(" âŗ Applying seed: {}", seed.name.bright_green());
67 self.apply_seed(seed).await?;
68 println!("â
Seed '{}' completed successfully", seed_name);
69 }
70 None => {
71 return Err(anyhow::anyhow!("Seed '{}' not found", seed_name));
72 }
73 }
74
75 Ok(())
76 }
77
78 pub async fn revert_seeds(&self) -> Result<()> {
80 println!("đ Reverting all seeds...");
81
82 self.create_seeds_table().await?;
84
85 let applied_seeds = self.get_applied_seeds_reversed().await?;
87
88 if applied_seeds.is_empty() {
89 println!(" âšī¸ No seeds to revert");
90 return Ok(());
91 }
92
93 for seed_name in applied_seeds {
94 println!(" âŗ Reverting seed: {}", seed_name.bright_yellow());
95
96 let revert_seed = self.find_revert_seed(&seed_name)?;
98
99 if let Some(revert_content) = revert_seed {
100 self.revert_seed(&seed_name, &revert_content).await?;
101 println!(" â Successfully reverted: {}", seed_name.bright_green());
102 } else {
103 println!(" â ī¸ No revert seed found for: {}", seed_name.yellow());
104 }
105 }
106
107 println!("â
Seed reversion completed");
108 Ok(())
109 }
110
111 pub async fn create_seed(&self, name: &str, seed_type: SeedType) -> Result<()> {
113 println!("đ Creating new seed: {}", name.bright_cyan());
114
115 let timestamp = chrono::Utc::now().format("%Y%m%d_%H%M%S");
116 let seed_name = format!("{}_{}.sql", timestamp, name);
117 let seeds_dir = "seeds";
118
119 fs::create_dir_all(seeds_dir)?;
121
122 let content = match seed_type {
123 SeedType::Data => self.generate_data_seed_template(name),
124 SeedType::Test => self.generate_test_seed_template(name),
125 SeedType::Reference => self.generate_reference_seed_template(name),
126 };
127
128 let seed_path = Path::new(seeds_dir).join(&seed_name);
129 fs::write(&seed_path, content)?;
130
131 println!(" â Created: {}", seed_path.display().to_string().bright_green());
132
133 let revert_name = format!("{}_{}_revert.sql", timestamp, name);
135 let revert_content = format!(
136 r#"-- Revert Seed: {name}
137-- Created: {timestamp}
138-- Type: {seed_type:?}
139
140-- Add your revert SQL here
141-- Example:
142-- DELETE FROM users WHERE email LIKE '%@test.local';
143"#,
144 name = name,
145 timestamp = timestamp,
146 seed_type = seed_type
147 );
148
149 let revert_path = Path::new(seeds_dir).join(&revert_name);
150 fs::write(&revert_path, revert_content)?;
151
152 println!(" â Created: {}", revert_path.display().to_string().bright_green());
153
154 Ok(())
155 }
156
157 pub async fn history(&self) -> Result<Vec<SeedRecord>> {
159 self.create_seeds_table().await?;
161
162 let rows = sqlx::query("SELECT id, name, seed_type, applied_at FROM schema_seeds ORDER BY applied_at DESC")
163 .fetch_all(&self.pool)
164 .await?;
165
166 let records: Result<Vec<SeedRecord>, sqlx::Error> = rows.into_iter()
167 .map(|row| {
168 Ok(SeedRecord {
169 id: row.get("id"),
170 name: row.get("name"),
171 seed_type: row.get("seed_type"),
172 applied_at: row.get("applied_at"),
173 })
174 })
175 .collect();
176
177 match records {
178 Ok(history) => {
179 if history.is_empty() {
180 println!(" âšī¸ No seeds have been applied yet");
181 } else {
182 println!("đ Seed History:");
183 for (i, record) in history.iter().enumerate() {
184 let status = if i == 0 {
185 "đš".to_string() } else {
187 " ".to_string()
188 };
189 println!(" {} {} [{}] - Applied at {}",
190 status,
191 record.name.bright_cyan(),
192 record.seed_type.bright_yellow(),
193 record.applied_at.format("%Y-%m-%d %H:%M:%S")
194 );
195 }
196 }
197 Ok(history)
198 }
199 Err(e) => Err(anyhow::anyhow!("Failed to load seed history: {}", e))
200 }
201 }
202
203 pub async fn status(&self) -> Result<SeedStatus> {
205 self.create_seeds_table().await?;
207
208 let seed_files = self.load_seed_files()?;
210
211 let applied_seeds = self.get_applied_seeds().await?;
213
214 let mut pending_seeds = Vec::new();
216 let mut applied_seed_details = Vec::new();
217
218 for seed_file in seed_files {
219 if applied_seeds.contains(&seed_file.name) {
220 applied_seed_details.push(seed_file.name);
221 } else {
222 pending_seeds.push(seed_file.name);
223 }
224 }
225
226 let status = SeedStatus {
227 total_seeds: applied_seed_details.len() + pending_seeds.len(),
228 applied_seeds: applied_seed_details,
229 pending_seeds,
230 last_seed: applied_seeds.last().cloned(),
231 };
232
233 println!("đ Seed Status:");
235 println!(" Total seeds: {}", status.total_seeds);
236 println!(" Applied seeds: {}", status.applied_seeds.len());
237 println!(" Pending seeds: {}", status.pending_seeds.len());
238
239 if !status.pending_seeds.is_empty() {
240 println!(" Pending:");
241 for seed in &status.pending_seeds {
242 println!(" - {}", seed.yellow());
243 }
244 }
245
246 Ok(status)
247 }
248
249 pub async fn load_seed_data<T>(&self, file_path: &str) -> Result<Vec<T>>
251 where
252 T: for<'de> Deserialize<'de> + Send,
253 {
254 let content = fs::read_to_string(file_path)?;
255 let data: Vec<T> = serde_json::from_str(&content)?;
256 Ok(data)
257 }
258
259 pub async fn bulk_insert(&self, table_name: &str, data: &[serde_json::Value]) -> Result<u64> {
261 if data.is_empty() {
262 return Ok(0);
263 }
264
265 let mut total_rows = 0;
266
267 let chunk_size = 1000;
269 for chunk in data.chunks(chunk_size) {
270 let sample_value = &chunk[0];
271 let fields: Vec<String> = sample_value.as_object()
272 .ok_or_else(|| anyhow::anyhow!("Data must be objects"))?
273 .keys()
274 .cloned()
275 .collect();
276
277 let placeholders: Vec<String> = chunk.iter()
279 .enumerate()
280 .map(|(i, _)| {
281 let value_placeholders: Vec<String> = (0..fields.len())
282 .map(|j| format!("${}", i * fields.len() + j + 1))
283 .collect();
284 format!("({})", value_placeholders.join(", "))
285 })
286 .collect();
287
288 let query = format!(
289 "INSERT INTO {} ({}) VALUES {}",
290 table_name,
291 fields.join(", "),
292 placeholders.join(", ")
293 );
294
295 let mut query_builder = sqlx::query(&query);
296
297 for item in chunk {
299 for field in &fields {
300 if let Some(value) = item.get(field) {
301 if let Some(str_val) = value.as_str() {
302 query_builder = query_builder.bind(str_val);
303 } else if let Some(int_val) = value.as_i64() {
304 query_builder = query_builder.bind(int_val);
305 } else if let Some(float_val) = value.as_f64() {
306 query_builder = query_builder.bind(float_val);
307 } else if let Some(bool_val) = value.as_bool() {
308 query_builder = query_builder.bind(bool_val);
309 } else {
310 query_builder = query_builder.bind(serde_json::to_string(value)?);
311 }
312 } else {
313 query_builder = query_builder.bind::<Option<String>>(None);
314 }
315 }
316 }
317
318 let result = query_builder.execute(&self.pool).await?;
319 total_rows += result.rows_affected();
320 }
321
322 Ok(total_rows)
323 }
324
325 async fn create_seeds_table(&self) -> Result<()> {
327 sqlx::query(
328 r#"
329 CREATE TABLE IF NOT EXISTS schema_seeds (
330 id SERIAL PRIMARY KEY,
331 name VARCHAR(255) NOT NULL UNIQUE,
332 seed_type VARCHAR(50) NOT NULL,
333 applied_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
334 )
335 "#
336 )
337 .execute(&self.pool)
338 .await?;
339
340 Ok(())
341 }
342
343 fn load_seed_files(&self) -> Result<Vec<SeedFile>> {
345 let seeds_dir = Path::new("seeds");
346 let mut seeds = Vec::new();
347
348 if seeds_dir.exists() {
349 for entry in fs::read_dir(seeds_dir)? {
350 let entry = entry?;
351 let path = entry.path();
352
353 if path.extension().and_then(|s| s.to_str()) == Some("sql") &&
354 !path.file_name().unwrap().to_string_lossy().contains("_revert") {
355
356 let name = path.file_stem()
357 .unwrap()
358 .to_string_lossy()
359 .to_string();
360
361 let content = fs::read_to_string(&path)?;
362
363 let seed_type = self.determine_seed_type(&name, &content);
365
366 seeds.push(SeedFile {
367 name: name.clone(),
368 content,
369 seed_type,
370 });
371 }
372 }
373 }
374
375 seeds.sort_by(|a, b| a.name.cmp(&b.name));
376 Ok(seeds)
377 }
378
379 fn determine_seed_type(&self, name: &str, content: &str) -> String {
381 if name.to_lowercase().contains("test") || content.to_lowercase().contains("test") {
382 "test".to_string()
383 } else if name.to_lowercase().contains("ref") || content.to_lowercase().contains("reference") {
384 "reference".to_string()
385 } else {
386 "data".to_string()
387 }
388 }
389
390 async fn get_applied_seeds(&self) -> Result<Vec<String>> {
392 let rows = sqlx::query("SELECT name FROM schema_seeds ORDER BY name")
393 .fetch_all(&self.pool)
394 .await?;
395
396 let seeds: Vec<String> = rows
397 .into_iter()
398 .map(|row| row.get::<String, _>("name"))
399 .collect();
400
401 Ok(seeds)
402 }
403
404 async fn get_applied_seeds_reversed(&self) -> Result<Vec<String>> {
406 let rows = sqlx::query("SELECT name FROM schema_seeds ORDER BY applied_at DESC")
407 .fetch_all(&self.pool)
408 .await?;
409
410 let seeds: Vec<String> = rows
411 .into_iter()
412 .map(|row| row.get::<String, _>("name"))
413 .collect();
414
415 Ok(seeds)
416 }
417
418 async fn apply_seed(&self, seed: &SeedFile) -> Result<()> {
420 let mut tx = self.pool.begin().await?;
421
422 sqlx::query(&seed.content)
424 .execute(&mut *tx)
425 .await?;
426
427 sqlx::query("INSERT INTO schema_seeds (name, seed_type) VALUES ($1, $2)")
429 .bind(&seed.name)
430 .bind(&seed.seed_type)
431 .execute(&mut *tx)
432 .await?;
433
434 tx.commit().await?;
435
436 Ok(())
437 }
438
439 fn find_revert_seed(&self, seed_name: &str) -> Result<Option<String>> {
441 let seeds_dir = Path::new("seeds");
442
443 let revert_name = format!("{}_revert", seed_name);
445
446 if seeds_dir.exists() {
447 for entry in fs::read_dir(seeds_dir)? {
448 let entry = entry?;
449 let path = entry.path();
450
451 if path.extension().and_then(|s| s.to_str()) == Some("sql") {
452 let file_stem = path.file_stem()
453 .unwrap()
454 .to_string_lossy()
455 .to_string();
456
457 if file_stem == revert_name {
458 let content = fs::read_to_string(&path)?;
459 return Ok(Some(content));
460 }
461 }
462 }
463 }
464
465 Ok(None)
466 }
467
468 async fn revert_seed(&self, seed_name: &str, revert_content: &str) -> Result<()> {
470 let mut tx = self.pool.begin().await?;
471
472 sqlx::query(revert_content)
474 .execute(&mut *tx)
475 .await?;
476
477 sqlx::query("DELETE FROM schema_seeds WHERE name = $1")
479 .bind(seed_name)
480 .execute(&mut *tx)
481 .await?;
482
483 tx.commit().await?;
484
485 Ok(())
486 }
487
488 fn generate_data_seed_template(&self, name: &str) -> String {
490 format!(
491 r#"-- Data Seed: {name}
492-- Type: Data
493-- Description: Initial data seed for production/staging
494
495-- Insert initial data here
496-- Example:
497-- INSERT INTO users (id, name, email, created_at, updated_at) VALUES
498-- ('admin-id', 'System Administrator', 'admin@company.com', NOW(), NOW()),
499-- ('user-id', 'Test User', 'user@company.com', NOW(), NOW());
500
501-- You can also use INSERT with multiple rows:
502-- INSERT INTO categories (id, name, description) VALUES
503-- ('cat1', 'Electronics', 'Electronic devices and accessories'),
504-- ('cat2', 'Books', 'Books and educational materials'),
505-- ('cat3', 'Clothing', 'Apparel and fashion items');
506
507-- For complex data, consider using JSON:
508-- INSERT INTO settings (key, value) VALUES
509-- ('app_config', '{{"theme": "dark", "language": "en", "timezone": "UTC"}}'),
510-- ('feature_flags', '{{"new_ui": true, "beta_features": false, "analytics": true}}');
511"#,
512 name = name
513 )
514 }
515
516 fn generate_test_seed_template(&self, name: &str) -> String {
518 format!(
519 r#"-- Test Seed: {name}
520-- Type: Test
521-- Description: Test data for development and testing
522
523-- Insert test data here
524-- Use identifiable test data with consistent patterns
525
526-- Example test users
527-- INSERT INTO users (id, name, email, created_at, updated_at) VALUES
528-- ('test-user-1', 'Test User 1', 'user1@test.local', NOW(), NOW()),
529-- ('test-user-2', 'Test User 2', 'user2@test.local', NOW(), NOW()),
530-- ('test-admin', 'Test Admin', 'admin@test.local', NOW(), NOW());
531
532-- Test categories
533-- INSERT INTO categories (id, name, description) VALUES
534-- ('test-cat-1', 'Test Category 1', 'Test category for testing'),
535-- ('test-cat-2', 'Test Category 2', 'Another test category');
536
537-- Test products
538-- INSERT INTO products (id, name, price, category_id, active, created_at, updated_at) VALUES
539-- ('test-prod-1', 'Test Product 1', 99.99, 'test-cat-1', true, NOW(), NOW()),
540-- ('test-prod-2', 'Test Product 2', 149.99, 'test-cat-1', false, NOW(), NOW()),
541-- ('test-prod-3', 'Test Product 3', 199.99, 'test-cat-2', true, NOW(), NOW());
542
543-- Note: Test data should use identifiable patterns like:
544-- - test-local domain (@test.local)
545-- - test- prefixes (test-user-1, test-cat-1, etc.)
546-- - Consistent test values for reproducibility
547"#,
548 name = name
549 )
550 }
551
552 fn generate_reference_seed_template(&self, name: &str) -> String {
554 format!(
555 r#"-- Reference Seed: {name}
556-- Type: Reference
557-- Description: Reference data and lookup tables
558
559-- Insert reference data here
560-- This is typically static data that applications depend on
561
562-- Example reference data
563-- INSERT INTO countries (id, code, name, iso3, currency) VALUES
564-- ('US', 'US', 'United States', 'USA', 'USD'),
565-- ('CA', 'CA', 'Canada', 'CAN', 'CAD'),
566-- ('GB', 'GB', 'United Kingdom', 'GBR', 'GBP');
567
568-- INSERT INTO languages (id, code, name, is_active) VALUES
569-- ('en', 'en', 'English', true),
570-- ('es', 'es', 'Spanish', true),
571-- ('fr', 'fr', 'French', true),
572-- ('de', 'de', 'German', true);
573
574-- INSERT INTO user_roles (id, name, description, permissions) VALUES
575-- ('admin', 'Administrator', 'Full system access', '["create", "read", "update", "delete", "admin"]'),
576-- ('moderator', 'Moderator', 'Content moderation access', '["read", "update", "moderate"]'),
577-- ('user', 'User', 'Standard user access', '["read", "update_profile"]');
578
579-- INSERT INTO system_settings (id, key, value, description, is_public) VALUES
580-- ('1', 'app_name', 'My Application', 'Application name', true),
581-- ('2', 'app_version', '1.0.0', 'Current application version', true),
582-- ('3', 'max_file_size', '10485760', 'Maximum file size in bytes', false),
583-- ('4', 'session_timeout', '3600', 'Session timeout in seconds', false);
584
585-- Reference data characteristics:
586-- - Static and rarely changes
587-- - Required for application functionality
588-- - Often has foreign key relationships
589-- - Usually loaded early in the application lifecycle
590"#,
591 name = name
592 )
593 }
594}
595
596struct SeedFile {
598 name: String,
599 content: String,
600 seed_type: String,
601}
602
603#[derive(Debug, Clone)]
605pub struct SeedRecord {
606 pub id: i32,
607 pub name: String,
608 pub seed_type: String,
609 pub applied_at: DateTime<Utc>,
610}
611
612#[derive(Debug, Clone)]
614pub struct SeedStatus {
615 pub total_seeds: usize,
616 pub applied_seeds: Vec<String>,
617 pub pending_seeds: Vec<String>,
618 pub last_seed: Option<String>,
619}
620
621#[derive(Debug, Clone, Serialize, Deserialize)]
623pub enum SeedType {
624 Data,
625 Test,
626 Reference,
627}
628
629pub trait Seed {
631 fn name(&self) -> &str;
632 fn seed_type(&self) -> SeedType;
633 fn up(&self) -> &str;
634 fn down(&self) -> &str;
635}