pg_embedded_setup_unpriv/cluster/
lifecycle.rs

1//! Database lifecycle operations for `TestClusterConnection`.
2//!
3//! This module provides methods for creating, dropping, and managing databases
4//! on a running `PostgreSQL` cluster.
5
6use color_eyre::eyre::WrapErr;
7use tracing::info_span;
8
9use super::{
10    connection::{TestClusterConnection, escape_identifier},
11    lifecycle_template::{
12        STD_TEMPLATE_LOCKS,
13        TemplateCreationOps,
14        ensure_template_exists_with_lock,
15    },
16    temporary_database::TemporaryDatabase,
17};
18use crate::error::BootstrapResult;
19
20/// A strongly-typed database name for use with lifecycle operations.
21///
22/// This newtype provides type safety for database name parameters, preventing
23/// accidental misuse of raw strings while still allowing convenient conversion
24/// from string literals.
25///
26/// # Examples
27///
28/// ```
29/// use pg_embedded_setup_unpriv::DatabaseName;
30///
31/// // From string literal
32/// let name: DatabaseName = "my_database".into();
33/// assert_eq!(name.as_str(), "my_database");
34///
35/// // From owned String
36/// let name: DatabaseName = String::from("another_db").into();
37/// assert_eq!(name.as_str(), "another_db");
38/// ```
39#[derive(Debug, Clone, PartialEq, Eq, Hash)]
40pub struct DatabaseName(String);
41
42impl DatabaseName {
43    /// Creates a new `DatabaseName` from a string.
44    #[must_use]
45    pub fn new(name: impl Into<String>) -> Self { Self(name.into()) }
46
47    /// Returns the database name as a string slice.
48    #[must_use]
49    pub fn as_str(&self) -> &str { &self.0 }
50}
51
52impl AsRef<str> for DatabaseName {
53    fn as_ref(&self) -> &str { &self.0 }
54}
55
56impl From<&str> for DatabaseName {
57    fn from(s: &str) -> Self { Self(s.to_owned()) }
58}
59
60impl From<String> for DatabaseName {
61    fn from(s: String) -> Self { Self(s) }
62}
63impl TestClusterConnection {
64    /// Executes a DDL command for database creation or deletion.
65    ///
66    /// This private helper consolidates the common pattern of escaping an
67    /// identifier, formatting SQL, and executing it via `batch_execute`.
68    fn execute_ddl_command(
69        &self,
70        sql_template: &str,
71        name: &str,
72        error_msg_verb: &str,
73    ) -> BootstrapResult<()> {
74        let mut client = self.admin_client()?;
75        let escaped = escape_identifier(name);
76        let sql = sql_template.replace("{}", &format!("\"{escaped}\""));
77        client
78            .batch_execute(&sql)
79            .wrap_err(format!("failed to {error_msg_verb} database '{name}'"))
80            .map_err(crate::error::BootstrapError::from)
81    }
82
83    /// Creates a new database with the given name.
84    ///
85    /// Connects to the `postgres` database as superuser and executes
86    /// `CREATE DATABASE`.
87    ///
88    /// # Errors
89    ///
90    /// Returns an error if the database already exists or if the connection
91    /// fails.
92    ///
93    /// # Examples
94    ///
95    /// ```no_run
96    /// use pg_embedded_setup_unpriv::TestCluster;
97    ///
98    /// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
99    /// let cluster = TestCluster::new()?;
100    /// cluster.connection().create_database("my_test_db")?;
101    /// # Ok(())
102    /// # }
103    /// ```
104    pub fn create_database(&self, name: impl Into<DatabaseName>) -> BootstrapResult<()> {
105        let db_name = name.into();
106        let _span = info_span!("create_database", db = %db_name.as_str()).entered();
107        self.execute_ddl_command("CREATE DATABASE {}", db_name.as_str(), "create")
108    }
109
110    /// Creates a new database by cloning an existing template.
111    ///
112    /// Connects to the `postgres` database as superuser and executes
113    /// `CREATE DATABASE ... TEMPLATE`. This is significantly faster than
114    /// creating an empty database and running migrations, as `PostgreSQL`
115    /// performs a filesystem-level copy.
116    ///
117    /// # Errors
118    ///
119    /// Returns an error if:
120    /// - The target database already exists
121    /// - The template database does not exist
122    /// - The template database has active connections
123    /// - The connection fails
124    ///
125    /// # Examples
126    ///
127    /// ```no_run
128    /// use pg_embedded_setup_unpriv::TestCluster;
129    ///
130    /// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
131    /// let cluster = TestCluster::new()?;
132    ///
133    /// // Create and set up a template database
134    /// cluster.connection().create_database("my_template")?;
135    /// // ... run migrations on my_template ...
136    ///
137    /// // Clone the template for a test
138    /// cluster
139    ///     .connection()
140    ///     .create_database_from_template("test_db", "my_template")?;
141    /// # Ok(())
142    /// # }
143    /// ```
144    pub fn create_database_from_template(
145        &self,
146        name: impl Into<DatabaseName>,
147        template: impl Into<DatabaseName>,
148    ) -> BootstrapResult<()> {
149        let db_name = name.into();
150        let template_name = template.into();
151        let _span =
152            info_span!("create_database_from_template", db = %db_name.as_str(), template = %template_name.as_str()).entered();
153        let mut client = self.admin_client()?;
154        let escaped_name = escape_identifier(db_name.as_str());
155        let escaped_template = escape_identifier(template_name.as_str());
156        let sql = format!("CREATE DATABASE \"{escaped_name}\" TEMPLATE \"{escaped_template}\"");
157        client
158            .batch_execute(&sql)
159            .wrap_err(format!(
160                "failed to create database '{}' from template '{}'",
161                db_name.as_str(),
162                template_name.as_str()
163            ))
164            .map_err(crate::error::BootstrapError::from)
165    }
166
167    /// Drops an existing database.
168    ///
169    /// Connects to the `postgres` database as superuser and executes
170    /// `DROP DATABASE`.
171    ///
172    /// # Errors
173    ///
174    /// Returns an error if the database does not exist, has active connections,
175    /// or if the connection fails.
176    ///
177    /// # Examples
178    ///
179    /// ```no_run
180    /// use pg_embedded_setup_unpriv::TestCluster;
181    ///
182    /// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
183    /// let cluster = TestCluster::new()?;
184    /// cluster.connection().create_database("temp_db")?;
185    /// cluster.connection().drop_database("temp_db")?;
186    /// # Ok(())
187    /// # }
188    /// ```
189    pub fn drop_database(&self, name: impl Into<DatabaseName>) -> BootstrapResult<()> {
190        let db_name = name.into();
191        let _span = info_span!("drop_database", db = %db_name.as_str()).entered();
192        self.execute_ddl_command("DROP DATABASE {}", db_name.as_str(), "drop")
193    }
194
195    /// Checks whether a database with the given name exists.
196    ///
197    /// # Errors
198    ///
199    /// Returns an error if the connection fails.
200    ///
201    /// # Examples
202    ///
203    /// ```no_run
204    /// use pg_embedded_setup_unpriv::TestCluster;
205    ///
206    /// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
207    /// let cluster = TestCluster::new()?;
208    /// assert!(cluster.connection().database_exists("postgres")?);
209    /// assert!(!cluster.connection().database_exists("nonexistent")?);
210    /// # Ok(())
211    /// # }
212    /// ```
213    pub fn database_exists(&self, name: impl Into<DatabaseName>) -> BootstrapResult<bool> {
214        let db_name = name.into();
215        let mut client = self.admin_client()?;
216        let row = client
217            .query_one(
218                "SELECT EXISTS(SELECT 1 FROM pg_database WHERE datname = $1)",
219                &[&db_name.as_str()],
220            )
221            .wrap_err("failed to query pg_database")
222            .map_err(crate::error::BootstrapError::from)?;
223        Ok(row.get(0))
224    }
225
226    /// Ensures a template database exists, creating it if necessary.
227    ///
228    /// Uses per-template locking to prevent concurrent creation attempts when
229    /// multiple tests race to initialize the same template. The `setup_fn` is
230    /// called only if the template does not already exist.
231    ///
232    /// # Errors
233    ///
234    /// Returns an error if database creation fails or if `setup_fn` returns
235    /// an error. If setup fails or panics after this call creates the template,
236    /// the partially created template is dropped before the error is returned
237    /// or the panic resumes.
238    ///
239    /// # Examples
240    ///
241    /// ```no_run
242    /// use pg_embedded_setup_unpriv::TestCluster;
243    ///
244    /// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
245    /// let cluster = TestCluster::new()?;
246    ///
247    /// // Ensure template exists, running migrations if needed
248    /// cluster
249    ///     .connection()
250    ///     .ensure_template_exists("my_template", |db_name| {
251    ///         // Run migrations on the newly created template database
252    ///         // e.g., diesel::migration::run(&mut conn)?;
253    ///         Ok(())
254    ///     })?;
255    ///
256    /// // Clone the template for each test
257    /// cluster
258    ///     .connection()
259    ///     .create_database_from_template("test_db_1", "my_template")?;
260    /// # Ok(())
261    /// # }
262    /// ```
263    pub fn ensure_template_exists<F>(
264        &self,
265        name: impl Into<DatabaseName>,
266        setup_fn: F,
267    ) -> BootstrapResult<()>
268    where
269        F: FnOnce(&str) -> BootstrapResult<()>,
270    {
271        let db_name = name.into();
272        let _span = info_span!("ensure_template_exists", template = %db_name.as_str()).entered();
273        ensure_template_exists_with_lock(
274            &STD_TEMPLATE_LOCKS,
275            db_name.as_str(),
276            TemplateCreationOps {
277                database_exists: || self.database_exists(db_name.as_str()),
278                create_database: || self.create_database(db_name.as_str()),
279                drop_database: || self.drop_database(db_name.as_str()),
280                setup_fn: || setup_fn(db_name.as_str()),
281            },
282        )
283    }
284
285    /// Creates a temporary database that is dropped when the guard is dropped.
286    ///
287    /// This is useful for test isolation where each test creates its own
288    /// database and the database is automatically cleaned up when the test
289    /// completes.
290    ///
291    /// # Errors
292    ///
293    /// Returns an error if the database already exists or if the connection
294    /// fails.
295    ///
296    /// # Examples
297    ///
298    /// ```no_run
299    /// use pg_embedded_setup_unpriv::TestCluster;
300    ///
301    /// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
302    /// let cluster = TestCluster::new()?;
303    /// let temp_db = cluster.connection().temporary_database("my_temp_db")?;
304    ///
305    /// // Database is dropped automatically when temp_db goes out of scope
306    /// let url = temp_db.url();
307    /// # Ok(())
308    /// # }
309    /// ```
310    pub fn temporary_database(
311        &self,
312        name: impl Into<DatabaseName>,
313    ) -> BootstrapResult<TemporaryDatabase> {
314        let db_name = name.into();
315        let _span = info_span!("temporary_database", db = %db_name.as_str()).entered();
316        self.create_database(db_name.as_str())?;
317        Ok(TemporaryDatabase::new(
318            db_name.as_str().to_owned(),
319            self.database_url("postgres"),
320            self.database_url(db_name.as_str()),
321        ))
322    }
323
324    /// Creates a temporary database from a template.
325    ///
326    /// Combines template cloning with RAII cleanup. The database is created
327    /// by cloning the template and is automatically dropped when the guard
328    /// goes out of scope.
329    ///
330    /// # Errors
331    ///
332    /// Returns an error if the target database already exists, the template
333    /// does not exist, the template has active connections, or if the
334    /// connection fails.
335    ///
336    /// # Examples
337    ///
338    /// ```no_run
339    /// use pg_embedded_setup_unpriv::TestCluster;
340    ///
341    /// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
342    /// let cluster = TestCluster::new()?;
343    ///
344    /// // Create and migrate a template once
345    /// cluster.ensure_template_exists("migrated_template", |_| Ok(()))?;
346    ///
347    /// // Each test gets its own database cloned from the template
348    /// let temp_db = cluster
349    ///     .connection()
350    ///     .temporary_database_from_template("test_db", "migrated_template")?;
351    ///
352    /// // Database is dropped automatically when temp_db goes out of scope
353    /// # Ok(())
354    /// # }
355    /// ```
356    pub fn temporary_database_from_template(
357        &self,
358        name: impl Into<DatabaseName>,
359        template: impl Into<DatabaseName>,
360    ) -> BootstrapResult<TemporaryDatabase> {
361        let db_name = name.into();
362        let template_name = template.into();
363        let _span =
364            info_span!("temporary_database_from_template", db = %db_name.as_str(), template = %template_name.as_str())
365                .entered();
366        self.create_database_from_template(db_name.as_str(), template_name.as_str())?;
367        Ok(TemporaryDatabase::new(
368            db_name.as_str().to_owned(),
369            self.database_url("postgres"),
370            self.database_url(db_name.as_str()),
371        ))
372    }
373}
374
375#[cfg(test)]
376mod tests {
377    //! Unit tests for the `DatabaseName` newtype.
378
379    use super::DatabaseName;
380
381    #[test]
382    fn database_name_constructors_agree() {
383        let from_new = DatabaseName::new("analytics");
384        let from_str: DatabaseName = "analytics".into();
385        let from_string: DatabaseName = String::from("analytics").into();
386
387        assert_eq!(from_new.as_str(), "analytics");
388        assert_eq!(from_new, from_str);
389        assert_eq!(from_str, from_string);
390        assert_eq!(AsRef::<str>::as_ref(&from_string), "analytics");
391    }
392}
pg_embedded_setup_unpriv/cluster/lifecycle.rs

pg_embedded_setup_unpriv/cluster/
lifecycle.rs
