max / makenotwork

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