db.rs

  1pub mod kvp;
  2pub mod query;
  3
  4// Re-export
  5pub use anyhow;
  6use anyhow::Context as _;
  7pub use gpui;
  8use gpui::{App, AppContext, Global};
  9pub use indoc::indoc;
 10pub use inventory;
 11pub use paths::database_dir;
 12pub use smol;
 13pub use sqlez;
 14pub use sqlez_macros;
 15pub use uuid;
 16
 17pub use release_channel::RELEASE_CHANNEL;
 18use release_channel::ReleaseChannel;
 19use sqlez::domain::Migrator;
 20use sqlez::thread_safe_connection::ThreadSafeConnection;
 21use sqlez_macros::sql;
 22use std::future::Future;
 23use std::path::{Path, PathBuf};
 24use std::sync::atomic::AtomicBool;
 25use std::sync::{LazyLock, atomic::Ordering};
 26use util::{ResultExt, maybe};
 27use zed_env_vars::ZED_STATELESS;
 28
 29/// A migration registered via `static_connection!` and collected at link time.
 30pub struct DomainMigration {
 31    pub name: &'static str,
 32    pub migrations: &'static [&'static str],
 33    pub dependencies: &'static [&'static str],
 34    pub should_allow_migration_change: fn(usize, &str, &str) -> bool,
 35}
 36
 37inventory::collect!(DomainMigration);
 38
 39/// The shared database connection backing all domain-specific DB wrappers.
 40/// Set as a GPUI global per-App. Falls back to a shared LazyLock if not set.
 41pub struct AppDatabase(pub ThreadSafeConnection);
 42
 43impl Global for AppDatabase {}
 44
 45/// Migrator that runs all inventory-registered domain migrations.
 46pub struct AppMigrator;
 47
 48impl Migrator for AppMigrator {
 49    fn migrate(connection: &sqlez::connection::Connection) -> anyhow::Result<()> {
 50        let registrations: Vec<&DomainMigration> = inventory::iter::<DomainMigration>().collect();
 51        let sorted = topological_sort(&registrations);
 52        for reg in &sorted {
 53            let mut should_allow = reg.should_allow_migration_change;
 54            connection.migrate(reg.name, reg.migrations, &mut should_allow)?;
 55        }
 56        Ok(())
 57    }
 58}
 59
 60impl AppDatabase {
 61    /// Opens the production database and runs all inventory-registered
 62    /// migrations in dependency order.
 63    pub fn new() -> Self {
 64        let db_dir = database_dir();
 65        let connection = smol::block_on(open_db::<AppMigrator>(db_dir, *RELEASE_CHANNEL));
 66        Self(connection)
 67    }
 68
 69    /// Creates a new in-memory database with a unique name and runs all
 70    /// inventory-registered migrations in dependency order.
 71    #[cfg(any(test, feature = "test-support"))]
 72    pub fn test_new() -> Self {
 73        let name = format!("test-db-{}", uuid::Uuid::new_v4());
 74        let connection = smol::block_on(open_test_db::<AppMigrator>(&name));
 75        Self(connection)
 76    }
 77
 78    /// Returns the per-App connection if set, otherwise falls back to
 79    /// the shared LazyLock.
 80    pub fn global(cx: &App) -> &ThreadSafeConnection {
 81        #[allow(unreachable_code)]
 82        if let Some(db) = cx.try_global::<Self>() {
 83            return &db.0;
 84        } else {
 85            #[cfg(any(feature = "test-support", test))]
 86            return &TEST_APP_DATABASE.0;
 87
 88            panic!("database not initialized")
 89        }
 90    }
 91}
 92
 93fn topological_sort<'a>(registrations: &[&'a DomainMigration]) -> Vec<&'a DomainMigration> {
 94    let mut sorted: Vec<&DomainMigration> = Vec::new();
 95    let mut visited: std::collections::HashSet<&str> = std::collections::HashSet::new();
 96
 97    fn visit<'a>(
 98        name: &str,
 99        registrations: &[&'a DomainMigration],
100        sorted: &mut Vec<&'a DomainMigration>,
101        visited: &mut std::collections::HashSet<&'a str>,
102    ) {
103        if visited.contains(name) {
104            return;
105        }
106        if let Some(reg) = registrations.iter().find(|r| r.name == name) {
107            for dep in reg.dependencies {
108                visit(dep, registrations, sorted, visited);
109            }
110            visited.insert(reg.name);
111            sorted.push(reg);
112        }
113    }
114
115    for reg in registrations {
116        visit(reg.name, registrations, &mut sorted, &mut visited);
117    }
118    sorted
119}
120
121/// Shared fallback `AppDatabase` used when no per-App global is set.
122#[cfg(any(test, feature = "test-support"))]
123static TEST_APP_DATABASE: LazyLock<AppDatabase> = LazyLock::new(AppDatabase::test_new);
124
125const CONNECTION_INITIALIZE_QUERY: &str = sql!(
126    PRAGMA foreign_keys=TRUE;
127);
128
129const DB_INITIALIZE_QUERY: &str = sql!(
130    PRAGMA journal_mode=WAL;
131    PRAGMA busy_timeout=500;
132    PRAGMA case_sensitive_like=TRUE;
133    PRAGMA synchronous=NORMAL;
134);
135
136const FALLBACK_DB_NAME: &str = "FALLBACK_MEMORY_DB";
137
138const DB_FILE_NAME: &str = "db.sqlite";
139
140pub static ALL_FILE_DB_FAILED: LazyLock<AtomicBool> = LazyLock::new(|| AtomicBool::new(false));
141
142/// A type that can be used as a database scope for path construction.
143pub trait DbScope {
144    fn scope_name(&self) -> &str;
145}
146
147impl DbScope for ReleaseChannel {
148    fn scope_name(&self) -> &str {
149        self.dev_name()
150    }
151}
152
153/// A database scope shared across all release channels.
154pub struct GlobalDbScope;
155
156impl DbScope for GlobalDbScope {
157    fn scope_name(&self) -> &str {
158        "global"
159    }
160}
161
162/// Returns the path to the `AppDatabase` SQLite file for the given scope
163/// under `db_dir`.
164pub fn db_path(db_dir: &Path, scope: impl DbScope) -> PathBuf {
165    db_dir
166        .join(format!("0-{}", scope.scope_name()))
167        .join(DB_FILE_NAME)
168}
169
170/// Open or create a database at the given directory path.
171/// This will retry a couple times if there are failures. If opening fails once, the db directory
172/// is moved to a backup folder and a new one is created. If that fails, a shared in memory db is created.
173/// In either case, static variables are set so that the user can be notified.
174pub async fn open_db<M: Migrator + 'static>(
175    db_dir: &Path,
176    scope: impl DbScope,
177) -> ThreadSafeConnection {
178    if *ZED_STATELESS {
179        return open_fallback_db::<M>().await;
180    }
181
182    let db_path = db_path(db_dir, scope);
183
184    let connection = maybe!(async {
185        if let Some(parent) = db_path.parent() {
186            smol::fs::create_dir_all(parent)
187                .await
188                .context("Could not create db directory")
189                .log_err()?;
190        }
191        open_main_db::<M>(&db_path).await
192    })
193    .await;
194
195    if let Some(connection) = connection {
196        return connection;
197    }
198
199    // Set another static ref so that we can escalate the notification
200    ALL_FILE_DB_FAILED.store(true, Ordering::Release);
201
202    // If still failed, create an in memory db with a known name
203    open_fallback_db::<M>().await
204}
205
206async fn open_main_db<M: Migrator>(db_path: &Path) -> Option<ThreadSafeConnection> {
207    log::trace!("Opening database {}", db_path.display());
208    ThreadSafeConnection::builder::<M>(db_path.to_string_lossy().as_ref(), true)
209        .with_db_initialization_query(DB_INITIALIZE_QUERY)
210        .with_connection_initialize_query(CONNECTION_INITIALIZE_QUERY)
211        .build()
212        .await
213        .log_err()
214}
215
216async fn open_fallback_db<M: Migrator>() -> ThreadSafeConnection {
217    log::warn!("Opening fallback in-memory database");
218    ThreadSafeConnection::builder::<M>(FALLBACK_DB_NAME, false)
219        .with_db_initialization_query(DB_INITIALIZE_QUERY)
220        .with_connection_initialize_query(CONNECTION_INITIALIZE_QUERY)
221        .build()
222        .await
223        .expect(
224            "Fallback in memory database failed. Likely initialization queries or migrations have fundamental errors",
225        )
226}
227
228#[cfg(any(test, feature = "test-support"))]
229pub async fn open_test_db<M: Migrator>(db_name: &str) -> ThreadSafeConnection {
230    use sqlez::thread_safe_connection::locking_queue;
231
232    ThreadSafeConnection::builder::<M>(db_name, false)
233        .with_db_initialization_query(DB_INITIALIZE_QUERY)
234        .with_connection_initialize_query(CONNECTION_INITIALIZE_QUERY)
235        // Serialize queued writes via a mutex and run them synchronously
236        .with_write_queue_constructor(locking_queue())
237        .build()
238        .await
239        .unwrap()
240}
241
242/// Implements a basic DB wrapper for a given domain
243///
244/// Arguments:
245/// - type of connection wrapper
246/// - dependencies, whose migrations should be run prior to this domain's migrations
247#[macro_export]
248macro_rules! static_connection {
249    ($t:ident, [ $($d:ty),* ]) => {
250        impl ::std::ops::Deref for $t {
251            type Target = $crate::sqlez::thread_safe_connection::ThreadSafeConnection;
252
253            fn deref(&self) -> &Self::Target {
254                &self.0
255            }
256        }
257
258        impl ::std::clone::Clone for $t {
259            fn clone(&self) -> Self {
260                $t(self.0.clone())
261            }
262        }
263
264        impl $t {
265            /// Returns an instance backed by the per-App database if set,
266            /// or the shared fallback connection otherwise.
267            pub fn global(cx: &$crate::gpui::App) -> Self {
268                $t($crate::AppDatabase::global(cx).clone())
269            }
270
271            #[cfg(any(test, feature = "test-support"))]
272            pub async fn open_test_db(name: &'static str) -> Self {
273                $t($crate::open_test_db::<$t>(name).await)
274            }
275        }
276
277        $crate::inventory::submit! {
278            $crate::DomainMigration {
279                name: <$t as $crate::sqlez::domain::Domain>::NAME,
280                migrations: <$t as $crate::sqlez::domain::Domain>::MIGRATIONS,
281                dependencies: &[$(<$d as $crate::sqlez::domain::Domain>::NAME),*],
282                should_allow_migration_change: <$t as $crate::sqlez::domain::Domain>::should_allow_migration_change,
283            }
284        }
285    }
286}
287
288pub fn write_and_log<F>(cx: &App, db_write: impl FnOnce() -> F + Send + 'static)
289where
290    F: Future<Output = anyhow::Result<()>> + Send,
291{
292    cx.background_spawn(async move { db_write().await.log_err() })
293        .detach()
294}
295
296#[cfg(test)]
297mod tests {
298    use std::thread;
299
300    use sqlez::domain::Domain;
301    use sqlez_macros::sql;
302
303    use crate::open_db;
304
305    // Test bad migration panics
306    #[gpui::test]
307    #[should_panic]
308    async fn test_bad_migration_panics() {
309        enum BadDB {}
310
311        impl Domain for BadDB {
312            const NAME: &str = "db_tests";
313            const MIGRATIONS: &[&str] = &[
314                sql!(CREATE TABLE test(value);),
315                // failure because test already exists
316                sql!(CREATE TABLE test(value);),
317            ];
318        }
319
320        let tempdir = tempfile::Builder::new()
321            .prefix("DbTests")
322            .tempdir()
323            .unwrap();
324        let _bad_db = open_db::<BadDB>(tempdir.path(), release_channel::ReleaseChannel::Dev).await;
325    }
326
327    /// Test that DB exists but corrupted (causing recreate)
328    #[gpui::test]
329    async fn test_db_corruption(cx: &mut gpui::TestAppContext) {
330        cx.executor().allow_parking();
331
332        enum CorruptedDB {}
333
334        impl Domain for CorruptedDB {
335            const NAME: &str = "db_tests";
336            const MIGRATIONS: &[&str] = &[sql!(CREATE TABLE test(value);)];
337        }
338
339        enum GoodDB {}
340
341        impl Domain for GoodDB {
342            const NAME: &str = "db_tests"; //Notice same name
343            const MIGRATIONS: &[&str] = &[sql!(CREATE TABLE test2(value);)];
344        }
345
346        let tempdir = tempfile::Builder::new()
347            .prefix("DbTests")
348            .tempdir()
349            .unwrap();
350        {
351            let corrupt_db =
352                open_db::<CorruptedDB>(tempdir.path(), release_channel::ReleaseChannel::Dev).await;
353            assert!(corrupt_db.persistent());
354        }
355
356        let good_db = open_db::<GoodDB>(tempdir.path(), release_channel::ReleaseChannel::Dev).await;
357        assert!(
358            good_db.select_row::<usize>("SELECT * FROM test2").unwrap()()
359                .unwrap()
360                .is_none()
361        );
362    }
363
364    /// Test that DB exists but corrupted (causing recreate)
365    #[gpui::test(iterations = 30)]
366    async fn test_simultaneous_db_corruption(cx: &mut gpui::TestAppContext) {
367        cx.executor().allow_parking();
368
369        enum CorruptedDB {}
370
371        impl Domain for CorruptedDB {
372            const NAME: &str = "db_tests";
373
374            const MIGRATIONS: &[&str] = &[sql!(CREATE TABLE test(value);)];
375        }
376
377        enum GoodDB {}
378
379        impl Domain for GoodDB {
380            const NAME: &str = "db_tests"; //Notice same name
381            const MIGRATIONS: &[&str] = &[sql!(CREATE TABLE test2(value);)]; // But different migration
382        }
383
384        let tempdir = tempfile::Builder::new()
385            .prefix("DbTests")
386            .tempdir()
387            .unwrap();
388        {
389            // Setup the bad database
390            let corrupt_db =
391                open_db::<CorruptedDB>(tempdir.path(), release_channel::ReleaseChannel::Dev).await;
392            assert!(corrupt_db.persistent());
393        }
394
395        // Try to connect to it a bunch of times at once
396        let mut guards = vec![];
397        for _ in 0..10 {
398            let tmp_path = tempdir.path().to_path_buf();
399            let guard = thread::spawn(move || {
400                let good_db = smol::block_on(open_db::<GoodDB>(
401                    tmp_path.as_path(),
402                    release_channel::ReleaseChannel::Dev,
403                ));
404                assert!(
405                    good_db.select_row::<usize>("SELECT * FROM test2").unwrap()()
406                        .unwrap()
407                        .is_none()
408                );
409            });
410
411            guards.push(guard);
412        }
413
414        for guard in guards.into_iter() {
415            assert!(guard.join().is_ok());
416        }
417    }
418}