Skip to main content

picodata_driver/
config.rs

1use std::{net::IpAddr, time::Duration};
2
3use pg::config::Host;
4
5use crate::instance::{Addr, Instance};
6
7/// Default max connections per instance pool.
8const DEFAULT_POOL_SIZE: usize = 10;
9/// Default interval between topology discovery refreshes.
10const DEFAULT_DISCOVERY_INTERVAL: Duration = Duration::from_millis(500);
11
12/// Picodata config
13#[derive(Clone, Debug, PartialEq, Eq)]
14pub struct Config {
15    /// Config to initialize connection.
16    pub postgres: pg::Config,
17    /// Max size of connection pool for one instance.
18    pub pool_max_size: usize,
19    /// Interval for making topology requests.
20    pub discorery_interval: Duration,
21    /// Path to cert file.
22    pub ssl_cert_file: Option<String>,
23    /// Path to key file.
24    pub ssl_key_file: Option<String>,
25    /// Path to ca file.
26    pub ssl_ca_file: Option<String>,
27    /// Tiers that are not used to create connections to.
28    pub excluded_tiers: Option<Vec<String>>,
29}
30
31impl Config {
32    /// Get hosts of instance.
33    ///
34    /// # Returns
35    ///
36    /// Hosts if it was set.
37    #[must_use]
38    pub fn hosts(&self) -> &[Host] {
39        self.postgres.get_hosts()
40    }
41    /// Get ports of instance.
42    ///
43    /// # Returns
44    ///
45    /// Ports if it was set.
46    #[must_use]
47    pub fn ports(&self) -> &[u16] {
48        self.postgres.get_ports()
49    }
50
51    /// Get a hostaddrs of instance.
52    ///
53    /// # Returns
54    ///
55    /// A slice of hostaddrs.
56    #[must_use]
57    pub fn hostaddrs(&self) -> &[IpAddr] {
58        self.postgres.get_hostaddrs()
59    }
60
61    /// Clones the config, pointed at `inst`.
62    pub(crate) fn clone_for(&self, inst: &Instance) -> Self {
63        let mut cfg = self.clone_pg_reset();
64        match &inst.addr {
65            Addr::Ip(addr) => {
66                cfg.postgres.hostaddr(addr.ip());
67                cfg.postgres.port(addr.port());
68            }
69            Addr::Name { host, port } => {
70                cfg.postgres.host(&**host);
71                cfg.postgres.port(*port);
72            }
73        }
74        cfg
75    }
76
77    /// Clones the config with resetting addr info of instance.
78    ///
79    /// # Returns
80    ///
81    /// An instance of `Config`.
82    #[must_use]
83    pub fn clone_pg_reset(&self) -> Self {
84        Self {
85            postgres: clone_pg_config(&self.postgres),
86            pool_max_size: self.pool_max_size,
87            discorery_interval: self.discorery_interval,
88            ssl_cert_file: self.ssl_cert_file.clone(),
89            ssl_key_file: self.ssl_key_file.clone(),
90            ssl_ca_file: self.ssl_ca_file.clone(),
91            excluded_tiers: self.excluded_tiers.clone(),
92        }
93    }
94}
95
96impl From<pg::Config> for Config {
97    fn from(postgres: pg::Config) -> Self {
98        Self {
99            postgres,
100            pool_max_size: DEFAULT_POOL_SIZE,
101            discorery_interval: DEFAULT_DISCOVERY_INTERVAL,
102            ssl_cert_file: None,
103            ssl_key_file: None,
104            ssl_ca_file: None,
105            excluded_tiers: None,
106        }
107    }
108}
109
110/// Clones a `pg::Config` preserving auth and TLS settings but resetting host/port.
111#[must_use]
112pub fn clone_pg_config(cfg: &pg::Config) -> pg::Config {
113    let mut result = pg::Config::new();
114
115    if let Some(v) = cfg.get_application_name() {
116        result.application_name(v);
117    }
118    if let Some(v) = cfg.get_user() {
119        result.user(v);
120    }
121    if let Some(v) = cfg.get_password() {
122        result.password(v);
123    }
124    if let Some(v) = cfg.get_dbname() {
125        result.dbname(v);
126    }
127    if let Some(v) = cfg.get_options() {
128        result.options(v);
129    }
130    if let Some(v) = cfg.get_connect_timeout() {
131        result.connect_timeout(*v);
132    }
133    if let Some(v) = cfg.get_tcp_user_timeout() {
134        result.tcp_user_timeout(*v);
135    }
136
137    if let Some(v) = cfg.get_keepalives_interval() {
138        result.keepalives_interval(v);
139    }
140    if let Some(v) = cfg.get_keepalives_retries() {
141        result.keepalives_retries(v);
142    }
143
144    result.ssl_mode(cfg.get_ssl_mode());
145    result.ssl_negotiation(cfg.get_ssl_negotiation());
146    result.keepalives(cfg.get_keepalives());
147    result.keepalives_idle(cfg.get_keepalives_idle());
148    result.target_session_attrs(cfg.get_target_session_attrs());
149    result.channel_binding(cfg.get_channel_binding());
150    result.load_balance_hosts(cfg.get_load_balance_hosts());
151
152    result
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158
159    /// An IP goes in as `hostaddr`, a name as `host`.
160    #[test]
161    fn clone_for_uses_hostaddr_for_an_ip_and_host_for_a_name() {
162        let base = Config::from(pg::Config::new());
163        let inst =
164            |addr| Instance::new_for_test("i1", addr, "default", uuid::Uuid::nil(), true, 3000);
165
166        let cfg = base.clone_for(&inst("127.0.0.1:5433"));
167        assert_eq!(cfg.hostaddrs(), [IpAddr::from([127, 0, 0, 1])]);
168        assert!(cfg.hosts().is_empty(), "an IP must not become a name");
169        assert_eq!(cfg.ports(), [5433]);
170
171        let cfg = base.clone_for(&inst("[::1]:5433"));
172        assert_eq!(cfg.hostaddrs(), [IpAddr::from([0, 0, 0, 0, 0, 0, 0, 1])]);
173
174        let cfg = base.clone_for(&inst("picodata-1:5432"));
175        assert!(cfg.hostaddrs().is_empty(), "a name has no address yet");
176        assert_eq!(cfg.hosts().len(), 1);
177        assert_eq!(cfg.ports(), [5432]);
178    }
179}