pg-embed-setup-unpriv 0.5.2

//! Database lifecycle operations for `TestClusterConnection`.
//!
//! This module provides methods for creating, dropping, and managing databases
//! on a running `PostgreSQL` cluster.

use color_eyre::eyre::WrapErr;
use tracing::info_span;

use super::{
    connection::{TestClusterConnection, escape_identifier},
    lifecycle_template::{
        STD_TEMPLATE_LOCKS,
        TemplateCreationOps,
        ensure_template_exists_with_lock,
    },
    temporary_database::TemporaryDatabase,
};
use crate::error::BootstrapResult;

/// A strongly-typed database name for use with lifecycle operations.
///
/// This newtype provides type safety for database name parameters, preventing
/// accidental misuse of raw strings while still allowing convenient conversion
/// from string literals.
///
/// # Examples
///
/// ```
/// use pg_embedded_setup_unpriv::DatabaseName;
///
/// // From string literal
/// let name: DatabaseName = "my_database".into();
/// assert_eq!(name.as_str(), "my_database");
///
/// // From owned String
/// let name: DatabaseName = String::from("another_db").into();
/// assert_eq!(name.as_str(), "another_db");
/// ```
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct DatabaseName(String);

impl DatabaseName {
    /// Creates a new `DatabaseName` from a string.
    #[must_use]
    pub fn new(name: impl Into<String>) -> Self { Self(name.into()) }

    /// Returns the database name as a string slice.
    #[must_use]
    pub fn as_str(&self) -> &str { &self.0 }
}

impl AsRef<str> for DatabaseName {
    fn as_ref(&self) -> &str { &self.0 }
}

impl From<&str> for DatabaseName {
    fn from(s: &str) -> Self { Self(s.to_owned()) }
}

impl From<String> for DatabaseName {
    fn from(s: String) -> Self { Self(s) }
}
impl TestClusterConnection {
    /// Executes a DDL command for database creation or deletion.
    ///
    /// This private helper consolidates the common pattern of escaping an
    /// identifier, formatting SQL, and executing it via `batch_execute`.
    fn execute_ddl_command(
        &self,
        sql_template: &str,
        name: &str,
        error_msg_verb: &str,
    ) -> BootstrapResult<()> {
        let mut client = self.admin_client()?;
        let escaped = escape_identifier(name);
        let sql = sql_template.replace("{}", &format!("\"{escaped}\""));
        client
            .batch_execute(&sql)
            .wrap_err(format!("failed to {error_msg_verb} database '{name}'"))
            .map_err(crate::error::BootstrapError::from)
    }

    /// Creates a new database with the given name.
    ///
    /// Connects to the `postgres` database as superuser and executes
    /// `CREATE DATABASE`.
    ///
    /// # Errors
    ///
    /// Returns an error if the database already exists or if the connection
    /// fails.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use pg_embedded_setup_unpriv::TestCluster;
    ///
    /// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
    /// let cluster = TestCluster::new()?;
    /// cluster.connection().create_database("my_test_db")?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn create_database(&self, name: impl Into<DatabaseName>) -> BootstrapResult<()> {
        let db_name = name.into();
        let _span = info_span!("create_database", db = %db_name.as_str()).entered();
        self.execute_ddl_command("CREATE DATABASE {}", db_name.as_str(), "create")
    }

    /// Creates a new database by cloning an existing template.
    ///
    /// Connects to the `postgres` database as superuser and executes
    /// `CREATE DATABASE ... TEMPLATE`. This is significantly faster than
    /// creating an empty database and running migrations, as `PostgreSQL`
    /// performs a filesystem-level copy.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - The target database already exists
    /// - The template database does not exist
    /// - The template database has active connections
    /// - The connection fails
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use pg_embedded_setup_unpriv::TestCluster;
    ///
    /// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
    /// let cluster = TestCluster::new()?;
    ///
    /// // Create and set up a template database
    /// cluster.connection().create_database("my_template")?;
    /// // ... run migrations on my_template ...
    ///
    /// // Clone the template for a test
    /// cluster
    ///     .connection()
    ///     .create_database_from_template("test_db", "my_template")?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn create_database_from_template(
        &self,
        name: impl Into<DatabaseName>,
        template: impl Into<DatabaseName>,
    ) -> BootstrapResult<()> {
        let db_name = name.into();
        let template_name = template.into();
        let _span =
            info_span!("create_database_from_template", db = %db_name.as_str(), template = %template_name.as_str()).entered();
        let mut client = self.admin_client()?;
        let escaped_name = escape_identifier(db_name.as_str());
        let escaped_template = escape_identifier(template_name.as_str());
        let sql = format!("CREATE DATABASE \"{escaped_name}\" TEMPLATE \"{escaped_template}\"");
        client
            .batch_execute(&sql)
            .wrap_err(format!(
                "failed to create database '{}' from template '{}'",
                db_name.as_str(),
                template_name.as_str()
            ))
            .map_err(crate::error::BootstrapError::from)
    }

    /// Drops an existing database.
    ///
    /// Connects to the `postgres` database as superuser and executes
    /// `DROP DATABASE`.
    ///
    /// # Errors
    ///
    /// Returns an error if the database does not exist, has active connections,
    /// or if the connection fails.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use pg_embedded_setup_unpriv::TestCluster;
    ///
    /// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
    /// let cluster = TestCluster::new()?;
    /// cluster.connection().create_database("temp_db")?;
    /// cluster.connection().drop_database("temp_db")?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn drop_database(&self, name: impl Into<DatabaseName>) -> BootstrapResult<()> {
        let db_name = name.into();
        let _span = info_span!("drop_database", db = %db_name.as_str()).entered();
        self.execute_ddl_command("DROP DATABASE {}", db_name.as_str(), "drop")
    }

    /// Checks whether a database with the given name exists.
    ///
    /// # Errors
    ///
    /// Returns an error if the connection fails.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use pg_embedded_setup_unpriv::TestCluster;
    ///
    /// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
    /// let cluster = TestCluster::new()?;
    /// assert!(cluster.connection().database_exists("postgres")?);
    /// assert!(!cluster.connection().database_exists("nonexistent")?);
    /// # Ok(())
    /// # }
    /// ```
    pub fn database_exists(&self, name: impl Into<DatabaseName>) -> BootstrapResult<bool> {
        let db_name = name.into();
        let mut client = self.admin_client()?;
        let row = client
            .query_one(
                "SELECT EXISTS(SELECT 1 FROM pg_database WHERE datname = $1)",
                &[&db_name.as_str()],
            )
            .wrap_err("failed to query pg_database")
            .map_err(crate::error::BootstrapError::from)?;
        Ok(row.get(0))
    }

    /// Ensures a template database exists, creating it if necessary.
    ///
    /// Uses per-template locking to prevent concurrent creation attempts when
    /// multiple tests race to initialize the same template. The `setup_fn` is
    /// called only if the template does not already exist.
    ///
    /// # Errors
    ///
    /// Returns an error if database creation fails or if `setup_fn` returns
    /// an error. If setup fails or panics after this call creates the template,
    /// the partially created template is dropped before the error is returned
    /// or the panic resumes.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use pg_embedded_setup_unpriv::TestCluster;
    ///
    /// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
    /// let cluster = TestCluster::new()?;
    ///
    /// // Ensure template exists, running migrations if needed
    /// cluster
    ///     .connection()
    ///     .ensure_template_exists("my_template", |db_name| {
    ///         // Run migrations on the newly created template database
    ///         // e.g., diesel::migration::run(&mut conn)?;
    ///         Ok(())
    ///     })?;
    ///
    /// // Clone the template for each test
    /// cluster
    ///     .connection()
    ///     .create_database_from_template("test_db_1", "my_template")?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn ensure_template_exists<F>(
        &self,
        name: impl Into<DatabaseName>,
        setup_fn: F,
    ) -> BootstrapResult<()>
    where
        F: FnOnce(&str) -> BootstrapResult<()>,
    {
        let db_name = name.into();
        let _span = info_span!("ensure_template_exists", template = %db_name.as_str()).entered();
        ensure_template_exists_with_lock(
            &STD_TEMPLATE_LOCKS,
            db_name.as_str(),
            TemplateCreationOps {
                database_exists: || self.database_exists(db_name.as_str()),
                create_database: || self.create_database(db_name.as_str()),
                drop_database: || self.drop_database(db_name.as_str()),
                setup_fn: || setup_fn(db_name.as_str()),
            },
        )
    }

    /// Creates a temporary database that is dropped when the guard is dropped.
    ///
    /// This is useful for test isolation where each test creates its own
    /// database and the database is automatically cleaned up when the test
    /// completes.
    ///
    /// # Errors
    ///
    /// Returns an error if the database already exists or if the connection
    /// fails.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use pg_embedded_setup_unpriv::TestCluster;
    ///
    /// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
    /// let cluster = TestCluster::new()?;
    /// let temp_db = cluster.connection().temporary_database("my_temp_db")?;
    ///
    /// // Database is dropped automatically when temp_db goes out of scope
    /// let url = temp_db.url();
    /// # Ok(())
    /// # }
    /// ```
    pub fn temporary_database(
        &self,
        name: impl Into<DatabaseName>,
    ) -> BootstrapResult<TemporaryDatabase> {
        let db_name = name.into();
        let _span = info_span!("temporary_database", db = %db_name.as_str()).entered();
        self.create_database(db_name.as_str())?;
        Ok(TemporaryDatabase::new(
            db_name.as_str().to_owned(),
            self.database_url("postgres"),
            self.database_url(db_name.as_str()),
        ))
    }

    /// Creates a temporary database from a template.
    ///
    /// Combines template cloning with RAII cleanup. The database is created
    /// by cloning the template and is automatically dropped when the guard
    /// goes out of scope.
    ///
    /// # Errors
    ///
    /// Returns an error if the target database already exists, the template
    /// does not exist, the template has active connections, or if the
    /// connection fails.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use pg_embedded_setup_unpriv::TestCluster;
    ///
    /// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
    /// let cluster = TestCluster::new()?;
    ///
    /// // Create and migrate a template once
    /// cluster.ensure_template_exists("migrated_template", |_| Ok(()))?;
    ///
    /// // Each test gets its own database cloned from the template
    /// let temp_db = cluster
    ///     .connection()
    ///     .temporary_database_from_template("test_db", "migrated_template")?;
    ///
    /// // Database is dropped automatically when temp_db goes out of scope
    /// # Ok(())
    /// # }
    /// ```
    pub fn temporary_database_from_template(
        &self,
        name: impl Into<DatabaseName>,
        template: impl Into<DatabaseName>,
    ) -> BootstrapResult<TemporaryDatabase> {
        let db_name = name.into();
        let template_name = template.into();
        let _span =
            info_span!("temporary_database_from_template", db = %db_name.as_str(), template = %template_name.as_str())
                .entered();
        self.create_database_from_template(db_name.as_str(), template_name.as_str())?;
        Ok(TemporaryDatabase::new(
            db_name.as_str().to_owned(),
            self.database_url("postgres"),
            self.database_url(db_name.as_str()),
        ))
    }
}

#[cfg(test)]
mod tests {
    //! Unit tests for the `DatabaseName` newtype.

    use super::DatabaseName;

    #[test]
    fn database_name_constructors_agree() {
        let from_new = DatabaseName::new("analytics");
        let from_str: DatabaseName = "analytics".into();
        let from_string: DatabaseName = String::from("analytics").into();

        assert_eq!(from_new.as_str(), "analytics");
        assert_eq!(from_new, from_str);
        assert_eq!(from_str, from_string);
        assert_eq!(AsRef::<str>::as_ref(&from_string), "analytics");
    }
}