aster_forge_test/
postgres.rs

1//! Shared reusable `PostgreSQL` container for integration tests.
2//!
3//! The container provides isolated databases with automatic stale-resource cleanup. Products own
4//! their migrations and seed data; this module owns database creation, connection retry, and
5//! teardown mechanics.
6
7use crate::database::connect_with_retry;
8use crate::state::{ContainerLease, ContainerStateLock};
9use crate::suite::TestContainerSuite;
10use sea_orm::{ConnectionTrait, DatabaseConnection};
11use testcontainers::core::{ContainerAsync, ContainerRequest, IntoContainerPort};
12use testcontainers::{GenericImage, ImageExt, ReuseDirective, runners::AsyncRunner};
13
14const POSTGRES_TEST_SHM_SIZE_BYTES: u64 = 1024 * 1024 * 1024;
15const POSTGRES_CONTAINER_SERVICE: &str = "postgres-shm-1g";
16
17/// Handle to the suite's shared `PostgreSQL` container.
18pub struct PostgresTestContainer {
19    admin_url: String,
20    suite: TestContainerSuite,
21    _container: ContainerAsync<GenericImage>,
22    _lease: ContainerLease,
23}
24
25/// Isolated `PostgreSQL` database owned by one test process.
26pub struct PostgresTestDatabase {
27    name: String,
28    url: String,
29    admin_url: String,
30    suite: TestContainerSuite,
31    ownership: DatabaseOwnership,
32}
33
34#[derive(Clone, Copy)]
35enum DatabaseOwnership {
36    Process,
37    Shared,
38}
39
40impl PostgresTestContainer {
41    /// Starts (or reuses) the shared `PostgreSQL` container with `postgres`/`postgres` credentials.
42    ///
43    /// # Panics
44    ///
45    /// Panics when shared state, container startup, port discovery, readiness, stale-database
46    /// cleanup, or connection shutdown fails.
47    pub async fn start(suite: &TestContainerSuite) -> Self {
48        let lock = ContainerStateLock::acquire(suite, "postgres");
49        let mut state = lock.load();
50        let stale_resources = state.prune_stale_during_current_execution();
51        state.register_current_process();
52        for resource in &stale_resources {
53            state.remember_current_process_resource(resource);
54        }
55        lock.save(&state);
56
57        drop(lock);
58
59        // Several nextest processes can enter this path at once. Docker's reusable-container
60        // lookup and create operation is not atomic, so one process may briefly observe a name
61        // conflict while another is creating the shared container. Retry those startup errors
62        // until the first process has published a reusable container that this process can attach
63        // to.
64        let container = start_postgres_container_with_retry(suite).await;
65        let port = container
66            .get_host_port_ipv4(IntoContainerPort::tcp(5432))
67            .await
68            .expect("PostgreSQL test port should be exposed");
69        let admin_url = format!("postgres://postgres:postgres@127.0.0.1:{port}/postgres");
70        connect_with_retry(&admin_url, "PostgreSQL")
71            .await
72            .close()
73            .await
74            .expect("failed to close PostgreSQL readiness probe connection");
75
76        let fixture = Self {
77            admin_url,
78            suite: suite.clone(),
79            _container: container,
80            _lease: ContainerLease::new(suite.clone(), "postgres"),
81        };
82        fixture.cleanup_databases(&stale_resources).await;
83        fixture
84    }
85
86    /// Returns the admin URL pointing at the default `postgres` database.
87    #[must_use]
88    pub fn admin_url(&self) -> &str {
89        &self.admin_url
90    }
91
92    /// Returns a stable identity for the running suite container.
93    ///
94    /// The resolved admin URL includes the checkout-scoped host port and therefore changes when
95    /// the reusable container belongs to a different suite instance.
96    #[must_use]
97    pub fn container_identity(&self) -> &str {
98        &self.admin_url
99    }
100
101    /// Creates and registers an isolated database for a product test.
102    ///
103    /// # Panics
104    ///
105    /// Panics when the database name is invalid, shared state fails, or the admin connection,
106    /// `CREATE DATABASE`, or connection shutdown fails.
107    pub async fn create_database(&self, name: &str) -> PostgresTestDatabase {
108        self.create_database_inner(name, None, DatabaseOwnership::Process)
109            .await
110    }
111
112    /// Creates and registers an isolated database cloned from `template`.
113    ///
114    /// Products still own the template contents, migrations, and seed data. This helper only
115    /// provides the product-neutral `PostgreSQL` database lifecycle and safe identifier handling.
116    ///
117    /// # Panics
118    ///
119    /// Panics when either database name is invalid, shared state fails, or the admin connection,
120    /// `CREATE DATABASE ... TEMPLATE ...`, or connection shutdown fails.
121    pub async fn create_database_from_template(
122        &self,
123        name: &str,
124        template: &str,
125    ) -> PostgresTestDatabase {
126        assert_valid_database_name(template);
127        self.create_database_inner(name, Some(template), DatabaseOwnership::Process)
128            .await
129    }
130
131    /// Creates a suite-scoped database for a product-owned reusable fixture.
132    ///
133    /// Unlike [`Self::create_database`], this resource survives the producer process. Products
134    /// must use a separate locked fixture-state protocol to validate or invalidate its contents.
135    pub async fn create_shared_database(&self, name: &str) -> PostgresTestDatabase {
136        self.create_database_inner(name, None, DatabaseOwnership::Shared)
137            .await
138    }
139
140    /// Drops a suite-scoped fixture database and unregisters it.
141    ///
142    /// # Panics
143    ///
144    /// Panics when the name is invalid or database cleanup, connection, or shutdown fails.
145    pub async fn drop_shared_database(&self, name: &str) {
146        assert_valid_database_name(name);
147        let admin = connect_with_retry(&self.admin_url, "PostgreSQL").await;
148        admin
149            .execute_unprepared(&format!(
150                "DROP DATABASE IF EXISTS {} WITH (FORCE)",
151                quote_identifier(name)
152            ))
153            .await
154            .unwrap_or_else(|error| {
155                panic!("failed to drop shared PostgreSQL test database {name}: {error}")
156            });
157        admin.close().await.unwrap_or_else(|error| {
158            panic!("failed to close PostgreSQL shared database admin connection: {error}")
159        });
160
161        self.forget_shared_resource(name);
162    }
163
164    /// Registers a suite-scoped product fixture that must outlive the producer process.
165    ///
166    /// This is for product-owned migrated template databases. Products must pair it with their
167    /// own fingerprint-based invalidation and call [`Self::forget_shared_resource`] after
168    /// dropping a superseded fixture.
169    pub fn remember_shared_resource(&self, resource: &str) {
170        let lock = ContainerStateLock::acquire(&self.suite, "postgres");
171        let mut state = lock.load();
172        state.remember_shared_resource(resource);
173        lock.save(&state);
174    }
175
176    /// Removes a suite-scoped product fixture after it was explicitly cleaned up.
177    pub fn forget_shared_resource(&self, resource: &str) {
178        let lock = ContainerStateLock::acquire(&self.suite, "postgres");
179        let mut state = lock.load();
180        state.forget_shared_resource(resource);
181        lock.save(&state);
182    }
183
184    async fn create_database_inner(
185        &self,
186        name: &str,
187        template: Option<&str>,
188        ownership: DatabaseOwnership,
189    ) -> PostgresTestDatabase {
190        assert_valid_database_name(name);
191        let lock = ContainerStateLock::acquire(&self.suite, "postgres");
192        let mut state = lock.load();
193        match ownership {
194            DatabaseOwnership::Process => state.remember_current_process_resource(name),
195            DatabaseOwnership::Shared => state.remember_shared_resource(name),
196        }
197        lock.save(&state);
198        drop(lock);
199
200        let admin = connect_with_retry(&self.admin_url, "PostgreSQL").await;
201        let create_database = create_database_statement(name, template);
202        admin
203            .execute_unprepared(&create_database)
204            .await
205            .unwrap_or_else(|error| {
206                panic!("failed to create PostgreSQL test database {name}: {error}")
207            });
208        admin
209            .close()
210            .await
211            .unwrap_or_else(|error| panic!("failed to close PostgreSQL admin connection: {error}"));
212
213        PostgresTestDatabase {
214            name: name.to_string(),
215            url: database_url(&self.admin_url, name),
216            admin_url: self.admin_url.clone(),
217            suite: self.suite.clone(),
218            ownership,
219        }
220    }
221
222    async fn cleanup_databases(&self, names: &[String]) {
223        if names.is_empty() {
224            return;
225        }
226        let admin = connect_with_retry(&self.admin_url, "PostgreSQL").await;
227        for name in names {
228            admin
229                .execute_unprepared(&format!(
230                    "DROP DATABASE IF EXISTS {} WITH (FORCE)",
231                    quote_identifier(name)
232                ))
233                .await
234                .unwrap_or_else(|error| {
235                    panic!("failed to drop stale PostgreSQL test database {name}: {error}")
236                });
237            let lock = ContainerStateLock::acquire(&self.suite, "postgres");
238            let mut state = lock.load();
239            state.forget_resource(std::process::id(), name);
240            lock.save(&state);
241        }
242        admin
243            .close()
244            .await
245            .unwrap_or_else(|error| panic!("failed to close PostgreSQL admin connection: {error}"));
246    }
247}
248
249async fn start_postgres_container_with_retry(
250    suite: &TestContainerSuite,
251) -> ContainerAsync<GenericImage> {
252    let mut last_error = None;
253    for _attempt in 0..240 {
254        match postgres_container_request(suite).start().await {
255            Ok(container) => return container,
256            Err(error) => {
257                last_error = Some(error.to_string());
258                tokio::time::sleep(std::time::Duration::from_millis(500)).await;
259            }
260        }
261    }
262
263    panic!(
264        "failed to start PostgreSQL test container after retries: {}",
265        last_error.unwrap_or_else(|| "unknown container startup error".to_string())
266    );
267}
268
269fn postgres_container_request(suite: &TestContainerSuite) -> ContainerRequest<GenericImage> {
270    GenericImage::new("postgres", "16")
271        .with_exposed_port(IntoContainerPort::tcp(5432))
272        .with_container_name(suite.container_name(POSTGRES_CONTAINER_SERVICE))
273        .with_reuse(ReuseDirective::Always)
274        .with_shm_size(POSTGRES_TEST_SHM_SIZE_BYTES)
275        .with_env_var("POSTGRES_USER", "postgres")
276        .with_env_var("POSTGRES_PASSWORD", "postgres")
277        .with_env_var("POSTGRES_DB", "postgres")
278}
279
280impl PostgresTestDatabase {
281    /// Returns the isolated database name.
282    #[must_use]
283    pub fn name(&self) -> &str {
284        &self.name
285    }
286
287    /// Returns the connection URL for this database.
288    #[must_use]
289    pub fn url(&self) -> &str {
290        &self.url
291    }
292
293    /// Connects to this database, retrying while the service becomes ready.
294    ///
295    /// # Panics
296    ///
297    /// Panics when the database does not accept a connection before the readiness timeout.
298    pub async fn connect(&self) -> DatabaseConnection {
299        connect_with_retry(&self.url, "PostgreSQL").await
300    }
301
302    /// Drops this database and removes it from the shared resource registry.
303    ///
304    /// # Panics
305    ///
306    /// Panics when the admin connection, database drop, connection shutdown, or shared-state
307    /// update fails.
308    pub async fn cleanup(&self) {
309        let admin = connect_with_retry(&self.admin_url, "PostgreSQL").await;
310        admin
311            .execute_unprepared(&format!(
312                "DROP DATABASE IF EXISTS {} WITH (FORCE)",
313                quote_identifier(&self.name)
314            ))
315            .await
316            .unwrap_or_else(|error| {
317                panic!(
318                    "failed to drop PostgreSQL test database {}: {error}",
319                    self.name
320                )
321            });
322        admin
323            .close()
324            .await
325            .unwrap_or_else(|error| panic!("failed to close PostgreSQL admin connection: {error}"));
326
327        let lock = ContainerStateLock::acquire(&self.suite, "postgres");
328        let mut state = lock.load();
329        match self.ownership {
330            DatabaseOwnership::Process => {
331                state.forget_resource(std::process::id(), &self.name);
332            }
333            DatabaseOwnership::Shared => state.forget_shared_resource(&self.name),
334        }
335        lock.save(&state);
336    }
337}
338
339fn database_url(admin_url: &str, name: &str) -> String {
340    admin_url.rsplit_once('/').map_or_else(
341        || admin_url.to_string(),
342        |(base, _)| format!("{base}/{name}"),
343    )
344}
345
346fn quote_identifier(value: &str) -> String {
347    format!("\"{}\"", value.replace('"', "\"\""))
348}
349
350fn create_database_statement(name: &str, template: Option<&str>) -> String {
351    template.map_or_else(
352        || format!("CREATE DATABASE {}", quote_identifier(name)),
353        |template| {
354            format!(
355                "CREATE DATABASE {} TEMPLATE {}",
356                quote_identifier(name),
357                quote_identifier(template)
358            )
359        },
360    )
361}
362
363fn assert_valid_database_name(name: &str) {
364    assert!(
365        !name.is_empty()
366            && name.len() <= 63
367            && name
368                .bytes()
369                .all(|byte| byte.is_ascii_alphanumeric() || byte == b'_'),
370        "PostgreSQL test database name must be 1-63 ASCII alphanumeric or '_' characters: {name:?}"
371    );
372}
373
374#[cfg(test)]
375mod tests {
376    use super::{
377        POSTGRES_CONTAINER_SERVICE, POSTGRES_TEST_SHM_SIZE_BYTES, assert_valid_database_name,
378        create_database_statement, database_url, postgres_container_request, quote_identifier,
379    };
380    use crate::suite::TestContainerSuite;
381    use testcontainers::{core::ExecCommand, runners::AsyncRunner};
382
383    #[test]
384    fn database_url_replaces_admin_database() {
385        assert_eq!(
386            database_url("postgres://user:pass@127.0.0.1:5432/postgres", "isolated"),
387            "postgres://user:pass@127.0.0.1:5432/isolated"
388        );
389    }
390
391    #[test]
392    fn identifier_quoting_escapes_quotes() {
393        assert_eq!(quote_identifier("test\"name"), "\"test\"\"name\"");
394    }
395
396    #[test]
397    fn postgres_container_request_sets_versioned_name_and_shared_memory() {
398        let suite = TestContainerSuite::new("forge-postgres-request");
399        let request = postgres_container_request(&suite);
400        let expected_name = suite.container_name(POSTGRES_CONTAINER_SERVICE);
401
402        assert_eq!(request.shm_size(), Some(POSTGRES_TEST_SHM_SIZE_BYTES));
403        assert_eq!(
404            request.container_name().as_deref(),
405            Some(expected_name.as_str())
406        );
407    }
408
409    #[tokio::test]
410    async fn postgres_container_exposes_configured_shared_memory() {
411        let suite = TestContainerSuite::new("forge-postgres-shm");
412        let container = postgres_container_request(&suite)
413            .start()
414            .await
415            .expect("PostgreSQL test container should start");
416        let mut command = container
417            .exec(ExecCommand::new(["df", "-B1", "--output=size", "/dev/shm"]))
418            .await
419            .expect("shared-memory capacity command should start");
420        let stdout = command
421            .stdout_to_vec()
422            .await
423            .expect("shared-memory capacity command should finish");
424        let output = String::from_utf8(stdout).expect("df output should be UTF-8");
425        let capacity = output
426            .lines()
427            .find_map(|line| line.trim().parse::<u64>().ok())
428            .expect("df output should contain the shared-memory capacity");
429
430        assert!(
431            capacity >= POSTGRES_TEST_SHM_SIZE_BYTES,
432            "PostgreSQL test container shared memory must be at least {POSTGRES_TEST_SHM_SIZE_BYTES} bytes, got {capacity}"
433        );
434    }
435
436    #[test]
437    fn database_creation_supports_an_optional_template() {
438        assert_eq!(
439            create_database_statement("isolated", None),
440            "CREATE DATABASE \"isolated\""
441        );
442        assert_eq!(
443            create_database_statement("isolated", Some("template")),
444            "CREATE DATABASE \"isolated\" TEMPLATE \"template\""
445        );
446    }
447
448    #[test]
449    fn database_name_accepts_boundaries() {
450        assert_valid_database_name("a");
451        assert_valid_database_name(&"a".repeat(63));
452        assert_valid_database_name("aster_product_123");
453    }
454
455    #[test]
456    fn database_name_rejects_unsafe_or_oversized_values() {
457        for name in ["", "has-hyphen", "has quote\"", "has space"] {
458            assert!(
459                std::panic::catch_unwind(|| assert_valid_database_name(name)).is_err(),
460                "database name {name:?} should be rejected"
461            );
462        }
463        let oversized = "a".repeat(64);
464        assert!(std::panic::catch_unwind(|| assert_valid_database_name(&oversized)).is_err());
465    }
466}
aster_forge_test/postgres.rs

aster_forge_test/
postgres.rs
