tempo_cli/db/
migrations.rs

1use anyhow::Result;
2use rusqlite::Connection;
3use std::collections::HashMap;
4
5const MIGRATION_001: &str = include_str!("../../migrations/001_minimal_schema.sql");
6const MIGRATION_002: &str = include_str!("../../migrations/002_features_schema.sql");
7
8pub fn run_migrations(conn: &Connection) -> Result<()> {
9    let current_version = get_current_version(conn)?;
10    let migrations = get_migrations();
11
12    let mut migrations: Vec<_> = migrations.into_iter().collect();
13    migrations.sort_by_key(|(version, _)| *version);
14
15    for (version, sql) in migrations {
16        if version > current_version {
17            log::info!("Running migration {}", version);
18
19            // Run migration in a transaction
20            let tx = conn.unchecked_transaction()?;
21
22            // Execute the SQL as a batch
23            log::debug!("Executing migration SQL: {}", sql);
24            tx.execute_batch(&sql)?;
25
26            // Update schema version
27            tx.execute(
28                "INSERT OR REPLACE INTO schema_version (version) VALUES (?1)",
29                [version],
30            )?;
31
32            tx.commit()?;
33            log::info!("Migration {} completed", version);
34        }
35    }
36
37    Ok(())
38}
39
40fn get_current_version(conn: &Connection) -> Result<i32> {
41    // Check if schema_version table exists
42    let table_exists: bool = conn.query_row(
43        "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='schema_version'",
44        [],
45        |row| row.get::<_, i32>(0),
46    )? > 0;
47
48    if !table_exists {
49        return Ok(0);
50    }
51
52    // Get the current version
53    let version = conn.query_row("SELECT MAX(version) FROM schema_version", [], |row| {
54        let version: Option<i32> = row.get(0)?;
55        Ok(version.unwrap_or(0))
56    })?;
57
58    Ok(version)
59}
60
61fn get_migrations() -> HashMap<i32, String> {
62    let mut migrations = HashMap::new();
63    migrations.insert(1, MIGRATION_001.to_string());
64    migrations.insert(2, MIGRATION_002.to_string());
65    migrations
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71    use rusqlite::Connection;
72
73    #[test]
74    fn test_migrations() {
75        let conn = Connection::open_in_memory()
76            .expect("Failed to create in-memory database for testing");
77        run_migrations(&conn)
78            .expect("Failed to run migrations in test");
79
80        // Verify tables exist
81        let tables: Vec<String> = conn
82            .prepare("SELECT name FROM sqlite_master WHERE type='table'")
83            .expect("Failed to prepare table query")
84            .query_map([], |row| Ok(row.get::<_, String>(0)?))
85            .expect("Failed to execute table query")
86            .collect::<Result<Vec<_>, _>>()
87            .expect("Failed to collect table names");
88
89        assert!(tables.contains(&"projects".to_string()));
90        assert!(tables.contains(&"sessions".to_string()));
91        assert!(tables.contains(&"tags".to_string()));
92        assert!(tables.contains(&"schema_version".to_string()));
93    }
94}