Skip to main content

max / makenotwork

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