1use super::*;
2use gpui::executor::Background;
3use parking_lot::Mutex;
4use sea_orm::ConnectionTrait;
5use sqlx::migrate::MigrateDatabase;
6use std::sync::Arc;
7
8pub struct TestDb {
9 pub db: Option<Arc<Database>>,
10 pub connection: Option<sqlx::AnyConnection>,
11}
12
13impl TestDb {
14 pub fn sqlite(background: Arc<Background>) -> Self {
15 let url = format!("sqlite::memory:");
16 let runtime = tokio::runtime::Builder::new_current_thread()
17 .enable_io()
18 .enable_time()
19 .build()
20 .unwrap();
21
22 let mut db = runtime.block_on(async {
23 let mut options = ConnectOptions::new(url);
24 options.max_connections(5);
25 let db = Database::new(options, Executor::Deterministic(background))
26 .await
27 .unwrap();
28 let sql = include_str!(concat!(
29 env!("CARGO_MANIFEST_DIR"),
30 "/migrations.sqlite/20221109000000_test_schema.sql"
31 ));
32 db.pool
33 .execute(sea_orm::Statement::from_string(
34 db.pool.get_database_backend(),
35 sql.into(),
36 ))
37 .await
38 .unwrap();
39 db
40 });
41
42 db.runtime = Some(runtime);
43
44 Self {
45 db: Some(Arc::new(db)),
46 connection: None,
47 }
48 }
49
50 pub fn postgres(background: Arc<Background>) -> Self {
51 static LOCK: Mutex<()> = Mutex::new(());
52
53 let _guard = LOCK.lock();
54 let mut rng = StdRng::from_entropy();
55 let url = format!(
56 "postgres://postgres@localhost/zed-test-{}",
57 rng.gen::<u128>()
58 );
59 let runtime = tokio::runtime::Builder::new_current_thread()
60 .enable_io()
61 .enable_time()
62 .build()
63 .unwrap();
64
65 let mut db = runtime.block_on(async {
66 sqlx::Postgres::create_database(&url)
67 .await
68 .expect("failed to create test db");
69 let mut options = ConnectOptions::new(url);
70 options
71 .max_connections(5)
72 .idle_timeout(Duration::from_secs(0));
73 let db = Database::new(options, Executor::Deterministic(background))
74 .await
75 .unwrap();
76 let migrations_path = concat!(env!("CARGO_MANIFEST_DIR"), "/migrations");
77 db.migrate(Path::new(migrations_path), false).await.unwrap();
78 db
79 });
80
81 db.runtime = Some(runtime);
82
83 Self {
84 db: Some(Arc::new(db)),
85 connection: None,
86 }
87 }
88
89 pub fn db(&self) -> &Arc<Database> {
90 self.db.as_ref().unwrap()
91 }
92}
93
94impl Drop for TestDb {
95 fn drop(&mut self) {
96 let db = self.db.take().unwrap();
97 if let sea_orm::DatabaseBackend::Postgres = db.pool.get_database_backend() {
98 db.runtime.as_ref().unwrap().block_on(async {
99 use util::ResultExt;
100 let query = "
101 SELECT pg_terminate_backend(pg_stat_activity.pid)
102 FROM pg_stat_activity
103 WHERE
104 pg_stat_activity.datname = current_database() AND
105 pid <> pg_backend_pid();
106 ";
107 db.pool
108 .execute(sea_orm::Statement::from_string(
109 db.pool.get_database_backend(),
110 query.into(),
111 ))
112 .await
113 .log_err();
114 sqlx::Postgres::drop_database(db.options.get_url())
115 .await
116 .log_err();
117 })
118 }
119 }
120}