Put some state about the migration INTO the database

glean-core/src/database/migration.rs

@@ -14,6 +14,7 @@ use crate::Lifetime;
 use crate::Result;
 
 use rkv::StoreOptions;
+use rusqlite::OptionalExtension;
 use rusqlite::Transaction;
 
 use super::sqlite;
@@ -264,6 +265,23 @@ pub fn try_migrate(data_path: &Path, db: &sqlite::Database) -> Result<Option<Mig
         return Ok(None);
     }
 
+    {
+        let conn = db.conn.lock();
+        let migration_status: Option<String> = conn
+            .query_row("SELECT state FROM migration", [], |row| row.get(0))
+            .optional()?;
+
+        if matches!(migration_status.as_deref(), Some("done")) {
+            log::debug!("Rkv already successfully migrated. Not doing it again.");
+            return Ok(None);
+        }
+    }
+
+    db.conn.write(|tx| {
+        tx.execute("INSERT INTO migration (id, state) VALUES (1, 'started') ON CONFLICT(id) DO UPDATE SET state = excluded.state", [])?;
+        Ok::<(), rusqlite::Error>(())
+    })?;
+
     let rkv = RkvDatabase::new(data_path)?;
 
     let state = db.conn.write(|tx| {

glean-core/src/database/sqlite.rs

@@ -158,7 +158,6 @@ impl Database {
 
         fs::create_dir_all(&path)?;
         let store_path = path.join(DEFAULT_DATABASE_FILE_NAME);
-        let sqlite_exists = store_path.exists();
         let (conn, load_state) = sqlite_open(&store_path)?;
 
         let mut db = Self {
@@ -169,26 +168,27 @@ impl Database {
             migration_error: MigrationResult::Unknown,
         };
 
-        if sqlite_exists {
-            log::debug!("SQLite database already exists. Not trying to migrate Rkv");
-            db.run_maintenance(false)?;
-        } else {
-            match migration::try_migrate(&path, &db) {
-                Ok(Some(state)) => {
-                    log::debug!("Migration done. state={state:?}");
-                    db.migration_state = Some(state);
-                    db.run_maintenance(true)?;
-                }
-                Ok(None) => {
-                    log::debug!("No migration.");
-                }
-                Err(e) => {
-                    db.migration_error = MigrationResult::Error;
-                    log::warn!("Migration failed! Continuing with SQLite backend without migrated data. Error: {e:?}")
-                }
+        match migration::try_migrate(&path, &db) {
+            Ok(Some(state)) => {
+                log::debug!("Migration done. state={state:?}");
+                db.migration_state = Some(state);
+                db.run_maintenance(true)?;
+            }
+            Ok(None) => {
+                log::debug!("No migration.");
+                db.run_maintenance(false)?;
+            }
+            Err(e) => {
+                db.migration_error = MigrationResult::Error;
+                log::warn!("Migration failed! Continuing with SQLite backend without migrated data. Error: {e:?}")
             }
         }
 
+        db.conn.write(|tx| {
+            tx.execute("INSERT INTO migration (id, state) VALUES (1, 'done') ON CONFLICT(id) DO UPDATE SET state = excluded.state", [])?;
+            Ok::<(), rusqlite::Error>(())
+        })?;
+
         Ok(db)
     }
 

glean-core/src/database/sqlite/schema.rs

@@ -6,7 +6,7 @@
 
 use std::num::NonZeroU32;
 
-use rusqlite::{config::DbConfig, Transaction};
+use rusqlite::{config::DbConfig, OptionalExtension, Transaction};
 
 use super::connection::ConnectionOpener;
 
@@ -16,7 +16,7 @@ use super::connection::ConnectionOpener;
 pub struct Schema;
 
 impl ConnectionOpener for Schema {
-    const MAX_SCHEMA_VERSION: u32 = 1;
+    const MAX_SCHEMA_VERSION: u32 = 2;
 
     type Error = SchemaError;
 
@@ -50,20 +50,48 @@ impl ConnectionOpener for Schema {
 
     fn create(tx: &mut Transaction<'_>) -> Result<(), Self::Error> {
         tx.execute_batch(
-            "CREATE TABLE telemetry(
+            "
+             CREATE TABLE telemetry(
                id TEXT NOT NULL,
                ping TEXT NOT NULL,
                lifetime TEXT NOT NULL,
                labels TEXT NOT NULL, -- can't be null or ON CONFLICT won't work
                value BLOB,
                UNIQUE(id, ping, labels)
-             );",
+             );
+             CREATE TABLE migration(id INTEGER PRIMARY KEY, state TEXT NOT NULL);
+            ",
         )?;
         Ok(())
     }
 
-    fn upgrade(_: &mut Transaction<'_>, to_version: NonZeroU32) -> Result<(), Self::Error> {
-        Err(SchemaError::UnsupportedSchemaVersion(to_version.get()))
+    fn upgrade(tx: &mut Transaction<'_>, to_version: NonZeroU32) -> Result<(), Self::Error> {
+        match to_version.get() {
+            2 => {
+                log::info!("Upgrading user_version to 2");
+                // Clients upgrading to schema 2 don't have the table.
+                // But they did run through the migration.
+                tx.execute_batch(
+                    "CREATE TABLE migration(
+                        id INTEGER PRIMARY KEY,
+                        state TEXT NOT NULL
+                    );",
+                )?;
+                let cid_exists: Option<i32> = tx
+                    .query_row(
+                        "SELECT 1 FROM telemetry WHERE id = 'client_id'",
+                        [],
+                        |row| row.get(0),
+                    )
+                    .optional()?;
+                if cid_exists.is_some() {
+                    log::info!("Client ID already exists. Marking migration as done.");
+                    tx.execute("INSERT INTO migration (id, state) VALUES (1, 'done') ON CONFLICT(id) DO UPDATE SET state = excluded.state", [])?;
+                }
+                Ok(())
+            }
+            to_version => Err(SchemaError::UnsupportedSchemaVersion(to_version)),
+        }
     }
 }
 

glean-core/tests/sqlite.rs

@@ -146,7 +146,7 @@ fn higher_user_version_upgrade_does_not_crash() {
     {
         let path = temp.path().join("db").join("glean.sqlite");
         let conn = rusqlite::Connection::open(path).unwrap();
-        conn.execute_batch("PRAGMA user_version = 2").unwrap();
+        conn.execute_batch("PRAGMA user_version = 99").unwrap();
     }
 
     let (glean, _temp) = new_glean(Some(temp));
@@ -230,3 +230,39 @@ fn database_externally_locked() {
     let client_id = clientid_metric().get_value(&glean, None);
     assert!(client_id.is_some());
 }
+
+#[test]
+fn schema_v2_is_applied() {
+    let (first_client_id, temp) = {
+        let (glean, temp) = new_glean(None);
+        let client_id = clientid_metric().get_value(&glean, None).unwrap();
+        drop(glean);
+
+        let db_path = temp.path().join("db").join("glean.sqlite");
+        let conn = rusqlite::Connection::open(db_path).unwrap();
+
+        conn.execute("DROP TABLE migration", []).unwrap();
+
+        // Reset to first schema version
+        conn.execute("PRAGMA user_version = 1", []).unwrap();
+
+        (client_id, temp)
+    };
+
+    let db_path = temp.path().join("db").join("glean.sqlite");
+    let (glean, _temp) = new_glean(Some(temp));
+
+    let client_id = clientid_metric().get_value(&glean, None).unwrap();
+    assert_eq!(first_client_id, client_id);
+
+    let conn = rusqlite::Connection::open(db_path).unwrap();
+    let cur_user_version: u32 = conn
+        .query_one("PRAGMA user_version", [], |row| row.get(0))
+        .unwrap();
+    assert_eq!(cur_user_version, 2);
+
+    let migration_state: String = conn
+        .query_one("SELECT state FROM migration", [], |row| row.get(0))
+        .unwrap();
+    assert_eq!(migration_state, "done");
+}

glean-core/tests/sqlite_migration.rs

@@ -258,6 +258,83 @@ fn migration_checkpoints() {
     // Unvacuumed the database is _smaller_, because the migrated data is in the WAL file.
     // It's around 20k bytes.
     // Vacuumed & checkpointed the WAL transactions are merged into the database.
-    let vacuumed_database_size = 32768;
+    let vacuumed_database_size = 36864;
     assert_eq!(db_file_size, vacuumed_database_size);
 }
+
+#[test]
+fn migration_reapplied_after_not_finishing() {
+    let temp = {
+        let (glean, temp) = new_glean(None);
+        drop(glean);
+
+        let db_path = temp.path().join("db").join("glean.sqlite");
+        let conn = rusqlite::Connection::open(db_path).unwrap();
+
+        // Let's start with an empty database.
+        conn.execute("DELETE FROM telemetry", []).unwrap();
+
+        // Ensure migration isn't marked as done.
+        conn.execute("DELETE FROM migration", []).unwrap();
+        temp
+    };
+
+    let db_path = temp.path().join("db");
+
+    let safe_bin = db_path.join("data.safe.bin");
+    // Reusing the same database file from above.
+    fs::write(safe_bin, RKV_DATABASE).unwrap();
+    let exp_client_id = uuid!("77ca0472-5124-4f6b-971d-4a2a928fb158");
+
+    let (glean, _temp) = new_glean(Some(temp));
+
+    let client_id = clientid_metric().get_value(&glean, None).unwrap();
+    assert_eq!(exp_client_id, client_id);
+
+    let metrics = MigrationMetrics::new();
+    assert_eq!(Some(13), metrics.migrated_metrics.get_value(&glean, None));
+    assert_eq!(Some(13), metrics.metrics_in_sqlite.get_value(&glean, None));
+    assert_eq!(None, metrics.failed_metrics.get_value(&glean, None));
+
+    assert!(metrics.migration_duration.get_value(&glean, None).is_some());
+    assert_eq!(None, metrics.migration_error.get_value(&glean, None));
+}
+
+#[test]
+fn migration_not_reapplied_if_marked_as_finished() {
+    let (first_client_id, temp) = {
+        let (glean, temp) = new_glean(None);
+        let client_id = clientid_metric().get_value(&glean, None).unwrap();
+        drop(glean);
+
+        let db_path = temp.path().join("db").join("glean.sqlite");
+        let conn = rusqlite::Connection::open(db_path).unwrap();
+
+        // Ensure migration is marked as done.
+        conn.execute("DELETE FROM migration", []).unwrap();
+        conn.execute("INSERT INTO migration (id, state) VALUES (1, 'done')", [])
+            .unwrap();
+        (client_id, temp)
+    };
+
+    let db_path = temp.path().join("db");
+
+    let safe_bin = db_path.join("data.safe.bin");
+    // Reusing the same database file from above.
+    fs::write(safe_bin, RKV_DATABASE).unwrap();
+    let rkv_client_id = uuid!("77ca0472-5124-4f6b-971d-4a2a928fb158");
+
+    let (glean, _temp) = new_glean(Some(temp));
+
+    let client_id = clientid_metric().get_value(&glean, None).unwrap();
+    assert_ne!(rkv_client_id, client_id);
+    assert_eq!(first_client_id, client_id);
+
+    // No migration happened.
+    let metrics = MigrationMetrics::new();
+    assert_eq!(None, metrics.migrated_metrics.get_value(&glean, None));
+    assert_eq!(None, metrics.metrics_in_sqlite.get_value(&glean, None));
+    assert_eq!(None, metrics.failed_metrics.get_value(&glean, None));
+    assert_eq!(None, metrics.migration_duration.get_value(&glean, None));
+    assert_eq!(None, metrics.migration_error.get_value(&glean, None));
+}