Skip to main content

max / makenotwork

17.0 KB · 422 lines History Blame Raw
1 //! Per-test database isolation using PostgreSQL template databases.
2 //!
3 //! A shared template database is created once (with all migrations) and each
4 //! test gets a cheap `CREATE DATABASE ... TEMPLATE` clone. Dropped
5 //! automatically when `TestDb` goes out of scope.
6
7 use sqlx::postgres::PgPoolOptions;
8 use sqlx::{Connection, Executor, PgConnection, PgPool};
9 use std::sync::{Once, OnceLock};
10 use std::time::Duration;
11 use uuid::Uuid;
12
13 /// Base name of the shared template database. The live name is suffixed with
14 /// the connecting role (see `template_name`).
15 const TEMPLATE_DB_BASE: &str = "mnw_test_template";
16
17 /// The live template name for this run, namespaced by the connecting postgres
18 /// role. Two suites run by different roles on a shared cluster (e.g. a
19 /// developer's `max` and the Sando gate's `sando`) get distinct templates, so
20 /// they never collide AND each role owns, and can therefore drop, its own.
21 /// This is what prevents a stale, foreign-owned `mnw_test_template` from
22 /// wedging a deploy's `cargo_test` gate. Set once by `ensure_template`.
23 static TEMPLATE_NAME: OnceLock<String> = OnceLock::new();
24
25 fn template_name() -> &'static str {
26 TEMPLATE_NAME.get().map_or(TEMPLATE_DB_BASE, String::as_str)
27 }
28
29 /// Reduce a postgres role to a safe identifier suffix (`[a-z0-9_]`).
30 fn sanitize_role(role: &str) -> String {
31 role.chars()
32 .map(|c| {
33 if c.is_ascii_alphanumeric() || c == '_' {
34 c.to_ascii_lowercase()
35 } else {
36 '_'
37 }
38 })
39 .collect()
40 }
41
42 /// Shared, deploy-stable role that owns all test databases when present. Both
43 /// the interactive (`max`) and Sando-gate (`sando`) logins are members, so
44 /// whoever runs the suite, every `mnw_test_*` DB ends up owned by this role and
45 /// any member can drop it, no superuser, no cross-user ownership clashes.
46 const SHARED_ROLE: &str = "mnw_test";
47
48 /// Best-effort `SET ROLE mnw_test` so subsequent `CREATE DATABASE`s are owned
49 /// by the shared role. Silently no-ops when the role is absent or the login
50 /// isn't a member (e.g. a fresh dev box without bootstrap), the per-role
51 /// template namespacing is the fallback in that case.
52 async fn assume_shared_role(conn: &mut PgConnection) {
53 let _ = conn
54 .execute(format!("SET ROLE \"{SHARED_ROLE}\"").as_str())
55 .await;
56 }
57
58 /// Ensures template creation runs exactly once, across all threads and runtimes.
59 static TEMPLATE_INIT: Once = Once::new();
60
61 fn admin_url() -> String {
62 std::env::var("TEST_DATABASE_URL")
63 .unwrap_or_else(|_| "postgres://localhost/postgres".to_string())
64 }
65
66 /// Advisory-lock key serializing template setup across every process/connection
67 /// on the cluster (`std::sync::Once` only covers one process). Concurrent test
68 /// binaries must not both drop+recreate the shared template. Arbitrary, stable.
69 const TEMPLATE_LOCK_KEY: i64 = 0x6D6E_775F_7470_6C00; // "mnw_tpl\0"
70
71 /// Latest migration version embedded in this binary (head of `migrations/`).
72 fn latest_migration_version() -> i64 {
73 sqlx::migrate!("./migrations")
74 .iter()
75 .map(|m| m.version)
76 .max()
77 .unwrap_or(0)
78 }
79
80 /// True if `template` already exists and its applied-migration head matches the
81 /// code's latest migration, a clone of it is up to date, so we reuse it instead
82 /// of dropping + rebuilding. Any error / missing table => "not current" =>
83 /// caller rebuilds.
84 async fn template_is_current(admin_url: &str, template: &str) -> bool {
85 let Ok(mut conn) = PgConnection::connect(admin_url).await else {
86 return false;
87 };
88 let exists = sqlx::query_as::<_, (i32,)>("SELECT 1 FROM pg_database WHERE datname = $1")
89 .bind(template)
90 .fetch_optional(&mut conn)
91 .await
92 .ok()
93 .flatten()
94 .is_some();
95 if !exists {
96 return false;
97 }
98 let tpl_url = replace_db_name(admin_url, template);
99 let Ok(mut tconn) = PgConnection::connect(&tpl_url).await else {
100 return false;
101 };
102 matches!(
103 sqlx::query_as::<_, (Option<i64>,)>("SELECT MAX(version) FROM _sqlx_migrations")
104 .fetch_one(&mut tconn)
105 .await,
106 Ok((Some(v),)) if v == latest_migration_version()
107 )
108 }
109
110 /// Create the template database with all migrations. Runs in a dedicated
111 /// single-threaded tokio runtime so it works from any context (including
112 /// inside `#[tokio::test]` and plain `#[test]`).
113 fn ensure_template() {
114 TEMPLATE_INIT.call_once(|| {
115 let rt = tokio::runtime::Builder::new_current_thread()
116 .enable_all()
117 .build()
118 .expect("build template setup runtime");
119
120 rt.block_on(async {
121 let t0 = std::time::Instant::now();
122 let admin = admin_url();
123 let mut conn = PgConnection::connect(&admin)
124 .await
125 .expect("connect to admin DB for template setup");
126
127 // Adopt the shared role so the template (and every clone) is owned
128 // by `mnw_test`, droppable by any member. No-ops on hosts without
129 // the role, falling back to per-role namespacing below.
130 assume_shared_role(&mut conn).await;
131
132 // Namespace the template by the *effective* role: with the shared
133 // role assumed this collapses to one template for everyone; without
134 // it, each login role gets its own so a developer's `max` run and
135 // the Sando gate's `sando` run never collide on a shared cluster.
136 let (role,): (String,) = sqlx::query_as("SELECT current_user")
137 .fetch_one(&mut conn)
138 .await
139 .expect("query current_user for template namespacing");
140 let template = format!("{TEMPLATE_DB_BASE}_{}", sanitize_role(&role));
141 let _ = TEMPLATE_NAME.set(template.clone());
142
143 // Serialize template setup across ALL processes on this cluster, not
144 // just this process's `Once`. Without it, concurrent test binaries
145 // each unconditionally drop+recreate the SAME shared-role template,
146 // so one drops the template another is mid-clone from, the mass
147 // "template database ... does not exist" flake on the deploy gate.
148 // Held until `conn` drops at the end of this block.
149 sqlx::query("SELECT pg_advisory_lock($1)")
150 .bind(TEMPLATE_LOCK_KEY)
151 .execute(&mut conn)
152 .await
153 .expect("acquire template advisory lock");
154
155 // Reuse the template when it's already migration-current. Dropping a
156 // live template is what races concurrent clones, and a FORCE drop
157 // also fails when another role holds connections we can't terminate
158 // ("permission denied to terminate process"). Only rebuild when the
159 // template is missing or its migration head is stale.
160 if template_is_current(&admin, &template).await {
161 eprintln!("[test-harness] Reusing current template DB {template}");
162 } else {
163 // We hold the cross-process lock, so no clone can be reading it.
164 conn.execute(
165 format!("DROP DATABASE IF EXISTS \"{template}\" WITH (FORCE)").as_str(),
166 )
167 .await
168 .unwrap_or_else(|e| {
169 panic!(
170 "drop stale template {template}: {e} \
171 (if owned by a different role, drop it as the postgres superuser)"
172 )
173 });
174
175 conn.execute(format!("CREATE DATABASE \"{template}\"").as_str())
176 .await
177 .expect("create template database");
178
179 // Connect to the template and run all migrations
180 let tpl_url = replace_db_name(&admin, &template);
181 let tpl_pool = PgPoolOptions::new()
182 .max_connections(2)
183 .acquire_timeout(Duration::from_secs(10))
184 .connect(&tpl_url)
185 .await
186 .expect("connect to template database");
187
188 let t_migrate = std::time::Instant::now();
189 sqlx::migrate!("./migrations")
190 .run(&tpl_pool)
191 .await
192 .expect("run migrations on template");
193 let migrate_ms = t_migrate.elapsed().as_millis();
194
195 // Also create the session store table
196 let session_store = tower_sessions_sqlx_store::PostgresStore::new(tpl_pool.clone());
197 session_store
198 .migrate()
199 .await
200 .expect("session store migration on template");
201
202 tpl_pool.close().await;
203
204 let total_ms = t0.elapsed().as_millis();
205 eprintln!(
206 "[test-harness] Template DB created in {total_ms}ms (migrations: {migrate_ms}ms)"
207 );
208 }
209
210 // Release explicitly (also released when `conn` drops); the lock was
211 // held across the whole reuse-or-rebuild window above.
212 let _ = sqlx::query("SELECT pg_advisory_unlock($1)")
213 .bind(TEMPLATE_LOCK_KEY)
214 .execute(&mut conn)
215 .await;
216 });
217 });
218 }
219
220 /// An isolated test database that cleans up after itself.
221 pub(crate) struct TestDb {
222 pub pool: PgPool,
223 db_name: String,
224 admin_url: String,
225 #[allow(dead_code)]
226 test_url: String,
227 /// Whether the session store table already exists (from template).
228 pub session_migrated: bool,
229 }
230
231 impl TestDb {
232 /// Create a fresh database cloned from the shared template.
233 pub(crate) async fn new() -> Self {
234 // ensure_template uses std::sync::Once + its own runtime, safe from any context.
235 // When called from an async context, we run it on a blocking thread to avoid
236 // nesting runtimes.
237 tokio::task::spawn_blocking(ensure_template)
238 .await
239 .expect("template setup panicked");
240
241 let t0 = std::time::Instant::now();
242 let admin = admin_url();
243 let db_name = format!("mnw_test_{}", Uuid::new_v4().simple());
244
245 let mut admin_conn = PgConnection::connect(&admin)
246 .await
247 .expect("Failed to connect to admin database");
248
249 // Own the clone via the shared role too, so cleanup (and any member)
250 // can drop it. Matches the template's ownership.
251 assume_shared_role(&mut admin_conn).await;
252
253 // Clone from the template, retrying on transient contention. Postgres
254 // serializes `CREATE DATABASE ... TEMPLATE t`: a concurrent clone of the
255 // same template (under the full-suite parallel stampede) fails with
256 // "source database ... is being accessed by other users". Retry rather
257 // than fail the test on that transient.
258 let create_sql = format!(
259 "CREATE DATABASE \"{db_name}\" TEMPLATE \"{}\"",
260 template_name()
261 );
262 let mut attempt = 0u32;
263 loop {
264 match admin_conn.execute(create_sql.as_str()).await {
265 Ok(_) => break,
266 Err(_) if attempt < 8 => {
267 attempt += 1;
268 tokio::time::sleep(Duration::from_millis(100 * attempt as u64)).await;
269 }
270 Err(e) => panic!(
271 "Failed to create test database from template after {attempt} retries: {e}"
272 ),
273 }
274 }
275
276 let test_url = replace_db_name(&admin, &db_name);
277
278 let pool = PgPoolOptions::new()
279 .max_connections(5)
280 .acquire_timeout(Duration::from_secs(5))
281 .connect(&test_url)
282 .await
283 .expect("Failed to connect to test database");
284
285 let clone_ms = t0.elapsed().as_millis();
286 if clone_ms > 500 {
287 eprintln!("[test-harness] SLOW DB clone: {clone_ms}ms for {db_name}");
288 }
289
290 TestDb {
291 pool,
292 db_name,
293 admin_url: admin,
294 test_url,
295 session_migrated: true,
296 }
297 }
298
299 /// The connection URL for this test database.
300 #[allow(dead_code)]
301 pub(crate) fn url(&self) -> &str {
302 &self.test_url
303 }
304 }
305
306 impl Drop for TestDb {
307 fn drop(&mut self) {
308 let admin_url = self.admin_url.clone();
309 let db_name = self.db_name.clone();
310 let pool = self.pool.clone();
311
312 std::thread::spawn(move || {
313 let rt = tokio::runtime::Builder::new_current_thread()
314 .enable_all()
315 .build()
316 .expect("Failed to build cleanup runtime");
317
318 rt.block_on(async {
319 // Actually close the pool, and wait for it. Every one of its
320 // live connections is a client of the database we're about to
321 // drop, and `DROP DATABASE` refuses while any exist. Bounded:
322 // `close()` waits for checked-out connections to come back, and
323 // a test that leaked a guard would otherwise hang the suite
324 // here forever. On timeout we fall through, the FORCE drop
325 // below evicts whatever is left.
326 if tokio::time::timeout(Duration::from_secs(5), pool.close())
327 .await
328 .is_err()
329 {
330 eprintln!(
331 "[test-harness] pool for {db_name} did not close in 5s; forcing"
332 );
333 }
334
335 let Ok(mut conn) = PgConnection::connect(&admin_url).await else {
336 eprintln!("[test-harness] LEAKED {db_name}: admin connect failed");
337 return;
338 };
339
340 // Deliberately NO `SET ROLE mnw_test` here, unlike the create
341 // path. Terminating a backend requires membership in the role
342 // that *opened* it, and these were opened by our login role
343 // (`max`, `sando`), not by `mnw_test`, which is a member of
344 // neither. Assuming the shared role therefore turns both the
345 // sweep below and FORCE's own eviction into "permission denied
346 // to terminate process", leaking the clone. Dropping needs only
347 // membership in the owning role (pg_has_role USAGE), which we
348 // already have as the login role, so staying put satisfies both.
349 let _ = conn
350 .execute(
351 format!(
352 "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '{db_name}'"
353 )
354 .as_str(),
355 )
356 .await;
357
358 // FORCE, and retry: a connection can still be mid-handshake
359 // when we terminate, so it survives the sweep above and lands
360 // in the split second before the drop. Bare `DROP DATABASE`
361 // fails outright on that race; FORCE evicts it, and the retries
362 // cover the case where another arrives.
363 let drop_sql = format!("DROP DATABASE IF EXISTS \"{db_name}\" WITH (FORCE)");
364 for attempt in 0..4u32 {
365 match conn.execute(drop_sql.as_str()).await {
366 Ok(_) => return,
367 Err(e) if attempt == 3 => {
368 // Never silent: a swallowed failure here is exactly
369 // how the cluster accumulated ~1000 stale clones.
370 eprintln!("[test-harness] LEAKED {db_name}: drop failed: {e}");
371 }
372 Err(_) => {
373 tokio::time::sleep(Duration::from_millis(50 * (attempt + 1) as u64))
374 .await;
375 }
376 }
377 }
378 });
379 })
380 .join()
381 .ok();
382 }
383 }
384
385 /// Replace the database name in a PostgreSQL connection URL.
386 fn replace_db_name(url: &str, new_db: &str) -> String {
387 if let Some(pos) = url.rfind('/') {
388 let base = &url[..pos];
389 let query = url[pos + 1..].find('?').map_or("", |q| &url[pos + 1 + q..]);
390 if query.is_empty() {
391 format!("{base}/{new_db}")
392 } else {
393 format!("{base}/{new_db}{query}")
394 }
395 } else {
396 panic!("Invalid database URL: no '/' found");
397 }
398 }
399
400 #[cfg(test)]
401 pub(crate) mod tests {
402 use super::*;
403
404 #[test]
405 fn replace_db_name_simple() {
406 let result = replace_db_name("postgres://localhost/postgres", "test_db");
407 assert_eq!(result, "postgres://localhost/test_db");
408 }
409
410 #[test]
411 fn replace_db_name_with_auth() {
412 let result = replace_db_name("postgres://user:pass@localhost:5432/mydb", "test_db");
413 assert_eq!(result, "postgres://user:pass@localhost:5432/test_db");
414 }
415
416 #[test]
417 fn replace_db_name_with_query() {
418 let result = replace_db_name("postgres://localhost/postgres?sslmode=disable", "test_db");
419 assert_eq!(result, "postgres://localhost/test_db?sslmode=disable");
420 }
421 }
422