headgate_testkit/
database.rs

1//! Per-test live-store isolation. SQL namespaces are installed through
2//! `headgate-migrate`; Redis uses a generated key prefix because its schema is the
3//! keyspace, not DDL.
4
5use std::fmt;
6use std::sync::atomic::{AtomicU64, Ordering};
7
8use headgate_migrate::{Direction, MigrateOptions, migrate_mysql, migrate_postgres_in_schema};
9use mysql_async::prelude::*;
10
11static NEXT_NAMESPACE: AtomicU64 = AtomicU64::new(1);
12
13fn unique_name(backend: &str) -> String {
14    // Generated from trusted components only, so it is safe to interpolate as a SQL
15    // identifier. Keeping it under 63 bytes also satisfies Postgres's identifier limit.
16    format!(
17        "hg_test_{backend}_{}_{}",
18        std::process::id(),
19        NEXT_NAMESPACE.fetch_add(1, Ordering::Relaxed)
20    )
21}
22
23#[derive(Debug)]
24pub struct TestDatabaseError(String);
25
26impl TestDatabaseError {
27    fn new(message: impl Into<String>) -> Self {
28        Self(message.into())
29    }
30}
31
32impl fmt::Display for TestDatabaseError {
33    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
34        f.write_str(&self.0)
35    }
36}
37
38impl std::error::Error for TestDatabaseError {}
39
40async fn pg_connect(
41    config: &tokio_postgres::Config,
42) -> Result<
43    (
44        tokio_postgres::Client,
45        tokio::task::JoinHandle<Result<(), tokio_postgres::Error>>,
46    ),
47    TestDatabaseError,
48> {
49    let (client, connection) = config
50        .connect(tokio_postgres::NoTls)
51        .await
52        .map_err(|error| TestDatabaseError::new(format!("postgres connect: {error}")))?;
53    Ok((client, tokio::spawn(connection)))
54}
55
56/// One migrated Postgres schema. Use [`config`](Self::config) to construct a pool/store;
57/// it carries the schema in the startup `search_path`, so every pooled connection sees
58/// the same isolated namespace.
59pub struct PostgresTestDatabase {
60    schema: String,
61    admin_config: tokio_postgres::Config,
62    test_config: tokio_postgres::Config,
63}
64
65impl PostgresTestDatabase {
66    pub async fn create(conninfo: &str) -> Result<Self, TestDatabaseError> {
67        let admin_config: tokio_postgres::Config = conninfo
68            .parse()
69            .map_err(|error| TestDatabaseError::new(format!("bad Postgres conninfo: {error}")))?;
70        let schema = unique_name("pg");
71        let (mut admin, admin_task) = pg_connect(&admin_config).await?;
72        admin
73            .batch_execute(&format!("CREATE SCHEMA {schema}"))
74            .await
75            .map_err(|error| TestDatabaseError::new(format!("create schema {schema}: {error}")))?;
76
77        let mut test_config = admin_config.clone();
78        test_config.options(&format!("-c search_path={schema}"));
79        let migrated = migrate_postgres_in_schema(
80            &mut admin,
81            &schema,
82            Direction::Up,
83            MigrateOptions::default(),
84        )
85        .await
86        .map_err(|error| TestDatabaseError::new(format!("{error}: {error:?}")));
87        if let Err(error) = migrated {
88            let _ = admin
89                .batch_execute(&format!("DROP SCHEMA {schema} CASCADE"))
90                .await;
91            drop(admin);
92            let _ = admin_task.await;
93            return Err(error);
94        }
95        drop(admin);
96        let _ = admin_task.await;
97        Ok(Self {
98            schema,
99            admin_config,
100            test_config,
101        })
102    }
103
104    pub fn schema(&self) -> &str {
105        &self.schema
106    }
107
108    pub fn config(&self) -> tokio_postgres::Config {
109        self.test_config.clone()
110    }
111
112    /// Drop only this helper's generated schema and every object the migrator installed
113    /// in it. The helper consumes itself so cleanup cannot accidentally run twice.
114    pub async fn cleanup(self) -> Result<(), TestDatabaseError> {
115        let (admin, task) = pg_connect(&self.admin_config).await?;
116        let result = admin
117            .batch_execute(&format!("DROP SCHEMA {} CASCADE", self.schema))
118            .await
119            .map_err(|error| TestDatabaseError::new(format!("drop schema: {error}")));
120        drop(admin);
121        let _ = task.await;
122        result
123    }
124}
125
126/// One migrated MySQL database. [`opts`](Self::opts) can be cloned into a
127/// `mysql_async::Pool` without re-parsing or exposing credentials.
128pub struct MysqlTestDatabase {
129    database: String,
130    admin_opts: mysql_async::Opts,
131    test_opts: mysql_async::Opts,
132}
133
134impl MysqlTestDatabase {
135    pub async fn create(url: &str) -> Result<Self, TestDatabaseError> {
136        let admin_opts = mysql_async::Opts::from_url(url)
137            .map_err(|error| TestDatabaseError::new(format!("bad MySQL URL: {error}")))?;
138        let database = unique_name("mysql");
139        let admin_pool = mysql_async::Pool::new(admin_opts.clone());
140        let mut admin = admin_pool
141            .get_conn()
142            .await
143            .map_err(|error| TestDatabaseError::new(format!("mysql connect: {error}")))?;
144        admin
145            .query_drop(format!("CREATE DATABASE {database}"))
146            .await
147            .map_err(|error| {
148                TestDatabaseError::new(format!("create database {database}: {error}"))
149            })?;
150
151        let test_opts: mysql_async::Opts = mysql_async::OptsBuilder::from_opts(admin_opts.clone())
152            .db_name(Some(database.clone()))
153            .into();
154        let pool = mysql_async::Pool::new(test_opts.clone());
155        let migrated = async {
156            let mut conn = pool
157                .get_conn()
158                .await
159                .map_err(|error| TestDatabaseError::new(format!("mysql connect: {error}")))?;
160            migrate_mysql(&mut conn, Direction::Up, MigrateOptions::default())
161                .await
162                .map_err(|error| TestDatabaseError::new(error.to_string()))
163        }
164        .await;
165        let _ = pool.disconnect().await;
166        if let Err(error) = migrated {
167            let _ = admin.query_drop(format!("DROP DATABASE {database}")).await;
168            drop(admin);
169            let _ = admin_pool.disconnect().await;
170            return Err(error);
171        }
172        drop(admin);
173        let _ = admin_pool.disconnect().await;
174        Ok(Self {
175            database,
176            admin_opts,
177            test_opts,
178        })
179    }
180
181    pub fn database(&self) -> &str {
182        &self.database
183    }
184
185    pub fn opts(&self) -> mysql_async::Opts {
186        self.test_opts.clone()
187    }
188
189    pub async fn cleanup(self) -> Result<(), TestDatabaseError> {
190        let pool = mysql_async::Pool::new(self.admin_opts);
191        let mut conn = pool
192            .get_conn()
193            .await
194            .map_err(|error| TestDatabaseError::new(format!("mysql connect: {error}")))?;
195        let result = conn
196            .query_drop(format!("DROP DATABASE {}", self.database))
197            .await
198            .map_err(|error| TestDatabaseError::new(format!("drop database: {error}")));
199        drop(conn);
200        let _ = pool.disconnect().await;
201        result
202    }
203}
204
205/// A process-unique Redis key prefix. Redis's schema boundary is already an explicit
206/// prefix in both drivers, so tests do not need a scarce numbered logical database and
207/// never call `FLUSHALL`/`FLUSHDB`.
208pub struct RedisTestNamespace {
209    prefix: String,
210    client: redis::Client,
211}
212
213impl RedisTestNamespace {
214    pub async fn create(url: &str) -> Result<Self, TestDatabaseError> {
215        let client = redis::Client::open(url)
216            .map_err(|error| TestDatabaseError::new(format!("bad Redis URL: {error}")))?;
217        let namespace = Self {
218            prefix: unique_name("redis"),
219            client,
220        };
221        // A generated name should be empty. Checking through SCAN catches the astronomic
222        // collision case without introducing KEYS (which is exactly the production trap
223        // the queue avoids).
224        if namespace.scan_keys().await?.is_empty() {
225            Ok(namespace)
226        } else {
227            Err(TestDatabaseError::new(format!(
228                "generated Redis prefix {} already exists",
229                namespace.prefix
230            )))
231        }
232    }
233
234    pub fn prefix(&self) -> &str {
235        &self.prefix
236    }
237
238    pub fn client(&self) -> redis::Client {
239        self.client.clone()
240    }
241
242    pub async fn connection_manager(
243        &self,
244    ) -> Result<redis::aio::ConnectionManager, TestDatabaseError> {
245        self.client
246            .get_connection_manager()
247            .await
248            .map_err(|error| TestDatabaseError::new(format!("redis connect: {error}")))
249    }
250
251    async fn scan_keys(&self) -> Result<Vec<String>, TestDatabaseError> {
252        let mut conn = self
253            .client
254            .get_multiplexed_async_connection()
255            .await
256            .map_err(|error| TestDatabaseError::new(format!("redis connect: {error}")))?;
257        let mut cursor = 0_u64;
258        let mut keys = Vec::new();
259        loop {
260            let (next, mut page): (u64, Vec<String>) = redis::cmd("SCAN")
261                .arg(cursor)
262                .arg("MATCH")
263                .arg(format!("{}:*", self.prefix))
264                .arg("COUNT")
265                .arg(100)
266                .query_async(&mut conn)
267                .await
268                .map_err(|error| TestDatabaseError::new(format!("redis scan: {error}")))?;
269            keys.append(&mut page);
270            cursor = next;
271            if cursor == 0 {
272                break;
273            }
274        }
275        Ok(keys)
276    }
277
278    pub async fn cleanup(self) -> Result<(), TestDatabaseError> {
279        let keys = self.scan_keys().await?;
280        if keys.is_empty() {
281            return Ok(());
282        }
283        let mut conn = self
284            .client
285            .get_multiplexed_async_connection()
286            .await
287            .map_err(|error| TestDatabaseError::new(format!("redis connect: {error}")))?;
288        for page in keys.chunks(100) {
289            redis::cmd("DEL")
290                .arg(page)
291                .query_async::<()>(&mut conn)
292                .await
293                .map_err(|error| TestDatabaseError::new(format!("redis cleanup: {error}")))?;
294        }
295        Ok(())
296    }
297}
headgate_testkit/database.rs

headgate_testkit/
database.rs
