integresql 0.2.0

use super::client::Client;
use crate::client::*;
use crate::server_models::*;
use http::Uri;
use log::{error, info, warn};
use std::env;
use std::fmt::Debug;
use std::fmt::Display;
use std::future::Future;
use std::hash::Hash;
use std::{collections::HashMap, error::Error, time::Duration};

const BASE_URL_ENV: &str = "INTEGRESQL_BASE_URL";
const DEFAULT_TIMEOUT_ENV: &str = "INTEGRESQL_TIMEOUT_SECONDS";
const DB_HOST_OVERRIDE_ENV: &str = "INTEGRESQL_DB_HOST";
const DB_PORT_OVERRIDE_ENV: &str = "INTEGRESQL_DB_PORT";
pub(crate) const DEFAULT_BASE_URL: &str = "http://integresql:5000/api";

/// The result type for template database initialization functions.
pub type InitializeResult = Result<(), Box<dyn Error>>;

/// A template database setup function that should be called synchronously.
///
/// When a new template database is created, this function will be called
/// with Postgres credentials for the template database. The function should connect
/// to the database, perform any setup needed for your tests (i.e. creating tables,
/// inserting data, setting up role and permissions, etc), then close the database
/// connection and return `Ok(())`. If the setup fails, it should return an error-
/// this will trigger IntegreSQL to discard the template database.
///
/// It is important that this function close the database connection before
/// returning- if any connections are left open, attempts to create test
/// databases from this template will fail.
pub trait TemplateInitializer {
    fn setup(self, config: ConnectionSettings) -> InitializeResult;
}

// Blanket implementation for any function/closure that matches the signature.
// We use `FnOnce` because the `setup` method consumes `self`.
impl<F> TemplateInitializer for F
where
    F: FnOnce(ConnectionSettings) -> InitializeResult,
{
    fn setup(self, config: ConnectionSettings) -> Result<(), Box<dyn Error>> {
        // `self` in this context is the function/closure itself, so we just call it.
        self(config)
    }
}

/// A template database setup function that should be called asynchronously.
///
/// When a new template database is created, this function will be called
/// with Postgres credentials for the template database. The function should connect
/// to the database, perform any setup needed for your tests (i.e. creating tables,
/// inserting data, setting up roles and permissions, etc), then close the database
/// connection and return `Ok(())`. If the setup fails, it should return an error-
/// this will trigger IntegreSQL to discard the template database.
///
/// It is important that this function close the database connection before
/// returning- if any connections are left open, attempts to create test
/// databases from this template will fail.
pub trait AsyncTemplateInitializer {
    fn setup(self, config: ConnectionSettings) -> impl Future<Output = InitializeResult> + Send;
}

// Blanket implementation that allows any async function with signature
// `async fn(DatabaseConfig) -> Result<(), Box<dyn Error>>` to be used as
// an AsyncTemplateInitializer.
impl<F, Fut> AsyncTemplateInitializer for F
where
    F: FnOnce(ConnectionSettings) -> Fut + Send,
    Fut: Future<Output = InitializeResult> + Send,
{
    fn setup(self, config: ConnectionSettings) -> impl Future<Output = InitializeResult> + Send {
        self(config)
    }
}

/// A manager for IntegreSQL template databases- the main entry point for this crate.
///
/// DbManager allows users to create and delete IntegreSQL template databases,
/// and to create suppliers that can be used to obtain test databases based on
/// those templates.
#[derive(Clone, Debug)]
pub struct DbManager {
    client: Client,
    host_override: Option<String>,
    port_override: Option<u16>,
}

#[derive(Clone)]
pub struct TemplateDb {
    client: Client,
    host_override: Option<String>,
    port_override: Option<u16>,
    template_hash: TemplateHash,
}

impl std::fmt::Debug for TemplateDb {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TemplateDb")
            .field("template_hash", &self.template_hash)
            .finish()
    }
}

impl Display for TemplateDb {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "TemplateDb({})", self.template_hash)
    }
}

impl TemplateDb {
    /// Returns a `TestDb` instance for the template database associated with this supplier.
    ///
    /// This method will create a new test database from the template, or return an existing
    /// one if it has already been created.
    pub fn get_writable_test_db(&self) -> Result<ConnectionSettings, IntegresqlError> {
        let response = self.client.get_test_db(self.template_hash)?;
        Ok(self.make_test_db(response, false))
    }

    /// Returns a `TestDb` instance for the template database associated with this supplier.
    ///
    /// This method will create a new test database from the template, or return an existing
    /// one if it has already been created.
    pub fn get_readonly_test_db(&self) -> Result<ConnectionSettings, IntegresqlError> {
        let response = self.client.get_test_db(self.template_hash)?;
        Ok(self.make_test_db(response, true))
    }

    /// Returns the hash-based identifier for this template database.
    ///
    /// The returned value is a 16 character hexadecimal string based on hashing
    /// the template's unique key.
    pub fn get_template_id(&self) -> String {
        self.template_hash.to_string()
    }

    fn make_test_db(&self, response: GetTestDbResponse, reuse: bool) -> ConnectionSettings {
        let drop_action = if reuse {
            DropAction::Unlock(self.client.clone(), response.id)
        } else {
            DropAction::Recreate(self.client.clone(), response.id)
        };
        ConnectionSettings {
            host: response.database.config.host,
            port: response.database.config.port,
            username: response.database.config.username,
            password: response.database.config.password,
            database: response.database.config.database,
            additional_params: response.database.config.additional_params,
            template_hash: response.database.template_hash,
            id: Some(response.id),
            drop_action: Some(drop_action),
        }
        .apply_overrides(self.host_override.as_deref(), self.port_override)
    }
}

/// Returns a DbManager configured from environment variables.
///
/// If any environment variable is set to an invalid value, this function
/// will panic.
impl Default for DbManager {
    fn default() -> Self {
        DbManager::from_env()
    }
}

impl DbManager {
    /// Create a new `DbManager` configured from environment variables.
    ///
    /// The IntegreSQL base URL and request timeout are read from the
    /// `INTEGRESQL_BASE_URL` and `INTEGRESQL_TIMEOUT_SECONDS` environment
    /// variables, respectively. If `INTEGRESQL_BASE_URL` is not set, it
    /// defaults to `http://integresql:5000/api`. If `INTEGRESQL_TIMEOUT` is not
    /// set, it defaults to `6`.
    ///
    /// If any environment variable is set to an invalid value, this function
    /// will panic. To create a `DbManager` without risk of panicing,
    /// use `DbManager::new` instead.
    pub fn from_env() -> Self {
        let base_url = env::var(BASE_URL_ENV)
            .unwrap_or_else(|_| DEFAULT_BASE_URL.to_string())
            .parse::<Uri>()
            .expect("Invalid URI in environment variable INTEGRESQL_CLIENT_BASE_URL");
        let timeout = match env::var(DEFAULT_TIMEOUT_ENV) {
            Ok(val) => val
                .parse::<u64>()
                .map(Duration::from_secs)
                .unwrap_or(DEFAULT_TIMEOUT),
            Err(_) => DEFAULT_TIMEOUT,
        };
        let host_override: Option<String> = env::var(DB_HOST_OVERRIDE_ENV).ok();
        let port_override: Option<u16> = env::var(DB_PORT_OVERRIDE_ENV)
            .map(|val| {
                val.parse::<u16>()
                    .expect("Invalid port number in INTEGRESQL_DB_PORT")
            })
            .ok();
        info!(
            "Configured Integresql client from environment: base_url={} timeout_seconds={}",
            base_url,
            timeout.as_secs()
        );
        let client = Client::new(base_url, timeout);
        DbManager {
            client,
            host_override,
            port_override,
        }
    }

    pub fn new(base_uri: Uri, timeout: Duration) -> Self {
        let client = Client::new(base_uri, timeout);
        DbManager {
            client,
            host_override: None,
            port_override: None,
        }
    }

    /// Returns a new `DbManager` with the provided override settings applied.
    ///
    /// This method consumes the current manager, replacing any existing overrides
    /// with the supplied values.
    pub fn with_overrides(
        mut self,
        host_override: Option<String>,
        port_override: Option<u16>,
    ) -> Self {
        self.host_override = host_override;
        self.port_override = port_override;
        self
    }

    /// Drops all test databases, and deregisters all template databases.
    ///
    /// This method drops all existing test databases from Postgres and
    /// deregisters all template databases from the IntegreSQL server. It does
    /// not *drop* the template databases themselves, though. As such, it will
    /// disallow creation of new test databases from any current template
    /// databases, but will when intitializing new templates, will reuse an
    /// existing template database if one already exists with the same template
    /// hash.
    ///
    /// To deregister *and* drop a template database, use `discard_template`.
    pub fn clear_db_tracking(&self) -> Result<(), IntegresqlError> {
        self.client.clear_db_tracking()
    }

    /// Deregisters and deletes the template database identified by the given
    /// template key, if one exists.
    ///
    /// The template_key should be the same value that was passed to
    /// `get_db_supplier_async` or `get_db_supplier_sync` when creating the
    /// template database.
    ///
    /// This prevents the creation of new test databases from this template, but
    /// does not delete any existing test databases that have already been
    /// created from this template.
    pub fn discard_template(&self, template_key: impl Hash) -> Result<(), IntegresqlError> {
        let template_hash = TemplateHash::from_hash(template_key);
        self.client.discard_template(template_hash)
    }

    /// Creates a new `DbSupplier` for the specified template key, using the
    /// provided synchronous initializer function to set up the template database.
    ///
    /// This method will create a new template database if one with the given
    /// key does not already exist, and will call the provided initializer
    /// to setup the database contents. If a template database with the same key
    /// already exists, the initializer will not be called and the existing
    /// template database will be used.
    ///
    /// The `template_key` should be a hashable value that uniquely identifies
    /// the initialization logic that will be used to setup the template
    /// database. For example, a `String` or `Vec<String>` containing the SQL
    /// scripts that will be used to setup the database.
    ///
    /// The `initializer` should return `Ok(())` if the template database was
    /// successfully set up, or an error if the setup failed.  A blanket
    /// implementation allows any non-async function or closure with signature
    /// `Result<(), Box<dyn Error + 'static>>` to be used as the initializer.
    pub async fn get_template_db_async(
        &self,
        template_key: impl Hash,
        initializer: impl AsyncTemplateInitializer,
    ) -> Result<TemplateDb, IntegresqlError> {
        let template_hash = TemplateHash::from_hash(template_key);
        let supplier = TemplateDb {
            client: self.client.clone(),
            host_override: self.host_override.clone(),
            port_override: self.port_override,
            template_hash,
        };

        let template_settings = match self.client.initialize_template(supplier.template_hash) {
            InitializeTemplateResult::Success(db) => db,
            InitializeTemplateResult::TemplateAlreadyInitialized => return Ok(supplier),
            InitializeTemplateResult::Err(e) => return Err(e),
        };

        let connection_settings = ConnectionSettings {
            host: template_settings.config.host,
            port: template_settings.config.port,
            username: template_settings.config.username,
            password: template_settings.config.password,
            database: template_settings.config.database,
            additional_params: template_settings.config.additional_params,
            template_hash: supplier.template_hash,
            id: None,          // This is a template, so no ID is assigned
            drop_action: None, // No drop action for templates
        };
        let connection_settings =
            connection_settings.apply_overrides(self.host_override.as_deref(), self.port_override);

        let result = initializer.setup(connection_settings).await;
        Self::handle_template_setup_result(supplier, result)
    }

    /// Creates a new `DbSupplier` for the specified template key, using the
    /// provided synchronous initializer function to set up the template database.
    ///
    /// This method will create a new template database if one with the given
    /// key does not already exist, and will call the provided initializer
    /// to setup the database contents. If a template database with the same key
    /// already exists, the initializer will not be called and the existing
    /// template database will be used.
    ///
    /// The `template_key` should be a hashable value that uniquely identifies
    /// the initialization logic that will be used to setup the template
    /// database. For example, a `String` or `Vec<String>` containing the SQL
    /// scripts that will be used to setup the database.
    ///
    /// The `initializer` should return `Ok(())` if the template database was
    /// successfully set up, or an error if the setup failed.  A blanket
    /// implementation allows any async function or closure with signature
    /// `Result<(), Box<dyn Error + 'static>>` to be used as the initializer.
    pub fn get_template_db_sync(
        &self,
        template_key: impl Hash,
        initializer: impl TemplateInitializer,
    ) -> Result<TemplateDb, IntegresqlError> {
        let template_hash = TemplateHash::from_hash(template_key);
        let supplier = TemplateDb {
            client: self.client.clone(),
            host_override: self.host_override.clone(),
            port_override: self.port_override,
            template_hash,
        };

        let template_settings = match self.client.initialize_template(supplier.template_hash) {
            InitializeTemplateResult::Success(db) => db,
            InitializeTemplateResult::TemplateAlreadyInitialized => return Ok(supplier),
            InitializeTemplateResult::Err(e) => return Err(e),
        };

        let connection_settings = ConnectionSettings {
            host: template_settings.config.host,
            port: template_settings.config.port,
            username: template_settings.config.username,
            password: template_settings.config.password,
            database: template_settings.config.database,
            additional_params: template_settings.config.additional_params,
            template_hash: supplier.template_hash,
            id: None,          // This is a template, so no ID is assigned
            drop_action: None, // No drop action for templates
        };
        let connection_settings =
            connection_settings.apply_overrides(self.host_override.as_deref(), self.port_override);

        let result = initializer.setup(connection_settings);
        Self::handle_template_setup_result(supplier, result)
    }

    fn handle_template_setup_result(
        supplier: TemplateDb,
        result: Result<(), Box<dyn Error>>,
    ) -> Result<TemplateDb, IntegresqlError> {
        match result {
            // If the setup was successful, finalize the template and return the supplier
            Ok(_) => {
                info!(
                    "Template setup completed successfully template_hash={}",
                    supplier.template_hash
                );
                supplier.client.finalize_template(supplier.template_hash)?;
                Ok(supplier)
            }
            // If it failed, discard the template and return an error
            Err(e) => {
                warn!(
                    "Template setup failed for hash template_hash={} error={}",
                    supplier.template_hash, e
                );
                supplier.client.discard_template(supplier.template_hash)?;
                Err(IntegresqlError::SetupError(e.to_string()))
            }
        }
    }
}

#[derive(Debug, Clone)]
enum DropAction {
    Recreate(Client, i32),
    Unlock(Client, i32),
}

/// Credentials to connect to a database managed by IntegreSQL- either a
/// template database or a test database.
///
/// If this is a writable test database, it will be deleted and recreated when
/// it goes out of scope. If this is a read-only test database, it will be
/// unlocked and made available to other tests when it goes out of scope.  If
/// this is a template database, no action will be taken when it goes out of
/// scope.
///
/// The `id` field can be used to determine if this is a test database or a
/// template database (if it's `None`, this is a template database).
#[derive(Debug, Clone)]
pub struct ConnectionSettings {
    /// The Postgres server hostname
    pub host: String,

    /// The Postgres server port
    pub port: u16,

    /// The username to connect as
    pub username: String,

    /// The password for this username
    pub password: String,

    /// The name of the test database
    pub database: String,

    /// Additional parameters to include in the connection string
    pub additional_params: Option<HashMap<String, String>>,

    /// If this is a test database, its unique test DB id. If it is a template
    /// database, `None`
    pub id: Option<i32>,

    /// The hash of this database's associated template database.
    ///
    /// If this is a template database, this will be the template hash of the
    /// database itself.  If this is a test database, this will be the template
    /// hash of the DB's source template DB.
    template_hash: TemplateHash,

    drop_action: Option<DropAction>,
}

/// # Warning
///
/// The `Drop` implementation for `TestDb` performs an HTTP request to the IntegreSQL server
/// to cleanup the test database when the `TestDb` instance goes out of scope. Any errors
/// encountered during this cleanup will be logged but not propagated.
impl Drop for ConnectionSettings {
    fn drop(&mut self) {
        if let Some(drop_action) = self.drop_action.take() {
            match drop_action {
                DropAction::Recreate(client, id) => {
                    client.recreate_test_db(self.template_hash, id)
                        .unwrap_or_else(|e| {
                            error!("Failed to recreate test database template_hash={} test_db_id={} error={}", self.template_hash, id, e);
                        });
                }
                DropAction::Unlock(client, id) => {
                    client.unlock_test_db(self.template_hash, id)
                        .unwrap_or_else(|e| {
                            error!("Failed to unlock test database template_hash={} test_db_id={} error={}", self.template_hash, id, e);
                        });
                }
            }
        }
    }
}

impl ConnectionSettings {
    /// Returns a libpq-compatible Postgres connection string based on this configuration.
    ///
    /// The returned string will be in this format: `postgres://username:password@host:port/database?param1=value1`
    pub fn to_libpq_url(&self) -> String {
        {
            let host: &str = &self.host;
            let port = self.port;
            let username: &str = &self.username;
            let password: &str = &self.password;
            let database: &str = &self.database;
            let additional_params = &self.additional_params;
            let additional_params = additional_params
                .as_ref()
                .map(|params| {
                    let mut params = params
                        .iter()
                        .filter(|(_, v)| !v.is_empty())
                        .map(|(k, v)| format!("{}={}", k, v))
                        .collect::<Vec<_>>();
                    params.sort();
                    params.join("&")
                })
                .unwrap_or_default();
            format!(
                "postgres://{}:{}@{}:{}/{}?{}",
                username, password, host, port, database, additional_params
            )
        }
    }

    /// Adds an additional connection parameter to the connection settings.
    ///
    /// If a parameter with the same key already exists, it will be overwritten.
    pub fn with_additional_param(&mut self, key: impl ToString, value: impl ToString) -> &mut Self {
        match self.additional_params {
            Some(ref mut params) => {
                params.insert(key.to_string(), value.to_string());
            }
            None => {
                let mut params = HashMap::new();
                params.insert(key.to_string(), value.to_string());
                self.additional_params = Some(params);
            }
        }
        self
    }

    /// The hash of this database's associated template database.
    ///
    /// If this is a template database, this will be the template hash of the
    /// database itself.  If this is a test database, this will be the template
    /// hash of the DB's source template DB.
    pub fn template_hash(&self) -> String {
        self.template_hash.to_string()
    }

    /// Returns a new `ConnectionSettings` with the provided host/port overrides applied.
    ///
    /// When an override is `None`, the original value is preserved.
    pub fn apply_overrides(
        mut self,
        host_override: Option<&str>,
        port_override: Option<u16>,
    ) -> Self {
        if let Some(host) = host_override {
            self.host = host.to_string();
        }
        if let Some(port) = port_override {
            self.port = port;
        }
        self
    }
}

mod test {
    use super::*;

    #[test]
    fn test_connection_settings_with_additional_param() {
        let mut settings = ConnectionSettings {
            host: "localhost".to_string(),
            port: 5432,
            username: "test_user".to_string(),
            password: "test_pass".to_string(),
            database: "test_db".to_string(),
            additional_params: None,
            id: None,
            template_hash: TemplateHash::from_hash("test_template"),
            drop_action: None,
        };
        let url = settings.to_libpq_url();
        assert_eq!(
            url,
            "postgres://test_user:test_pass@localhost:5432/test_db?"
        );
        let url = settings
            .with_additional_param("key1", "value1")
            .with_additional_param("key2", "value2")
            .to_libpq_url();
        assert_eq!(
            url,
            "postgres://test_user:test_pass@localhost:5432/test_db?key1=value1&key2=value2"
        );

        let url = settings.with_additional_param("key1", "").to_libpq_url();
        assert_eq!(
            url,
            "postgres://test_user:test_pass@localhost:5432/test_db?key2=value2"
        );
    }

    #[test]
    fn test_connection_settings_apply_overrides_with_values() {
        let settings = ConnectionSettings {
            host: "integresql-internal".to_string(),
            port: 5432,
            username: "test_user".to_string(),
            password: "test_pass".to_string(),
            database: "test_db".to_string(),
            additional_params: None,
            id: Some(1),
            template_hash: TemplateHash::from_hash("abc"),
            drop_action: None,
        };

        let overridden = settings.apply_overrides(Some("external-host"), Some(6543));

        assert_eq!(overridden.host, "external-host");
        assert_eq!(overridden.port, 6543);
        assert_eq!(overridden.username, "test_user");
        assert_eq!(overridden.database, "test_db");
    }

    #[test]
    fn test_connection_settings_apply_overrides_no_values() {
        let settings = ConnectionSettings {
            host: "integresql-internal".to_string(),
            port: 5432,
            username: "test_user".to_string(),
            password: "test_pass".to_string(),
            database: "test_db".to_string(),
            additional_params: None,
            id: Some(1),
            template_hash: TemplateHash::from_hash("abc"),
            drop_action: None,
        };

        let overridden = settings.apply_overrides(None, None);

        assert_eq!(overridden.host, "integresql-internal");
        assert_eq!(overridden.port, 5432);
    }
}