picodata-driver 3.0.0

use std::{net::IpAddr, time::Duration};

use pg::config::Host;

use crate::instance::{Addr, Instance};

/// Default max connections per instance pool.
const DEFAULT_POOL_SIZE: usize = 10;
/// Default interval between topology discovery refreshes.
const DEFAULT_DISCOVERY_INTERVAL: Duration = Duration::from_millis(500);

/// Picodata config
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Config {
    /// Config to initialize connection.
    pub postgres: pg::Config,
    /// Max size of connection pool for one instance.
    pub pool_max_size: usize,
    /// Interval for making topology requests.
    pub discorery_interval: Duration,
    /// Path to cert file.
    pub ssl_cert_file: Option<String>,
    /// Path to key file.
    pub ssl_key_file: Option<String>,
    /// Path to ca file.
    pub ssl_ca_file: Option<String>,
    /// Tiers that are not used to create connections to.
    pub excluded_tiers: Option<Vec<String>>,
}

impl Config {
    /// Get hosts of instance.
    ///
    /// # Returns
    ///
    /// Hosts if it was set.
    #[must_use]
    pub fn hosts(&self) -> &[Host] {
        self.postgres.get_hosts()
    }
    /// Get ports of instance.
    ///
    /// # Returns
    ///
    /// Ports if it was set.
    #[must_use]
    pub fn ports(&self) -> &[u16] {
        self.postgres.get_ports()
    }

    /// Get a hostaddrs of instance.
    ///
    /// # Returns
    ///
    /// A slice of hostaddrs.
    #[must_use]
    pub fn hostaddrs(&self) -> &[IpAddr] {
        self.postgres.get_hostaddrs()
    }

    /// Clones the config, pointed at `inst`.
    pub(crate) fn clone_for(&self, inst: &Instance) -> Self {
        let mut cfg = self.clone_pg_reset();
        match &inst.addr {
            Addr::Ip(addr) => {
                cfg.postgres.hostaddr(addr.ip());
                cfg.postgres.port(addr.port());
            }
            Addr::Name { host, port } => {
                cfg.postgres.host(&**host);
                cfg.postgres.port(*port);
            }
        }
        cfg
    }

    /// Clones the config with resetting addr info of instance.
    ///
    /// # Returns
    ///
    /// An instance of `Config`.
    #[must_use]
    pub fn clone_pg_reset(&self) -> Self {
        Self {
            postgres: clone_pg_config(&self.postgres),
            pool_max_size: self.pool_max_size,
            discorery_interval: self.discorery_interval,
            ssl_cert_file: self.ssl_cert_file.clone(),
            ssl_key_file: self.ssl_key_file.clone(),
            ssl_ca_file: self.ssl_ca_file.clone(),
            excluded_tiers: self.excluded_tiers.clone(),
        }
    }
}

impl From<pg::Config> for Config {
    fn from(postgres: pg::Config) -> Self {
        Self {
            postgres,
            pool_max_size: DEFAULT_POOL_SIZE,
            discorery_interval: DEFAULT_DISCOVERY_INTERVAL,
            ssl_cert_file: None,
            ssl_key_file: None,
            ssl_ca_file: None,
            excluded_tiers: None,
        }
    }
}

/// Clones a `pg::Config` preserving auth and TLS settings but resetting host/port.
#[must_use]
pub fn clone_pg_config(cfg: &pg::Config) -> pg::Config {
    let mut result = pg::Config::new();

    if let Some(v) = cfg.get_application_name() {
        result.application_name(v);
    }
    if let Some(v) = cfg.get_user() {
        result.user(v);
    }
    if let Some(v) = cfg.get_password() {
        result.password(v);
    }
    if let Some(v) = cfg.get_dbname() {
        result.dbname(v);
    }
    if let Some(v) = cfg.get_options() {
        result.options(v);
    }
    if let Some(v) = cfg.get_connect_timeout() {
        result.connect_timeout(*v);
    }
    if let Some(v) = cfg.get_tcp_user_timeout() {
        result.tcp_user_timeout(*v);
    }

    if let Some(v) = cfg.get_keepalives_interval() {
        result.keepalives_interval(v);
    }
    if let Some(v) = cfg.get_keepalives_retries() {
        result.keepalives_retries(v);
    }

    result.ssl_mode(cfg.get_ssl_mode());
    result.ssl_negotiation(cfg.get_ssl_negotiation());
    result.keepalives(cfg.get_keepalives());
    result.keepalives_idle(cfg.get_keepalives_idle());
    result.target_session_attrs(cfg.get_target_session_attrs());
    result.channel_binding(cfg.get_channel_binding());
    result.load_balance_hosts(cfg.get_load_balance_hosts());

    result
}

#[cfg(test)]
mod tests {
    use super::*;

    /// An IP goes in as `hostaddr`, a name as `host`.
    #[test]
    fn clone_for_uses_hostaddr_for_an_ip_and_host_for_a_name() {
        let base = Config::from(pg::Config::new());
        let inst =
            |addr| Instance::new_for_test("i1", addr, "default", uuid::Uuid::nil(), true, 3000);

        let cfg = base.clone_for(&inst("127.0.0.1:5433"));
        assert_eq!(cfg.hostaddrs(), [IpAddr::from([127, 0, 0, 1])]);
        assert!(cfg.hosts().is_empty(), "an IP must not become a name");
        assert_eq!(cfg.ports(), [5433]);

        let cfg = base.clone_for(&inst("[::1]:5433"));
        assert_eq!(cfg.hostaddrs(), [IpAddr::from([0, 0, 0, 0, 0, 0, 0, 1])]);

        let cfg = base.clone_for(&inst("picodata-1:5432"));
        assert!(cfg.hostaddrs().is_empty(), "a name has no address yet");
        assert_eq!(cfg.hosts().len(), 1);
        assert_eq!(cfg.ports(), [5432]);
    }
}