From b32a42ba3711bf8af55c5c3221393f86236c1cc7 Mon Sep 17 00:00:00 2001 From: "chandr-andr (Kiselev Aleksandr)" Date: Sun, 28 Apr 2024 19:36:14 +0200 Subject: [PATCH 1/7] Added more parameters to configure connection Signed-off-by: chandr-andr (Kiselev Aleksandr) --- src/driver/common_options.rs | 20 ++++++++ src/driver/connection_pool.rs | 74 +++++++++++++++++++--------- src/driver/mod.rs | 1 + src/driver/utils.rs | 90 +++++++++++++++++++++++++++++++++++ 4 files changed, 162 insertions(+), 23 deletions(-) create mode 100644 src/driver/utils.rs diff --git a/src/driver/common_options.rs b/src/driver/common_options.rs index cd8cb362..ec4c9e8d 100644 --- a/src/driver/common_options.rs +++ b/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, + } + } +} diff --git a/src/driver/connection_pool.rs b/src/driver/connection_pool.rs index fad7cb9e..491d6892 100644 --- a/src/driver/connection_pool.rs +++ b/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, port: Option, db_name: Option, + options: Option, + application_name: Option, + connect_timeout: Option, + tcp_user_timeout: Option, + keepalives: Option, + keepalives_idle: Option, + keepalives_interval: Option, + keepalives_retries: Option, + load_balance_hosts: Option, + max_db_pool_size: Option, conn_recycling_method: Option, ) -> RustPSQLDriverPyResult { @@ -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, port: Option, db_name: Option, + options: Option, + application_name: Option, + connect_timeout: Option, + tcp_user_timeout: Option, + keepalives: Option, + keepalives_idle: Option, + keepalives_interval: Option, + keepalives_retries: Option, + load_balance_hosts: Option, max_db_pool_size: Option, conn_recycling_method: Option, ) -> RustPSQLDriverPyResult { @@ -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, ) diff --git a/src/driver/mod.rs b/src/driver/mod.rs index aec33d5b..1ba0c203 100644 --- a/src/driver/mod.rs +++ b/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; diff --git a/src/driver/utils.rs b/src/driver/utils.rs new file mode 100644 index 00000000..0e3ee0f3 --- /dev/null +++ b/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, + username: Option, + password: Option, + host: Option, + port: Option, + db_name: Option, + options: Option, + application_name: Option, + connect_timeout: Option, + tcp_user_timeout: Option, + keepalives: Option, + keepalives_idle: Option, + keepalives_interval: Option, + keepalives_retries: Option, + load_balance_hosts: Option, +) -> RustPSQLDriverPyResult { + 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) +} From f85a09265454f36c9b7b569842db51e28255bec9 Mon Sep 17 00:00:00 2001 From: "chandr-andr (Kiselev Aleksandr)" Date: Sun, 28 Apr 2024 22:06:36 +0200 Subject: [PATCH 2/7] Added more parameters to configure connection Signed-off-by: chandr-andr (Kiselev Aleksandr) --- python/psqlpy/_internal/__init__.pyi | 180 +++++++++++++++++++++++---- src/driver/common_options.rs | 24 +++- src/driver/connection_pool.rs | 64 +++++++--- src/driver/utils.rs | 116 +++++++++++++---- src/lib.rs | 2 + 5 files changed, 319 insertions(+), 67 deletions(-) diff --git a/python/psqlpy/_internal/__init__.pyi b/python/psqlpy/_internal/__init__.pyi index df1b05d5..3033a864 100644 --- a/python/psqlpy/_internal/__init__.pyi +++ b/python/psqlpy/_internal/__init__.pyi @@ -1,6 +1,6 @@ import types from enum import Enum -from typing import Any, Callable, Optional, TypeVar +from typing import Any, Callable, List, Optional, TypeVar from typing_extensions import Self @@ -104,6 +104,14 @@ class IsolationLevel(Enum): RepeatableRead = 3 Serializable = 4 +class ConnLoadBalanceHosts(Enum): + """Load balancing configuration.""" + + # Make connection attempts to hosts in the order provided. + Disable = (1,) + # Make connection attempts to hosts in a random order. + Random = (2,) + class ReadVariant(Enum): """Class for Read Variant for transaction.""" @@ -869,8 +877,23 @@ class ConnectionPool: username: Optional[str] = None, password: Optional[str] = None, host: Optional[str] = None, + hosts: Optional[List[str]] = None, port: Optional[int] = None, + ports: Optional[List[int]] = None, db_name: Optional[str] = None, + options: Optional[str] = None, + application_name: Optional[str] = None, + connect_timeout_sec: Optional[int] = None, + connect_timeout_nanosec: Optional[int] = None, + tcp_user_timeout_sec: Optional[int] = None, + tcp_user_timeout_nanosec: Optional[int] = None, + keepalives: Optional[bool] = None, + keepalives_idle_sec: Optional[int] = None, + keepalives_idle_nanosec: Optional[int] = None, + keepalives_interval_sec: Optional[int] = None, + keepalives_interval_nanosec: Optional[int] = None, + keepalives_retries: Optional[int] = None, + load_balance_hosts: Optional[ConnLoadBalanceHosts] = None, max_db_pool_size: int = 2, conn_recycling_method: Optional[ConnRecyclingMethod] = None, ) -> None: @@ -879,22 +902,67 @@ class ConnectionPool: It connects to the database and create pool. You cannot set the minimum size for the connection - pool, by default it is 1. + pool, by it is 0. + `ConnectionPool` doesn't create connections on startup. + It makes new connection on demand. - This connection pool can: - - Startup itself with `startup` method - - Execute queries and return `QueryResult` class as a result - - Create new instance of `Transaction` + If you specify `dsn` parameter then `username`, `password`, + `host`, `hosts`, `port`, `ports`, `db_name` and `target_session_attrs` + parameters will be ignored. ### Parameters: - - `dsn`: full dsn connection string. + - `dsn`: Full dsn connection string. `postgres://postgres:postgres@localhost:5432/postgres?target_session_attrs=read-write` - - `username`: username of the user in postgres - - `password`: password of the user in postgres - - `host`: host of postgres - - `port`: port of postgres - - `db_name`: name of the database in postgres - - `max_db_pool_size`: maximum size of the connection pool + - `username`: Username of the user in the PostgreSQL + - `password`: Password of the user in the PostgreSQL + - `host`: Host of the PostgreSQL + - `hosts`: Hosts of the PostgreSQL + - `port`: Port of the PostgreSQL + - `ports`: Ports of the PostgreSQL + - `db_name`: Name of the database in PostgreSQL + - `target_session_attrs`: Specifies requirements of the session. + - `options`: Command line options used to configure the server + - `application_name`: Sets the application_name parameter on the server. + - `connect_timeout_sec`: The time limit in seconds applied to each socket-level + connection attempt. + Note that hostnames can resolve to multiple IP addresses, + and this limit is applied to each address. Defaults to no timeout. + - `connect_timeout_nanosec`: nanosec for connection timeout, + can be used only with connect_timeout_sec. + - `tcp_user_timeout_sec`: The time limit that + transmitted data may remain unacknowledged + before a connection is forcibly closed. + This is ignored for Unix domain socket connections. + It is only supported on systems where TCP_USER_TIMEOUT + is available and will default to the system default if omitted + or set to 0; on other systems, it has no effect. + - `tcp_user_timeout_nanosec`: nanosec for cp_user_timeout, + can be used only with tcp_user_timeout_sec. + - `keepalives`: Controls the use of TCP keepalive. + This option is ignored when connecting with Unix sockets. + Defaults to on. + - `keepalives_idle_sec`: The number of seconds of inactivity after + which a keepalive message is sent to the server. + This option is ignored when connecting with Unix sockets. + Defaults to 2 hours. + - `keepalives_idle_nanosec`: Nanosec for keepalives_idle_sec. + - `keepalives_interval_sec`: The time interval between TCP keepalive probes. + This option is ignored when connecting with Unix sockets. + - `keepalives_interval_nanosec`: Nanosec for keepalives_interval_sec. + - `keepalives_retries`: The maximum number of TCP keepalive probes + that will be sent before dropping a connection. + This option is ignored when connecting with Unix sockets. + - `load_balance_hosts`: Controls the order in which the client tries to connect + to the available hosts and addresses. + Once a connection attempt is successful no other + hosts and addresses will be tried. + This parameter is typically used in combination with multiple host names + or a DNS record that returns multiple IPs. + If set to disable, hosts and addresses will be tried in the order provided. + If set to random, hosts will be tried in a random order, and the IP addresses + resolved from a hostname will also be tried in a random order. + Defaults to disable. + - `max_db_pool_size`: maximum size of the connection pool. - `conn_recycling_method`: how a connection is recycled. """ async def execute( @@ -945,21 +1013,91 @@ def connect( username: Optional[str] = None, password: Optional[str] = None, host: Optional[str] = None, + hosts: Optional[List[str]] = None, port: Optional[int] = None, + ports: Optional[List[int]] = None, db_name: Optional[str] = None, + options: Optional[str] = None, + application_name: Optional[str] = None, + connect_timeout_sec: Optional[int] = None, + connect_timeout_nanosec: Optional[int] = None, + tcp_user_timeout_sec: Optional[int] = None, + tcp_user_timeout_nanosec: Optional[int] = None, + keepalives: Optional[bool] = None, + keepalives_idle_sec: Optional[int] = None, + keepalives_idle_nanosec: Optional[int] = None, + keepalives_interval_sec: Optional[int] = None, + keepalives_interval_nanosec: Optional[int] = None, + keepalives_retries: Optional[int] = None, + load_balance_hosts: Optional[ConnLoadBalanceHosts] = None, max_db_pool_size: int = 2, conn_recycling_method: Optional[ConnRecyclingMethod] = None, ) -> ConnectionPool: - """Create new connection pool. + """Create new PostgreSQL connection pool. + + It connects to the database and create pool. + + You cannot set the minimum size for the connection + pool, by it is 0. + `ConnectionPool` doesn't create connections on startup. + It makes new connection on demand. + + If you specify `dsn` parameter then `username`, `password`, + `host`, `hosts`, `port`, `ports`, `db_name` and `target_session_attrs` + parameters will be ignored. ### Parameters: - - `dsn`: full dsn connection string. + - `dsn`: Full dsn connection string. `postgres://postgres:postgres@localhost:5432/postgres?target_session_attrs=read-write` - - `username`: username of the user in postgres - - `password`: password of the user in postgres - - `host`: host of postgres - - `port`: port of postgres - - `db_name`: name of the database in postgres - - `max_db_pool_size`: maximum size of the connection pool + - `username`: Username of the user in the PostgreSQL + - `password`: Password of the user in the PostgreSQL + - `host`: Host of the PostgreSQL + - `hosts`: Hosts of the PostgreSQL + - `port`: Port of the PostgreSQL + - `ports`: Ports of the PostgreSQL + - `db_name`: Name of the database in PostgreSQL + - `target_session_attrs`: Specifies requirements of the session. + - `options`: Command line options used to configure the server + - `application_name`: Sets the application_name parameter on the server. + - `connect_timeout_sec`: The time limit in seconds applied to each socket-level + connection attempt. + Note that hostnames can resolve to multiple IP addresses, + and this limit is applied to each address. Defaults to no timeout. + - `connect_timeout_nanosec`: nanosec for connection timeout, + can be used only with connect_timeout_sec. + - `tcp_user_timeout_sec`: The time limit that + transmitted data may remain unacknowledged + before a connection is forcibly closed. + This is ignored for Unix domain socket connections. + It is only supported on systems where TCP_USER_TIMEOUT + is available and will default to the system default if omitted + or set to 0; on other systems, it has no effect. + - `tcp_user_timeout_nanosec`: nanosec for cp_user_timeout, + can be used only with tcp_user_timeout_sec. + - `keepalives`: Controls the use of TCP keepalive. + This option is ignored when connecting with Unix sockets. + Defaults to on. + - `keepalives_idle_sec`: The number of seconds of inactivity after + which a keepalive message is sent to the server. + This option is ignored when connecting with Unix sockets. + Defaults to 2 hours. + - `keepalives_idle_nanosec`: Nanosec for keepalives_idle_sec. + - `keepalives_interval_sec`: The time interval between TCP keepalive probes. + This option is ignored when connecting with Unix sockets. + - `keepalives_interval_nanosec`: Nanosec for keepalives_interval_sec. + - `keepalives_retries`: The maximum number of TCP keepalive probes + that will be sent before dropping a connection. + This option is ignored when connecting with Unix sockets. + - `load_balance_hosts`: Controls the order in which the client tries to connect + to the available hosts and addresses. + Once a connection attempt is successful no other + hosts and addresses will be tried. + This parameter is typically used in combination with multiple host names + or a DNS record that returns multiple IPs. + If set to disable, hosts and addresses will be tried in the order provided. + If set to random, hosts will be tried in a random order, and the IP addresses + resolved from a hostname will also be tried in a random order. + Defaults to disable. + - `max_db_pool_size`: maximum size of the connection pool. - `conn_recycling_method`: how a connection is recycled. """ diff --git a/src/driver/common_options.rs b/src/driver/common_options.rs index ec4c9e8d..31155e84 100644 --- a/src/driver/common_options.rs +++ b/src/driver/common_options.rs @@ -1,6 +1,6 @@ use deadpool_postgres::RecyclingMethod; use pyo3::pyclass; -use tokio_postgres::config::LoadBalanceHosts; +use tokio_postgres::config::{LoadBalanceHosts, TargetSessionAttrs}; #[pyclass] #[derive(Clone, Copy)] @@ -39,3 +39,25 @@ impl ConnLoadBalanceHosts { } } } + +#[pyclass] +#[derive(Clone, Copy)] +pub enum ConnTargetSessionAttrs { + /// No special properties are required. + Any, + /// The session must allow writes. + ReadWrite, + /// The session allow only reads. + ReadOnly, +} + +impl ConnTargetSessionAttrs { + #[must_use] + pub fn to_internal(&self) -> TargetSessionAttrs { + match self { + ConnTargetSessionAttrs::Any => TargetSessionAttrs::Any, + ConnTargetSessionAttrs::ReadWrite => TargetSessionAttrs::ReadWrite, + ConnTargetSessionAttrs::ReadOnly => TargetSessionAttrs::ReadOnly, + } + } +} diff --git a/src/driver/connection_pool.rs b/src/driver/connection_pool.rs index 491d6892..aaf56eda 100644 --- a/src/driver/connection_pool.rs +++ b/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::{time::Duration, vec}; +use std::vec; use tokio_postgres::{NoTls, Row}; use crate::{ @@ -11,7 +11,7 @@ use crate::{ }; use super::{ - common_options::{ConnLoadBalanceHosts, ConnRecyclingMethod}, + common_options::{ConnLoadBalanceHosts, ConnRecyclingMethod, ConnTargetSessionAttrs}, connection::Connection, utils::build_connection_config, }; @@ -27,15 +27,22 @@ pub fn connect( username: Option, password: Option, host: Option, + hosts: Option>, port: Option, + ports: Option>, db_name: Option, + target_session_attrs: Option, options: Option, application_name: Option, - connect_timeout: Option, - tcp_user_timeout: Option, + connect_timeout_sec: Option, + connect_timeout_nanosec: Option, + tcp_user_timeout_sec: Option, + tcp_user_timeout_nanosec: Option, keepalives: Option, - keepalives_idle: Option, - keepalives_interval: Option, + keepalives_idle_sec: Option, + keepalives_idle_nanosec: Option, + keepalives_interval_sec: Option, + keepalives_interval_nanosec: Option, keepalives_retries: Option, load_balance_hosts: Option, @@ -55,15 +62,22 @@ pub fn connect( username, password, host, + hosts, port, + ports, db_name, + target_session_attrs, options, application_name, - connect_timeout, - tcp_user_timeout, + connect_timeout_sec, + connect_timeout_nanosec, + tcp_user_timeout_sec, + tcp_user_timeout_nanosec, keepalives, - keepalives_idle, - keepalives_interval, + keepalives_idle_sec, + keepalives_idle_nanosec, + keepalives_interval_sec, + keepalives_interval_nanosec, keepalives_retries, load_balance_hosts, )?; @@ -106,15 +120,22 @@ impl ConnectionPool { username: Option, password: Option, host: Option, + hosts: Option>, port: Option, + ports: Option>, db_name: Option, + target_session_attrs: Option, options: Option, application_name: Option, - connect_timeout: Option, - tcp_user_timeout: Option, + connect_timeout_sec: Option, + connect_timeout_nanosec: Option, + tcp_user_timeout_sec: Option, + tcp_user_timeout_nanosec: Option, keepalives: Option, - keepalives_idle: Option, - keepalives_interval: Option, + keepalives_idle_sec: Option, + keepalives_idle_nanosec: Option, + keepalives_interval_sec: Option, + keepalives_interval_nanosec: Option, keepalives_retries: Option, load_balance_hosts: Option, max_db_pool_size: Option, @@ -125,15 +146,22 @@ impl ConnectionPool { username, password, host, + hosts, port, + ports, db_name, + target_session_attrs, options, application_name, - connect_timeout, - tcp_user_timeout, + connect_timeout_sec, + connect_timeout_nanosec, + tcp_user_timeout_sec, + tcp_user_timeout_nanosec, keepalives, - keepalives_idle, - keepalives_interval, + keepalives_idle_sec, + keepalives_idle_nanosec, + keepalives_interval_sec, + keepalives_interval_nanosec, keepalives_retries, load_balance_hosts, max_db_pool_size, diff --git a/src/driver/utils.rs b/src/driver/utils.rs index 0e3ee0f3..8c175afc 100644 --- a/src/driver/utils.rs +++ b/src/driver/utils.rs @@ -1,8 +1,8 @@ use std::{str::FromStr, time::Duration}; -use crate::exceptions::rust_errors::RustPSQLDriverPyResult; +use crate::exceptions::rust_errors::{RustPSQLDriverError, RustPSQLDriverPyResult}; -use super::common_options::ConnLoadBalanceHosts; +use super::common_options::{ConnLoadBalanceHosts, ConnTargetSessionAttrs}; /// Create new config. /// @@ -14,38 +14,88 @@ pub fn build_connection_config( username: Option, password: Option, host: Option, + hosts: Option>, port: Option, + ports: Option>, db_name: Option, + target_session_attrs: Option, options: Option, application_name: Option, - connect_timeout: Option, - tcp_user_timeout: Option, + connect_timeout_sec: Option, + connect_timeout_nanosec: Option, + tcp_user_timeout_sec: Option, + tcp_user_timeout_nanosec: Option, keepalives: Option, - keepalives_idle: Option, - keepalives_interval: Option, + keepalives_idle_sec: Option, + keepalives_idle_nanosec: Option, + keepalives_interval_sec: Option, + keepalives_interval_nanosec: Option, keepalives_retries: Option, load_balance_hosts: Option, ) -> RustPSQLDriverPyResult { - if let Some(dsn_string) = dsn { - return Ok(tokio_postgres::Config::from_str(&dsn_string)?); + if tcp_user_timeout_nanosec.is_some() && tcp_user_timeout_sec.is_none() { + return Err(RustPSQLDriverError::DataBasePoolConfigurationError( + "tcp_user_timeout_nanosec must be used with tcp_user_timeout_sec param.".into(), + )); } - 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 connect_timeout_nanosec.is_some() && connect_timeout_sec.is_none() { + return Err(RustPSQLDriverError::DataBasePoolConfigurationError( + "connect_timeout_nanosec must be used with connect_timeout_sec param.".into(), + )); } - if let Some(host) = host { - pg_config.host(&host); + + if keepalives_idle_nanosec.is_some() && keepalives_idle_sec.is_none() { + return Err(RustPSQLDriverError::DataBasePoolConfigurationError( + "keepalives_idle_nanosec must be used with keepalives_idle_sec param.".into(), + )); } - if let Some(port) = port { - pg_config.port(port); + if keepalives_interval_nanosec.is_some() && keepalives_interval_sec.is_none() { + return Err(RustPSQLDriverError::DataBasePoolConfigurationError( + "keepalives_interval_nanosec must be used with keepalives_interval_sec param.".into(), + )); } - if let Some(db_name) = db_name { - pg_config.dbname(&db_name); + 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(hosts) = hosts { + for single_host in hosts { + pg_config.host(&single_host); + } + } + + if let Some(host) = host { + pg_config.host(&host); + } + + if let Some(ports) = ports { + for single_port in ports { + pg_config.port(single_port); + } + } + + if let Some(port) = port { + pg_config.port(port); + } + + if let Some(db_name) = db_name { + pg_config.dbname(&db_name); + } + + if let Some(target_session_attrs) = target_session_attrs { + pg_config.target_session_attrs(target_session_attrs.to_internal()); + } } if let Some(options) = options { @@ -56,24 +106,36 @@ pub fn build_connection_config( pg_config.application_name(&application_name); } - if let Some(connect_timeout) = connect_timeout { - pg_config.connect_timeout(connect_timeout); + if let Some(connect_timeout_sec) = connect_timeout_sec { + pg_config.connect_timeout(Duration::new( + connect_timeout_sec, + connect_timeout_nanosec.unwrap_or_default(), + )); } - if let Some(tcp_user_timeout) = tcp_user_timeout { - pg_config.tcp_user_timeout(tcp_user_timeout); + if let Some(tcp_user_timeout_sec) = tcp_user_timeout_sec { + pg_config.tcp_user_timeout(Duration::new( + tcp_user_timeout_sec, + tcp_user_timeout_nanosec.unwrap_or_default(), + )); } 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_idle_sec) = keepalives_idle_sec { + pg_config.keepalives_idle(Duration::new( + keepalives_idle_sec, + keepalives_idle_nanosec.unwrap_or_default(), + )); } - if let Some(keepalives_interval) = keepalives_interval { - pg_config.keepalives_interval(keepalives_interval); + if let Some(keepalives_interval_sec) = keepalives_interval_sec { + pg_config.keepalives_interval(Duration::new( + keepalives_interval_sec, + keepalives_interval_nanosec.unwrap_or_default(), + )); } if let Some(keepalives_retries) = keepalives_retries { diff --git a/src/lib.rs b/src/lib.rs index 2392f098..0bf922a7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -23,6 +23,8 @@ fn psqlpy(py: Python<'_>, pymod: &Bound<'_, PyModule>) -> PyResult<()> { pymod.add_class::()?; pymod.add_class::()?; pymod.add_class::()?; + pymod.add_class::()?; + pymod.add_class::()?; pymod.add_class::()?; pymod.add_class::()?; add_module(py, pymod, "extra_types", extra_types_module)?; From d938c502dfd8bd9160ed605bf2fb920164dda5d8 Mon Sep 17 00:00:00 2001 From: "chandr-andr (Kiselev Aleksandr)" Date: Sun, 28 Apr 2024 22:08:41 +0200 Subject: [PATCH 3/7] Added more parameters to configure connection Signed-off-by: chandr-andr (Kiselev Aleksandr) --- python/psqlpy/_internal/__init__.pyi | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/python/psqlpy/_internal/__init__.pyi b/python/psqlpy/_internal/__init__.pyi index 3033a864..85d8eee3 100644 --- a/python/psqlpy/_internal/__init__.pyi +++ b/python/psqlpy/_internal/__init__.pyi @@ -108,9 +108,19 @@ class ConnLoadBalanceHosts(Enum): """Load balancing configuration.""" # Make connection attempts to hosts in the order provided. - Disable = (1,) + Disable = 1 # Make connection attempts to hosts in a random order. - Random = (2,) + Random = 2 + +class ConnTargetSessionAttrs(Enum): + """Properties required of a session.""" + + # No special properties are required. + Any = 1 + # The session must allow writes. + ReadWrite = 2 + # The session allow only reads. + ReadOnly = 3 class ReadVariant(Enum): """Class for Read Variant for transaction.""" From bb05b9fd695be71905bc111721b7d7b539856e5b Mon Sep 17 00:00:00 2001 From: "chandr-andr (Kiselev Aleksandr)" Date: Sun, 28 Apr 2024 22:08:45 +0200 Subject: [PATCH 4/7] Added more parameters to configure connection Signed-off-by: chandr-andr (Kiselev Aleksandr) --- python/psqlpy/_internal/__init__.pyi | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/psqlpy/_internal/__init__.pyi b/python/psqlpy/_internal/__init__.pyi index 85d8eee3..be344b5e 100644 --- a/python/psqlpy/_internal/__init__.pyi +++ b/python/psqlpy/_internal/__init__.pyi @@ -891,6 +891,7 @@ class ConnectionPool: port: Optional[int] = None, ports: Optional[List[int]] = None, db_name: Optional[str] = None, + target_session_attrs: Optional[ConnTargetSessionAttrs] = None, options: Optional[str] = None, application_name: Optional[str] = None, connect_timeout_sec: Optional[int] = None, @@ -1027,6 +1028,7 @@ def connect( port: Optional[int] = None, ports: Optional[List[int]] = None, db_name: Optional[str] = None, + target_session_attrs: Optional[ConnTargetSessionAttrs] = None, options: Optional[str] = None, application_name: Optional[str] = None, connect_timeout_sec: Optional[int] = None, From a019bdcd888267330b421c06b2e91ffba0ec6cee Mon Sep 17 00:00:00 2001 From: "chandr-andr (Kiselev Aleksandr)" Date: Sun, 28 Apr 2024 22:17:38 +0200 Subject: [PATCH 5/7] Added more parameters to configure connection Signed-off-by: chandr-andr (Kiselev Aleksandr) --- python/psqlpy/__init__.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/psqlpy/__init__.py b/python/psqlpy/__init__.py index 98a7c154..04201e72 100644 --- a/python/psqlpy/__init__.py +++ b/python/psqlpy/__init__.py @@ -1,7 +1,9 @@ from psqlpy._internal import ( Connection, ConnectionPool, + ConnLoadBalanceHosts, ConnRecyclingMethod, + ConnTargetSessionAttrs, Cursor, IsolationLevel, QueryResult, @@ -22,4 +24,6 @@ "IsolationLevel", "ReadVariant", "connect", + "ConnLoadBalanceHosts", + "ConnTargetSessionAttrs", ] From 2b23c4572316aba5e208f4d63d1f960f48855df3 Mon Sep 17 00:00:00 2001 From: "chandr-andr (Kiselev Aleksandr)" Date: Sun, 28 Apr 2024 23:10:36 +0200 Subject: [PATCH 6/7] Added more parameters to configure connection, added tests Signed-off-by: chandr-andr (Kiselev Aleksandr) --- python/tests/test_connection_pool.py | 79 +++++++++++++++++++++++++++- src/driver/utils.rs | 5 +- 2 files changed, 81 insertions(+), 3 deletions(-) diff --git a/python/tests/test_connection_pool.py b/python/tests/test_connection_pool.py index 467b2899..a22d4728 100644 --- a/python/tests/test_connection_pool.py +++ b/python/tests/test_connection_pool.py @@ -1,7 +1,15 @@ import pytest -from psqlpy import Connection, ConnectionPool, ConnRecyclingMethod, QueryResult, connect -from psqlpy.exceptions import RustPSQLDriverPyBaseError +from psqlpy import ( + Connection, + ConnectionPool, + ConnLoadBalanceHosts, + ConnRecyclingMethod, + ConnTargetSessionAttrs, + QueryResult, + connect, +) +from psqlpy.exceptions import DBPoolConfigurationError, RustPSQLDriverPyBaseError pytestmark = pytest.mark.anyio @@ -68,6 +76,73 @@ async def test_pool_conn_recycling_method( await pg_pool.execute("SELECT 1") +async def test_build_pool_failure() -> None: + with pytest.raises(expected_exception=DBPoolConfigurationError): + ConnectionPool( + dsn="postgres://postgres:postgres@localhost:5432/psqlpy_test", + connect_timeout_nanosec=12, + ) + with pytest.raises(expected_exception=DBPoolConfigurationError): + ConnectionPool( + dsn="postgres://postgres:postgres@localhost:5432/psqlpy_test", + connect_timeout_nanosec=12, + ) + with pytest.raises(expected_exception=DBPoolConfigurationError): + ConnectionPool( + dsn="postgres://postgres:postgres@localhost:5432/psqlpy_test", + keepalives_idle_nanosec=12, + ) + with pytest.raises(expected_exception=DBPoolConfigurationError): + ConnectionPool( + dsn="postgres://postgres:postgres@localhost:5432/psqlpy_test", + keepalives_interval_nanosec=12, + ) + + +@pytest.mark.parametrize( + "target_session_attrs", + [ + ConnTargetSessionAttrs.Any, + ConnTargetSessionAttrs.ReadWrite, + ConnTargetSessionAttrs.ReadOnly, + ], +) +async def test_pool_target_session_attrs( + target_session_attrs: ConnTargetSessionAttrs, +) -> None: + pg_pool = ConnectionPool( + db_name="psqlpy_test", + host="localhost", + username="postgres", + password="postgres", # noqa: S106 + target_session_attrs=target_session_attrs, + ) + + if target_session_attrs == ConnTargetSessionAttrs.ReadOnly: + with pytest.raises(expected_exception=RustPSQLDriverPyBaseError): + await pg_pool.execute("SELECT 1") + else: + await pg_pool.execute("SELECT 1") + + +@pytest.mark.parametrize( + "load_balance_hosts", + [ + ConnLoadBalanceHosts.Disable, + ConnLoadBalanceHosts.Random, + ], +) +async def test_pool_load_balance_hosts( + load_balance_hosts: ConnLoadBalanceHosts, +) -> None: + pg_pool = ConnectionPool( + dsn="postgres://postgres:postgres@localhost:5432/psqlpy_test", + load_balance_hosts=load_balance_hosts, + ) + + await pg_pool.execute("SELECT 1") + + async def test_close_connection_pool() -> None: """Test that `close` method closes connection pool.""" pg_pool = ConnectionPool( diff --git a/src/driver/utils.rs b/src/driver/utils.rs index 8c175afc..5e44c07c 100644 --- a/src/driver/utils.rs +++ b/src/driver/utils.rs @@ -64,8 +64,11 @@ pub fn build_connection_config( } else { pg_config = tokio_postgres::Config::new(); - if let (Some(password), Some(username)) = (password, username) { + if let Some(password) = password { pg_config.password(&password); + } + + if let Some(username) = username { pg_config.user(&username); } From ac5366d3647aaa96ce489a2724115e978e3d13d8 Mon Sep 17 00:00:00 2001 From: "chandr-andr (Kiselev Aleksandr)" Date: Mon, 29 Apr 2024 16:52:50 +0200 Subject: [PATCH 7/7] Added more parameters to connection pool docs Signed-off-by: chandr-andr (Kiselev Aleksandr) --- python/psqlpy/__init__.py | 8 ++++---- python/psqlpy/_internal/__init__.pyi | 12 ++++++------ python/tests/test_connection_pool.py | 20 ++++++++++---------- src/driver/common_options.rs | 23 +++++++++++------------ src/driver/connection_pool.rs | 10 +++++----- src/driver/utils.rs | 6 +++--- src/lib.rs | 4 ++-- 7 files changed, 41 insertions(+), 42 deletions(-) diff --git a/python/psqlpy/__init__.py b/python/psqlpy/__init__.py index 04201e72..baaa76f8 100644 --- a/python/psqlpy/__init__.py +++ b/python/psqlpy/__init__.py @@ -1,14 +1,14 @@ from psqlpy._internal import ( Connection, ConnectionPool, - ConnLoadBalanceHosts, ConnRecyclingMethod, - ConnTargetSessionAttrs, Cursor, IsolationLevel, + LoadBalanceHosts, QueryResult, ReadVariant, SingleQueryResult, + TargetSessionAttrs, Transaction, connect, ) @@ -24,6 +24,6 @@ "IsolationLevel", "ReadVariant", "connect", - "ConnLoadBalanceHosts", - "ConnTargetSessionAttrs", + "LoadBalanceHosts", + "TargetSessionAttrs", ] diff --git a/python/psqlpy/_internal/__init__.pyi b/python/psqlpy/_internal/__init__.pyi index be344b5e..af6a8374 100644 --- a/python/psqlpy/_internal/__init__.pyi +++ b/python/psqlpy/_internal/__init__.pyi @@ -104,7 +104,7 @@ class IsolationLevel(Enum): RepeatableRead = 3 Serializable = 4 -class ConnLoadBalanceHosts(Enum): +class LoadBalanceHosts(Enum): """Load balancing configuration.""" # Make connection attempts to hosts in the order provided. @@ -112,7 +112,7 @@ class ConnLoadBalanceHosts(Enum): # Make connection attempts to hosts in a random order. Random = 2 -class ConnTargetSessionAttrs(Enum): +class TargetSessionAttrs(Enum): """Properties required of a session.""" # No special properties are required. @@ -891,7 +891,7 @@ class ConnectionPool: port: Optional[int] = None, ports: Optional[List[int]] = None, db_name: Optional[str] = None, - target_session_attrs: Optional[ConnTargetSessionAttrs] = None, + target_session_attrs: Optional[TargetSessionAttrs] = None, options: Optional[str] = None, application_name: Optional[str] = None, connect_timeout_sec: Optional[int] = None, @@ -904,7 +904,7 @@ class ConnectionPool: keepalives_interval_sec: Optional[int] = None, keepalives_interval_nanosec: Optional[int] = None, keepalives_retries: Optional[int] = None, - load_balance_hosts: Optional[ConnLoadBalanceHosts] = None, + load_balance_hosts: Optional[LoadBalanceHosts] = None, max_db_pool_size: int = 2, conn_recycling_method: Optional[ConnRecyclingMethod] = None, ) -> None: @@ -1028,7 +1028,7 @@ def connect( port: Optional[int] = None, ports: Optional[List[int]] = None, db_name: Optional[str] = None, - target_session_attrs: Optional[ConnTargetSessionAttrs] = None, + target_session_attrs: Optional[TargetSessionAttrs] = None, options: Optional[str] = None, application_name: Optional[str] = None, connect_timeout_sec: Optional[int] = None, @@ -1041,7 +1041,7 @@ def connect( keepalives_interval_sec: Optional[int] = None, keepalives_interval_nanosec: Optional[int] = None, keepalives_retries: Optional[int] = None, - load_balance_hosts: Optional[ConnLoadBalanceHosts] = None, + load_balance_hosts: Optional[LoadBalanceHosts] = None, max_db_pool_size: int = 2, conn_recycling_method: Optional[ConnRecyclingMethod] = None, ) -> ConnectionPool: diff --git a/python/tests/test_connection_pool.py b/python/tests/test_connection_pool.py index a22d4728..e2962dbf 100644 --- a/python/tests/test_connection_pool.py +++ b/python/tests/test_connection_pool.py @@ -3,10 +3,10 @@ from psqlpy import ( Connection, ConnectionPool, - ConnLoadBalanceHosts, ConnRecyclingMethod, - ConnTargetSessionAttrs, + LoadBalanceHosts, QueryResult, + TargetSessionAttrs, connect, ) from psqlpy.exceptions import DBPoolConfigurationError, RustPSQLDriverPyBaseError @@ -102,13 +102,13 @@ async def test_build_pool_failure() -> None: @pytest.mark.parametrize( "target_session_attrs", [ - ConnTargetSessionAttrs.Any, - ConnTargetSessionAttrs.ReadWrite, - ConnTargetSessionAttrs.ReadOnly, + TargetSessionAttrs.Any, + TargetSessionAttrs.ReadWrite, + TargetSessionAttrs.ReadOnly, ], ) async def test_pool_target_session_attrs( - target_session_attrs: ConnTargetSessionAttrs, + target_session_attrs: TargetSessionAttrs, ) -> None: pg_pool = ConnectionPool( db_name="psqlpy_test", @@ -118,7 +118,7 @@ async def test_pool_target_session_attrs( target_session_attrs=target_session_attrs, ) - if target_session_attrs == ConnTargetSessionAttrs.ReadOnly: + if target_session_attrs == TargetSessionAttrs.ReadOnly: with pytest.raises(expected_exception=RustPSQLDriverPyBaseError): await pg_pool.execute("SELECT 1") else: @@ -128,12 +128,12 @@ async def test_pool_target_session_attrs( @pytest.mark.parametrize( "load_balance_hosts", [ - ConnLoadBalanceHosts.Disable, - ConnLoadBalanceHosts.Random, + LoadBalanceHosts.Disable, + LoadBalanceHosts.Random, ], ) async def test_pool_load_balance_hosts( - load_balance_hosts: ConnLoadBalanceHosts, + load_balance_hosts: LoadBalanceHosts, ) -> None: pg_pool = ConnectionPool( dsn="postgres://postgres:postgres@localhost:5432/psqlpy_test", diff --git a/src/driver/common_options.rs b/src/driver/common_options.rs index 31155e84..2d906234 100644 --- a/src/driver/common_options.rs +++ b/src/driver/common_options.rs @@ -1,6 +1,5 @@ use deadpool_postgres::RecyclingMethod; use pyo3::pyclass; -use tokio_postgres::config::{LoadBalanceHosts, TargetSessionAttrs}; #[pyclass] #[derive(Clone, Copy)] @@ -23,26 +22,26 @@ impl ConnRecyclingMethod { #[pyclass] #[derive(Clone, Copy)] -pub enum ConnLoadBalanceHosts { +pub enum LoadBalanceHosts { /// Make connection attempts to hosts in the order provided. Disable, /// Make connection attempts to hosts in a random order. Random, } -impl ConnLoadBalanceHosts { +impl LoadBalanceHosts { #[must_use] - pub fn to_internal(&self) -> LoadBalanceHosts { + pub fn to_internal(&self) -> tokio_postgres::config::LoadBalanceHosts { match self { - ConnLoadBalanceHosts::Disable => LoadBalanceHosts::Disable, - ConnLoadBalanceHosts::Random => LoadBalanceHosts::Random, + LoadBalanceHosts::Disable => tokio_postgres::config::LoadBalanceHosts::Disable, + LoadBalanceHosts::Random => tokio_postgres::config::LoadBalanceHosts::Random, } } } #[pyclass] #[derive(Clone, Copy)] -pub enum ConnTargetSessionAttrs { +pub enum TargetSessionAttrs { /// No special properties are required. Any, /// The session must allow writes. @@ -51,13 +50,13 @@ pub enum ConnTargetSessionAttrs { ReadOnly, } -impl ConnTargetSessionAttrs { +impl TargetSessionAttrs { #[must_use] - pub fn to_internal(&self) -> TargetSessionAttrs { + pub fn to_internal(&self) -> tokio_postgres::config::TargetSessionAttrs { match self { - ConnTargetSessionAttrs::Any => TargetSessionAttrs::Any, - ConnTargetSessionAttrs::ReadWrite => TargetSessionAttrs::ReadWrite, - ConnTargetSessionAttrs::ReadOnly => TargetSessionAttrs::ReadOnly, + TargetSessionAttrs::Any => tokio_postgres::config::TargetSessionAttrs::Any, + TargetSessionAttrs::ReadWrite => tokio_postgres::config::TargetSessionAttrs::ReadWrite, + TargetSessionAttrs::ReadOnly => tokio_postgres::config::TargetSessionAttrs::ReadOnly, } } } diff --git a/src/driver/connection_pool.rs b/src/driver/connection_pool.rs index aaf56eda..caa36e32 100644 --- a/src/driver/connection_pool.rs +++ b/src/driver/connection_pool.rs @@ -11,7 +11,7 @@ use crate::{ }; use super::{ - common_options::{ConnLoadBalanceHosts, ConnRecyclingMethod, ConnTargetSessionAttrs}, + common_options::{ConnRecyclingMethod, LoadBalanceHosts, TargetSessionAttrs}, connection::Connection, utils::build_connection_config, }; @@ -31,7 +31,7 @@ pub fn connect( port: Option, ports: Option>, db_name: Option, - target_session_attrs: Option, + target_session_attrs: Option, options: Option, application_name: Option, connect_timeout_sec: Option, @@ -44,7 +44,7 @@ pub fn connect( keepalives_interval_sec: Option, keepalives_interval_nanosec: Option, keepalives_retries: Option, - load_balance_hosts: Option, + load_balance_hosts: Option, max_db_pool_size: Option, conn_recycling_method: Option, @@ -124,7 +124,7 @@ impl ConnectionPool { port: Option, ports: Option>, db_name: Option, - target_session_attrs: Option, + target_session_attrs: Option, options: Option, application_name: Option, connect_timeout_sec: Option, @@ -137,7 +137,7 @@ impl ConnectionPool { keepalives_interval_sec: Option, keepalives_interval_nanosec: Option, keepalives_retries: Option, - load_balance_hosts: Option, + load_balance_hosts: Option, max_db_pool_size: Option, conn_recycling_method: Option, ) -> RustPSQLDriverPyResult { diff --git a/src/driver/utils.rs b/src/driver/utils.rs index 5e44c07c..df4cf1b2 100644 --- a/src/driver/utils.rs +++ b/src/driver/utils.rs @@ -2,7 +2,7 @@ use std::{str::FromStr, time::Duration}; use crate::exceptions::rust_errors::{RustPSQLDriverError, RustPSQLDriverPyResult}; -use super::common_options::{ConnLoadBalanceHosts, ConnTargetSessionAttrs}; +use super::common_options::{LoadBalanceHosts, TargetSessionAttrs}; /// Create new config. /// @@ -18,7 +18,7 @@ pub fn build_connection_config( port: Option, ports: Option>, db_name: Option, - target_session_attrs: Option, + target_session_attrs: Option, options: Option, application_name: Option, connect_timeout_sec: Option, @@ -31,7 +31,7 @@ pub fn build_connection_config( keepalives_interval_sec: Option, keepalives_interval_nanosec: Option, keepalives_retries: Option, - load_balance_hosts: Option, + load_balance_hosts: Option, ) -> RustPSQLDriverPyResult { if tcp_user_timeout_nanosec.is_some() && tcp_user_timeout_sec.is_none() { return Err(RustPSQLDriverError::DataBasePoolConfigurationError( diff --git a/src/lib.rs b/src/lib.rs index 0bf922a7..cf26c4da 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -23,8 +23,8 @@ fn psqlpy(py: Python<'_>, pymod: &Bound<'_, PyModule>) -> PyResult<()> { pymod.add_class::()?; pymod.add_class::()?; pymod.add_class::()?; - pymod.add_class::()?; - pymod.add_class::()?; + pymod.add_class::()?; + pymod.add_class::()?; pymod.add_class::()?; pymod.add_class::()?; add_module(py, pymod, "extra_types", extra_types_module)?;