fix(sqlite): Ensure that database migration runs sequentially

Otherwise, we might run migration from multiple db-connections.

3 files changed, 106 insertions, 52 deletions

diff --git a/crates/turtle/src/atuin_client/database.rs b/crates/turtle/src/atuin_client/database.rs
index c730b1d4..c3130c4c 100644
--- a/crates/turtle/src/atuin_client/database.rs
+++ b/crates/turtle/src/atuin_client/database.rs
@@ -2,19 +2,15 @@ use std::{
     env,
     path::{Path, PathBuf},
     str::FromStr,
-    time::Duration,
 };
 
-use crate::atuin_common::utils;
-use fs_err as fs;
+use crate::{atuin_client::utils::setup_db, atuin_common::utils};
+use fs_err::{self as fs};
 use itertools::Itertools;
 use sql_builder::{SqlBuilder, SqlName, bind::Bind, esc, quote};
 use sqlx::{
     Result, Row,
-    sqlite::{
-        SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqlitePoolOptions, SqliteRow,
-        SqliteSynchronous,
-    },
+    sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqliteRow, SqliteSynchronous},
 };
 use time::OffsetDateTime;
 use tracing::debug;
@@ -103,6 +99,17 @@ pub(crate) struct ClientSqlite {
 
 impl ClientSqlite {
     pub(crate) async fn new(path: impl AsRef<Path>, timeout: f64) -> Result<Self> {
+        fn mk_opts(path: &str) -> Result<SqliteConnectOptions> {
+            let opts = SqliteConnectOptions::from_str(path)?
+                .journal_mode(SqliteJournalMode::Wal)
+                .optimize_on_close(true, None)
+                .synchronous(SqliteSynchronous::Normal)
+                .with_regexp()
+                .create_if_missing(true);
+
+            Ok(opts)
+        }
+
         let path = path.as_ref();
         debug!("opening sqlite database at {path:?}");
 
@@ -120,30 +127,10 @@ impl ClientSqlite {
             fs::create_dir_all(dir)?;
         }
 
-        let opts = SqliteConnectOptions::from_str(path.as_os_str().to_str().unwrap())?
-            .journal_mode(SqliteJournalMode::Wal)
-            .optimize_on_close(true, None)
-            .synchronous(SqliteSynchronous::Normal)
-            .with_regexp()
-            .create_if_missing(true);
-
-        let pool = SqlitePoolOptions::new()
-            .acquire_timeout(Duration::from_secs_f64(timeout))
-            .connect_with(opts)
-            .await?;
-
-        Self::setup_db(&pool).await?;
+        let pool = setup_db!(path, timeout, mk_opts, "./db/client-migrations").await?;
         Ok(Self { pool })
     }
 
-    async fn setup_db(pool: &SqlitePool) -> Result<()> {
-        debug!("running sqlite database setup");
-
-        sqlx::migrate!("./db/client-migrations").run(pool).await?;
-
-        Ok(())
-    }
-
     async fn save_raw(tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, h: &History) -> Result<()> {
         sqlx::query(
             "insert or ignore into history(id, timestamp, duration, exit, command, cwd, session, hostname, author, intent, deleted_at)
diff --git a/crates/turtle/src/atuin_client/record/sqlite_store.rs b/crates/turtle/src/atuin_client/record/sqlite_store.rs
index 8a9afe45..18f5c869 100644
--- a/crates/turtle/src/atuin_client/record/sqlite_store.rs
+++ b/crates/turtle/src/atuin_client/record/sqlite_store.rs
@@ -2,18 +2,19 @@
 // Multiple stores of multiple types are all stored in one chonky table (for now), and we just index
 // by tag/host
 
+use std::path::Path;
 use std::str::FromStr;
-use std::{path::Path, time::Duration};
 
 use eyre::{Result, eyre};
 use fs_err as fs;
 
 use sqlx::{
     Row,
-    sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqlitePoolOptions, SqliteRow},
+    sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqliteRow},
 };
 use tracing::debug;
 
+use crate::atuin_client::utils::setup_db;
 use crate::atuin_common::record::{
     EncryptedData, Host, HostId, Record, RecordId, RecordIdx, RecordStatus,
 };
@@ -29,6 +30,15 @@ pub(crate) struct SqliteStore {
 
 impl SqliteStore {
     pub(crate) async fn new(path: impl AsRef<Path>, timeout: f64) -> Result<Self> {
+        fn mk_opts(path: &str) -> sqlx::Result<SqliteConnectOptions> {
+            let opts = SqliteConnectOptions::from_str(path)?
+                .journal_mode(SqliteJournalMode::Wal)
+                .foreign_keys(true)
+                .create_if_missing(true);
+
+            Ok(opts)
+        }
+
         let path = path.as_ref();
 
         debug!("opening sqlite database at {path:?}");
@@ -47,31 +57,11 @@ impl SqliteStore {
             fs::create_dir_all(dir)?;
         }
 
-        let opts = SqliteConnectOptions::from_str(path.as_os_str().to_str().unwrap())?
-            .journal_mode(SqliteJournalMode::Wal)
-            .foreign_keys(true)
-            .create_if_missing(true);
-
-        let pool = SqlitePoolOptions::new()
-            .acquire_timeout(Duration::from_secs_f64(timeout))
-            .connect_with(opts)
-            .await?;
-
-        Self::setup_db(&pool).await?;
+        let pool = setup_db!(path, timeout, mk_opts, "./db/client-record-migrations").await?;
 
         Ok(Self { pool })
     }
 
-    async fn setup_db(pool: &SqlitePool) -> Result<()> {
-        debug!("running sqlite database setup");
-
-        sqlx::migrate!("./db/client-record-migrations")
-            .run(pool)
-            .await?;
-
-        Ok(())
-    }
-
     async fn save_raw(
         tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
         r: &Record<EncryptedData>,
diff --git a/crates/turtle/src/atuin_client/utils.rs b/crates/turtle/src/atuin_client/utils.rs
index 6d178b77..989f9fc1 100644
--- a/crates/turtle/src/atuin_client/utils.rs
+++ b/crates/turtle/src/atuin_client/utils.rs
@@ -1,4 +1,3 @@
-
 pub(crate) fn get_hostname() -> String {
     std::env::var("ATUIN_HOST_NAME")
         .unwrap_or_else(|_| whoami::hostname().unwrap_or_else(|_| "unknown-host".to_string()))
@@ -13,3 +12,81 @@ pub(crate) fn get_username() -> String {
 pub(crate) fn get_host_user() -> String {
     format!("{}:{}", get_hostname(), get_username())
 }
+
+/// Setup a [`SQLite`] database.
+///
+/// This takes care of correct locking, so that we avoid a race when setting up the database.
+macro_rules! setup_db {
+    (
+        $db_path:expr,
+        $a_timeout:expr,
+        $opts:expr,
+        $m_name:literal $(,)?
+    ) => {{
+        async fn migrate(pool: &SqlitePool) -> sqlx::Result<()> {
+            { sqlx::sqlx_macros::migrate!($m_name) }.run(pool).await?;
+            Ok(())
+        }
+
+        crate::atuin_client::utils::setup_db_inner($db_path, $a_timeout, $opts, migrate)
+    }};
+}
+pub(crate) use setup_db;
+
+use std::{os::fd::AsRawFd, path::Path, time::Duration};
+
+use fs_err::OpenOptions;
+use sqlx::{
+    SqlitePool,
+    sqlite::{SqliteConnectOptions, SqlitePoolOptions},
+};
+use tracing::debug;
+
+/// Helper for `setup_db!`
+pub(crate) async fn setup_db_inner(
+    db_path: &Path,
+    acquire_timeout: f64,
+    mk_opts: fn(&str) -> sqlx::Result<SqliteConnectOptions>,
+    migrate: impl AsyncFn(&SqlitePool) -> sqlx::Result<()>,
+) -> sqlx::Result<SqlitePool> {
+    async fn open_db(timeout: f64, opts: SqliteConnectOptions) -> sqlx::Result<SqlitePool> {
+        let pool = SqlitePoolOptions::new()
+            .acquire_timeout(Duration::from_secs_f64(timeout))
+            .connect_with(opts)
+            .await?;
+
+        Ok(pool)
+    }
+
+    {
+        let file = OpenOptions::new()
+            .read(true)
+            .write(true)
+            .create(true)
+            .open(db_path)?;
+
+        // Lock the db file while we are running the migrations.
+        // Why? Because there is a small chance that we start running migrations (e.g. as the daemon)
+        // and then another process is started, which will also try to run migrations.
+        // Essentially, one of the processes will receive with a SQLite UNIQUE constraint failure.
+        // So let's avoid that possibility from the start.
+        file.lock()?;
+
+        let pool = open_db(
+            acquire_timeout,
+            mk_opts(format!("/proc/self/fd/{}", file.as_raw_fd()).as_str())?,
+        )
+        .await?;
+
+        debug!("running sqlite database setup");
+
+        migrate(&pool).await?;
+
+        file.unlock()?;
+    }
+
+    let real_opts = mk_opts(db_path.to_str().expect("Should be utf-8"))?;
+    let pool = open_db(acquire_timeout, real_opts).await?;
+
+    Ok(pool)
+}