tempo_cli/db/
seed.rs

1use crate::db::queries::{ProjectQueries, SessionQueries};
2use crate::models::{Project, Session, SessionContext};
3use anyhow::Result;
4use chrono::{Duration, Utc};
5use rand::Rng;
6use rusqlite::Connection;
7use std::path::PathBuf;
8
9pub fn seed_database(conn: &Connection) -> Result<()> {
10    println!("Seeding database...");
11
12    // 1. Clear existing data
13    conn.execute("DELETE FROM sessions", [])?;
14    conn.execute("DELETE FROM projects", [])?;
15    conn.execute("DELETE FROM tags", [])?;
16    conn.execute("DELETE FROM goals", [])?;
17
18    // Reset auto-increment counters
19    conn.execute("DELETE FROM sqlite_sequence WHERE name='sessions'", [])?;
20    conn.execute("DELETE FROM sqlite_sequence WHERE name='projects'", [])?;
21    conn.execute("DELETE FROM sqlite_sequence WHERE name='tags'", [])?;
22    conn.execute("DELETE FROM sqlite_sequence WHERE name='goals'", [])?;
23
24    println!("Database cleared.");
25
26    // 2. Create dummy projects
27    let projects = vec![
28        (
29            "LLM Fine-tuning",
30            "Fine-tuning Llama 3 on custom dataset",
31            "/Users/darylwanji/projects/llm-finetune",
32        ),
33        (
34            "RAG Pipeline",
35            "Retrieval Augmented Generation for internal docs",
36            "/Users/darylwanji/projects/rag-pipeline",
37        ),
38        (
39            "Computer Vision Model",
40            "Object detection for autonomous drone",
41            "/Users/darylwanji/projects/cv-drone",
42        ),
43        (
44            "Reinforcement Learning Agent",
45            "PPO agent for trading bot",
46            "/Users/darylwanji/projects/rl-trading",
47        ),
48        (
49            "Data Pipeline ETL",
50            "ETL pipeline for customer data",
51            "/Users/darylwanji/projects/etl-pipeline",
52        ),
53    ];
54
55    let mut project_ids = Vec::new();
56
57    for (name, desc, path) in projects {
58        let project = Project {
59            id: None,
60            name: name.to_string(),
61            path: PathBuf::from(path),
62            git_hash: Some("a1b2c3d".to_string()), // Dummy hash
63            description: Some(desc.to_string()),
64            created_at: Utc::now(),
65            updated_at: Utc::now(),
66            is_archived: false,
67        };
68
69        let id = ProjectQueries::create(conn, &project)?;
70        project_ids.push(id);
71        println!("Created project: {}", name);
72    }
73
74    // 3. Create dummy sessions
75    let mut rng = rand::thread_rng();
76    let contexts = vec![
77        SessionContext::Terminal,
78        SessionContext::IDE,
79        SessionContext::Linked,
80        SessionContext::Manual,
81    ];
82    let notes = vec![
83        "Debugging model convergence",
84        "Implementing attention mechanism",
85        "Optimizing data loader",
86        "Refactoring training loop",
87        "Writing documentation",
88        "Code review",
89        "Meeting with team",
90        "Testing inference speed",
91    ];
92
93    for &project_id in &project_ids {
94        // Create 10-20 sessions per project
95        let num_sessions = rng.gen_range(10..21);
96
97        for _ in 0..num_sessions {
98            // Random start time in the last 30 days
99            let days_ago = rng.gen_range(0..30);
100            let hours_ago = rng.gen_range(0..24);
101            let start_time = Utc::now() - Duration::days(days_ago) - Duration::hours(hours_ago);
102
103            // Random duration between 15 mins and 4 hours
104            let duration_mins = rng.gen_range(15..240);
105            let end_time = start_time + Duration::minutes(duration_mins);
106
107            let context = contexts[rng.gen_range(0..contexts.len())];
108            let note = if rng.gen_bool(0.3) {
109                Some(notes[rng.gen_range(0..notes.len())].to_string())
110            } else {
111                None
112            };
113
114            let session = Session {
115                id: None,
116                project_id,
117                start_time,
118                end_time: Some(end_time),
119                context,
120                paused_duration: Duration::zero(),
121                notes: note,
122                created_at: start_time,
123            };
124
125            SessionQueries::create(conn, &session)?;
126        }
127        println!(
128            "Created {} sessions for project {}",
129            num_sessions, project_id
130        );
131    }
132
133    println!("Database seeding completed successfully.");
134    Ok(())
135}
tempo_cli/db/seed.rs

tempo_cli/db/
seed.rs
