Added more parameters to configure connection

Signed-off-by: chandr-andr (Kiselev Aleksandr) <chandr@chandr.net>

Author: chandr-andr

src/driver/common_options.rs

@@ -1,5 +1,6 @@
 use deadpool_postgres::RecyclingMethod;
 use pyo3::pyclass;
+use tokio_postgres::config::LoadBalanceHosts;
 
 #[pyclass]
 #[derive(Clone, Copy)]
@@ -19,3 +20,22 @@ impl ConnRecyclingMethod {
         }
     }
 }
+
+#[pyclass]
+#[derive(Clone, Copy)]
+pub enum ConnLoadBalanceHosts {
+    /// Make connection attempts to hosts in the order provided.
+    Disable,
+    /// Make connection attempts to hosts in a random order.
+    Random,
+}
+
+impl ConnLoadBalanceHosts {
+    #[must_use]
+    pub fn to_internal(&self) -> LoadBalanceHosts {
+        match self {
+            ConnLoadBalanceHosts::Disable => LoadBalanceHosts::Disable,
+            ConnLoadBalanceHosts::Random => LoadBalanceHosts::Random,
+        }
+    }
+}

src/driver/connection_pool.rs

@@ -1,7 +1,7 @@
 use crate::runtime::tokio_runtime;
 use deadpool_postgres::{Manager, ManagerConfig, Object, Pool, RecyclingMethod};
 use pyo3::{pyclass, pyfunction, pymethods, PyAny};
-use std::{str::FromStr, vec};
+use std::{time::Duration, vec};
 use tokio_postgres::{NoTls, Row};
 
 use crate::{
@@ -10,7 +10,11 @@ use crate::{
     value_converter::{convert_parameters, PythonDTO, QueryParameter},
 };
 
-use super::{common_options::ConnRecyclingMethod, connection::Connection};
+use super::{
+    common_options::{ConnLoadBalanceHosts, ConnRecyclingMethod},
+    connection::Connection,
+    utils::build_connection_config,
+};
 
 /// Make new connection pool.
 ///
@@ -25,6 +29,16 @@ pub fn connect(
     host: Option<String>,
     port: Option<u16>,
     db_name: Option<String>,
+    options: Option<String>,
+    application_name: Option<String>,
+    connect_timeout: Option<Duration>,
+    tcp_user_timeout: Option<Duration>,
+    keepalives: Option<bool>,
+    keepalives_idle: Option<Duration>,
+    keepalives_interval: Option<Duration>,
+    keepalives_retries: Option<u32>,
+    load_balance_hosts: Option<ConnLoadBalanceHosts>,
+
     max_db_pool_size: Option<usize>,
     conn_recycling_method: Option<ConnRecyclingMethod>,
 ) -> RustPSQLDriverPyResult<ConnectionPool> {
@@ -36,27 +50,23 @@ pub fn connect(
         }
     }
 
-    let mut pg_config: tokio_postgres::Config;
-    if let Some(dsn_string) = dsn {
-        pg_config = tokio_postgres::Config::from_str(&dsn_string)?;
-    } else {
-        pg_config = tokio_postgres::Config::new();
-        if let (Some(password), Some(username)) = (password, username) {
-            pg_config.password(&password);
-            pg_config.user(&username);
-        }
-        if let Some(host) = host {
-            pg_config.host(&host);
-        }
-
-        if let Some(port) = port {
-            pg_config.port(port);
-        }
-
-        if let Some(db_name) = db_name {
-            pg_config.dbname(&db_name);
-        }
-    }
+    let pg_config = build_connection_config(
+        dsn,
+        username,
+        password,
+        host,
+        port,
+        db_name,
+        options,
+        application_name,
+        connect_timeout,
+        tcp_user_timeout,
+        keepalives,
+        keepalives_idle,
+        keepalives_interval,
+        keepalives_retries,
+        load_balance_hosts,
+    )?;
 
     let mgr_config: ManagerConfig;
     if let Some(conn_recycling_method) = conn_recycling_method {
@@ -98,6 +108,15 @@ impl ConnectionPool {
         host: Option<String>,
         port: Option<u16>,
         db_name: Option<String>,
+        options: Option<String>,
+        application_name: Option<String>,
+        connect_timeout: Option<Duration>,
+        tcp_user_timeout: Option<Duration>,
+        keepalives: Option<bool>,
+        keepalives_idle: Option<Duration>,
+        keepalives_interval: Option<Duration>,
+        keepalives_retries: Option<u32>,
+        load_balance_hosts: Option<ConnLoadBalanceHosts>,
         max_db_pool_size: Option<usize>,
         conn_recycling_method: Option<ConnRecyclingMethod>,
     ) -> RustPSQLDriverPyResult<Self> {
@@ -108,6 +127,15 @@ impl ConnectionPool {
             host,
             port,
             db_name,
+            options,
+            application_name,
+            connect_timeout,
+            tcp_user_timeout,
+            keepalives,
+            keepalives_idle,
+            keepalives_interval,
+            keepalives_retries,
+            load_balance_hosts,
             max_db_pool_size,
             conn_recycling_method,
         )

src/driver/mod.rs

@@ -4,3 +4,4 @@ pub mod connection_pool;
 pub mod cursor;
 pub mod transaction;
 pub mod transaction_options;
+pub mod utils;

src/driver/utils.rs

@@ -0,0 +1,90 @@
+use std::{str::FromStr, time::Duration};
+
+use crate::exceptions::rust_errors::RustPSQLDriverPyResult;
+
+use super::common_options::ConnLoadBalanceHosts;
+
+/// Create new config.
+///
+/// # Errors
+/// May return Err Result if cannot build new config.
+#[allow(clippy::too_many_arguments)]
+pub fn build_connection_config(
+    dsn: Option<String>,
+    username: Option<String>,
+    password: Option<String>,
+    host: Option<String>,
+    port: Option<u16>,
+    db_name: Option<String>,
+    options: Option<String>,
+    application_name: Option<String>,
+    connect_timeout: Option<Duration>,
+    tcp_user_timeout: Option<Duration>,
+    keepalives: Option<bool>,
+    keepalives_idle: Option<Duration>,
+    keepalives_interval: Option<Duration>,
+    keepalives_retries: Option<u32>,
+    load_balance_hosts: Option<ConnLoadBalanceHosts>,
+) -> RustPSQLDriverPyResult<tokio_postgres::Config> {
+    if let Some(dsn_string) = dsn {
+        return Ok(tokio_postgres::Config::from_str(&dsn_string)?);
+    }
+
+    let mut pg_config = tokio_postgres::Config::new();
+
+    if let (Some(password), Some(username)) = (password, username) {
+        pg_config.password(&password);
+        pg_config.user(&username);
+    }
+    if let Some(host) = host {
+        pg_config.host(&host);
+    }
+
+    if let Some(port) = port {
+        pg_config.port(port);
+    }
+
+    if let Some(db_name) = db_name {
+        pg_config.dbname(&db_name);
+    }
+
+    if let Some(options) = options {
+        pg_config.options(&options);
+    }
+
+    if let Some(application_name) = application_name {
+        pg_config.application_name(&application_name);
+    }
+
+    if let Some(connect_timeout) = connect_timeout {
+        pg_config.connect_timeout(connect_timeout);
+    }
+
+    if let Some(tcp_user_timeout) = tcp_user_timeout {
+        pg_config.tcp_user_timeout(tcp_user_timeout);
+    }
+
+    if let Some(keepalives) = keepalives {
+        if keepalives {
+            pg_config.keepalives(keepalives);
+
+            if let Some(keepalives_idle) = keepalives_idle {
+                pg_config.keepalives_idle(keepalives_idle);
+            }
+
+            if let Some(keepalives_interval) = keepalives_interval {
+                pg_config.keepalives_interval(keepalives_interval);
+            }
+
+            if let Some(keepalives_retries) = keepalives_retries {
+                pg_config.keepalives_retries(keepalives_retries);
+            }
+        }
+    }
+
+    if let Some(load_balance_hosts) = load_balance_hosts {
+        pg_config.load_balance_hosts(load_balance_hosts.to_internal());
+    }
+
+    Ok(pg_config)
+}