pg_embedded_setup_unpriv/cluster/
temporary_database.rs

1//! RAII guard for automatic database cleanup.
2//!
3//! `TemporaryDatabase` drops its associated database when the guard goes out of
4//! scope, mirroring the `TestCluster` lifecycle semantics.
5
6use color_eyre::eyre::WrapErr;
7use tracing::info_span;
8
9use super::connection::{connect_admin, escape_identifier};
10use crate::{error::BootstrapResult, observability::LOG_TARGET};
11
12/// RAII guard that drops a database when it goes out of scope.
13///
14/// The guard stores the database name and connection URL rather than borrowing
15/// a connection, avoiding lifetime issues and allowing reconnection in `Drop`.
16///
17/// # Examples
18///
19/// ```no_run
20/// use pg_embedded_setup_unpriv::TestCluster;
21///
22/// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
23/// let cluster = TestCluster::new()?;
24///
25/// // Create a temporary database that is dropped when the guard is dropped
26/// let temp_db = cluster.connection().temporary_database("my_temp_db")?;
27///
28/// // Use the database
29/// let url = temp_db.url();
30/// // ... run queries ...
31///
32/// // Database is dropped automatically when temp_db goes out of scope
33/// drop(temp_db);
34/// # Ok(())
35/// # }
36/// ```
37#[derive(Debug)]
38pub struct TemporaryDatabase {
39    name: String,
40    admin_url: String,
41    database_url: String,
42}
43
44impl TemporaryDatabase {
45    /// Creates a new `TemporaryDatabase` guard.
46    ///
47    /// This constructor is intended for internal use. Prefer using
48    /// [`TestClusterConnection::temporary_database`] or
49    /// [`TestClusterConnection::temporary_database_from_template`].
50    pub(crate) const fn new(name: String, admin_url: String, database_url: String) -> Self {
51        Self {
52            name,
53            admin_url,
54            database_url,
55        }
56    }
57
58    /// Returns the database name.
59    #[must_use]
60    pub fn name(&self) -> &str { &self.name }
61
62    /// Returns the connection URL for this database.
63    #[must_use]
64    pub fn url(&self) -> &str { &self.database_url }
65
66    /// Drops the database, failing if connections exist.
67    ///
68    /// This mirrors `PostgreSQL`'s native behaviour where `DROP DATABASE` fails
69    /// if there are active connections to the database.
70    ///
71    /// # Errors
72    ///
73    /// Returns an error if:
74    /// - The database has active connections
75    /// - The database does not exist
76    /// - The connection to the admin database fails
77    ///
78    /// # Examples
79    ///
80    /// ```no_run
81    /// use pg_embedded_setup_unpriv::TestCluster;
82    ///
83    /// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
84    /// let cluster = TestCluster::new()?;
85    /// let temp_db = cluster.connection().temporary_database("my_temp_db")?;
86    ///
87    /// // Explicitly drop (consumes the guard)
88    /// temp_db.drop_database()?;
89    /// # Ok(())
90    /// # }
91    /// ```
92    pub fn drop_database(self) -> BootstrapResult<()> { self.try_drop() }
93
94    /// Drops the database, terminating any active connections first.
95    ///
96    /// This is useful when you need to ensure the database is dropped even if
97    /// there are lingering connections (e.g., from connection pools that
98    /// haven't been drained).
99    ///
100    /// # Errors
101    ///
102    /// Returns an error if:
103    /// - The database does not exist
104    /// - The connection to the admin database fails
105    /// - Terminating connections fails
106    ///
107    /// # Examples
108    ///
109    /// ```no_run
110    /// use pg_embedded_setup_unpriv::TestCluster;
111    ///
112    /// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
113    /// let cluster = TestCluster::new()?;
114    /// let temp_db = cluster.connection().temporary_database("my_temp_db")?;
115    ///
116    /// // Force drop even if connections exist
117    /// temp_db.force_drop()?;
118    /// # Ok(())
119    /// # }
120    /// ```
121    pub fn force_drop(self) -> BootstrapResult<()> {
122        let _span = info_span!("force_drop_database", db = %self.name).entered();
123        let mut client = connect_admin(&self.admin_url)?;
124
125        // Terminate active connections using parameterized query
126        client
127            .execute(
128                "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = $1 AND \
129                 pid <> pg_backend_pid()",
130                &[&self.name],
131            )
132            .wrap_err(format!(
133                "failed to terminate connections to database '{}'",
134                self.name
135            ))
136            .map_err(crate::error::BootstrapError::from)?;
137
138        // Drop the database with escaped identifier
139        let escaped = escape_identifier(&self.name);
140        let drop_sql = format!("DROP DATABASE \"{escaped}\"");
141        client
142            .batch_execute(&drop_sql)
143            .wrap_err(format!("failed to drop database '{}'", self.name))
144            .map_err(crate::error::BootstrapError::from)?;
145
146        Ok(())
147    }
148
149    /// Attempts to drop the database without consuming self.
150    ///
151    /// Used by the `Drop` implementation for best-effort cleanup.
152    fn try_drop(&self) -> BootstrapResult<()> {
153        let _span = info_span!("drop_database", db = %self.name).entered();
154        let mut client = connect_admin(&self.admin_url)?;
155
156        let escaped = escape_identifier(&self.name);
157        let sql = format!("DROP DATABASE \"{escaped}\"");
158        client
159            .batch_execute(&sql)
160            .wrap_err(format!("failed to drop database '{}'", self.name))
161            .map_err(crate::error::BootstrapError::from)?;
162
163        Ok(())
164    }
165}
166
167impl Drop for TemporaryDatabase {
168    fn drop(&mut self) {
169        if let Err(e) = self.try_drop() {
170            tracing::warn!(
171                target: LOG_TARGET,
172                db = %self.name,
173                error = ?e,
174                "failed to drop temporary database"
175            );
176        }
177    }
178}
179
180#[cfg(test)]
181mod tests {
182    //! Tests for temporary database handling.
183    use super::*;
184
185    #[test]
186    fn temporary_database_accessors() {
187        let temp = TemporaryDatabase::new(
188            "test_db".to_owned(),
189            "postgresql://user:pass@localhost:5432/postgres".to_owned(),
190            "postgresql://user:pass@localhost:5432/test_db".to_owned(),
191        );
192
193        assert_eq!(temp.name(), "test_db");
194        assert!(temp.url().contains("test_db"));
195    }
196
197    /// An admin URL nothing listens on, with the attempt bounded.
198    ///
199    /// The port is zero because no server can listen there: it is the
200    /// kernel's request for an ephemeral port, never a stable listening
201    /// one. A high fixed port is only probably free, and a developer
202    /// whose own `PostgreSQL` happened to hold it would have these tests
203    /// connect to it and run `DROP DATABASE test_db` against a real
204    /// server. That needs the URL's credentials to be accepted as well,
205    /// so it is a narrow local risk rather than a likely one, but the
206    /// port costs nothing to make impossible.
207    ///
208    /// Two further details, both load-bearing on Windows. The host is
209    /// `127.0.0.1` rather than `localhost`, because the name resolves to
210    /// `::1` as well and the client tries the addresses in turn, so a
211    /// host that refuses one slowly is waited on before the other is
212    /// reached. `connect_timeout` then bounds what is left.
213    ///
214    /// Without them the test relies on the operating system refusing a
215    /// connection promptly, which Linux and macOS do in 0.01 to 0.04 s
216    /// and Windows does not: it measured 8.07 s on run 34161062265 and
217    /// then exceeded the profile's 180 s per-test allowance on run
218    /// 34274541426, ending the whole Windows lane.
219    const UNREACHABLE_ADMIN_URL: &str =
220        "postgresql://user:pass@127.0.0.1:0/postgres?connect_timeout=2";
221
222    /// The matching database URL. See [`UNREACHABLE_ADMIN_URL`].
223    const UNREACHABLE_DB_URL: &str = "postgresql://user:pass@127.0.0.1:0/test_db?connect_timeout=2";
224
225    #[test]
226    fn drop_database_returns_error_on_connection_failure() {
227        let temp = TemporaryDatabase::new(
228            "test_db".to_owned(),
229            UNREACHABLE_ADMIN_URL.to_owned(),
230            UNREACHABLE_DB_URL.to_owned(),
231        );
232
233        let result = temp.drop_database();
234        let Err(err) = result else {
235            panic!("expected error when database unreachable");
236        };
237        let err_str = err.to_string();
238        assert!(
239            err_str.contains("failed to connect"),
240            "expected connection failure, got: {err_str}"
241        );
242    }
243
244    #[test]
245    fn force_drop_returns_error_on_connection_failure() {
246        let temp = TemporaryDatabase::new(
247            "test_db".to_owned(),
248            UNREACHABLE_ADMIN_URL.to_owned(),
249            UNREACHABLE_DB_URL.to_owned(),
250        );
251
252        let result = temp.force_drop();
253        let Err(err) = result else {
254            panic!("expected error when database unreachable");
255        };
256        let err_str = err.to_string();
257        assert!(
258            err_str.contains("failed to connect"),
259            "expected connection failure, got: {err_str}"
260        );
261    }
262
263    #[test]
264    fn drop_trait_does_not_panic_on_connection_failure() {
265        // Create a TemporaryDatabase with an unreachable URL
266        let temp = TemporaryDatabase::new(
267            "test_db".to_owned(),
268            UNREACHABLE_ADMIN_URL.to_owned(),
269            UNREACHABLE_DB_URL.to_owned(),
270        );
271
272        // Dropping should not panic even when cleanup fails
273        // The Drop impl logs a warning but does not propagate errors
274        drop(temp);
275    }
276}
pg_embedded_setup_unpriv/cluster/temporary_database.rs

pg_embedded_setup_unpriv/cluster/
temporary_database.rs
