diff --git a/python/psqlpy/__init__.py b/python/psqlpy/__init__.py index bc268f3c..98a7c154 100644 --- a/python/psqlpy/__init__.py +++ b/python/psqlpy/__init__.py @@ -8,6 +8,7 @@ ReadVariant, SingleQueryResult, Transaction, + connect, ) __all__ = [ @@ -20,4 +21,5 @@ "ConnRecyclingMethod", "IsolationLevel", "ReadVariant", + "connect", ] diff --git a/python/psqlpy/_internal/__init__.pyi b/python/psqlpy/_internal/__init__.pyi index 770a5c8a..cc6fc052 100644 --- a/python/psqlpy/_internal/__init__.pyi +++ b/python/psqlpy/_internal/__init__.pyi @@ -936,3 +936,27 @@ class ConnectionPool: """ def close(self: Self) -> None: """Close the connection pool.""" + +def connect( + dsn: Optional[str] = None, + username: Optional[str] = None, + password: Optional[str] = None, + host: Optional[str] = None, + port: Optional[int] = None, + db_name: Optional[str] = None, + max_db_pool_size: int = 2, + conn_recycling_method: Optional[ConnRecyclingMethod] = None, +) -> ConnectionPool: + """Create new connection pool. + + ### Parameters: + - `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 + - `conn_recycling_method`: how a connection is recycled. + """ diff --git a/python/tests/test_connection_pool.py b/python/tests/test_connection_pool.py index a3eba491..467b2899 100644 --- a/python/tests/test_connection_pool.py +++ b/python/tests/test_connection_pool.py @@ -1,11 +1,20 @@ import pytest -from psqlpy import Connection, ConnectionPool, ConnRecyclingMethod, QueryResult +from psqlpy import Connection, ConnectionPool, ConnRecyclingMethod, QueryResult, connect from psqlpy.exceptions import RustPSQLDriverPyBaseError pytestmark = pytest.mark.anyio +async def test_connect_func() -> None: + """Test that connect function makes new connection pool.""" + pg_pool = connect( + dsn="postgres://postgres:postgres@localhost:5432/psqlpy_test", + ) + + await pg_pool.execute("SELECT 1") + + async def test_pool_dsn_startup() -> None: """Test that connection pool can startup with dsn.""" pg_pool = ConnectionPool( diff --git a/src/driver/connection_pool.rs b/src/driver/connection_pool.rs index 7ca1280c..fad7cb9e 100644 --- a/src/driver/connection_pool.rs +++ b/src/driver/connection_pool.rs @@ -1,303 +1,84 @@ use crate::runtime::tokio_runtime; use deadpool_postgres::{Manager, ManagerConfig, Object, Pool, RecyclingMethod}; -use pyo3::{pyclass, pymethods, PyAny}; +use pyo3::{pyclass, pyfunction, pymethods, PyAny}; use std::{str::FromStr, vec}; use tokio_postgres::{NoTls, Row}; use crate::{ - // common::rustdriver_future, exceptions::rust_errors::{RustPSQLDriverError, RustPSQLDriverPyResult}, query_result::PSQLDriverPyQueryResult, value_converter::{convert_parameters, PythonDTO, QueryParameter}, }; -use super::{ - common_options::ConnRecyclingMethod, - connection::Connection, - // connection::{Connection, RustConnection}, -}; +use super::{common_options::ConnRecyclingMethod, connection::Connection}; -/// `PSQLPool` is for internal use only. +/// Make new connection pool. /// -/// It is not exposed to python. -// pub struct RustPSQLPool { -// dsn: Option, -// username: Option, -// password: Option, -// host: Option, -// port: Option, -// db_name: Option, -// max_db_pool_size: Option, -// conn_recycling_method: Option, -// db_pool: Option, -// } - -// impl RustPSQLPool { -// /// Create new `RustPSQLPool`. -// #[must_use] -// #[allow(clippy::too_many_arguments)] -// pub fn new( -// dsn: Option, -// username: Option, -// password: Option, -// host: Option, -// port: Option, -// db_name: Option, -// max_db_pool_size: Option, -// conn_recycling_method: Option, -// ) -> Self { -// RustPSQLPool { -// dsn, -// username, -// password, -// host, -// port, -// db_name, -// max_db_pool_size, -// conn_recycling_method, -// db_pool: None, -// } -// } -// } - -// impl RustPSQLPool { -// /// Return new single connection. -// /// -// /// # Errors -// /// May return Err Result if cannot get new connection from the pool. -// pub async fn inner_connection(&self) -> RustPSQLDriverPyResult { -// let db_pool_manager = self -// .db_pool -// .as_ref() -// .ok_or(RustPSQLDriverError::DatabasePoolError( -// "Database pool is not initialized".into(), -// ))? -// .get() -// .await?; - -// Ok(Connection::new(Arc::new(RustConnection::new(Arc::new( -// tokio::sync::RwLock::new(db_pool_manager), -// ))))) -// } -// /// Execute querystring with parameters. -// /// -// /// Prepare statement and cache it, then execute. -// /// -// /// # Errors -// /// May return Err Result if cannot retrieve new connection -// /// or prepare statement or execute statement. -// pub async fn inner_execute( -// &self, -// querystring: String, -// parameters: Vec, -// prepared: bool, -// ) -> RustPSQLDriverPyResult { -// let db_pool_manager = self -// .db_pool -// .as_ref() -// .ok_or(RustPSQLDriverError::DatabasePoolError( -// "Database pool is not initialized".into(), -// ))? -// .get() -// .await?; - -// let vec_parameters: Vec<&QueryParameter> = parameters -// .iter() -// .map(|param| param as &QueryParameter) -// .collect(); - -// let result = if prepared { -// db_pool_manager -// .query( -// &db_pool_manager.prepare_cached(&querystring).await?, -// &vec_parameters.into_boxed_slice(), -// ) -// .await? -// } else { -// db_pool_manager -// .query(&querystring, &vec_parameters.into_boxed_slice()) -// .await? -// }; -// Ok(PSQLDriverPyQueryResult::new(result)) -// } - -// /// Create new Database pool. -// /// -// /// # Errors -// /// May return Err Result if Database pool is already initialized, -// /// `max_db_pool_size` is less than 2 or it's impossible to build db pool. -// pub fn inner_startup(mut self) -> RustPSQLDriverPyResult { -// let dsn = self.dsn.clone(); -// let password = self.password.clone(); -// let username = self.username.clone(); -// let db_host = self.host.clone(); -// let db_port = self.port; -// let db_name = self.db_name.clone(); -// let conn_recycling_method = self.conn_recycling_method; -// let max_db_pool_size = self.max_db_pool_size; - -// if self.db_pool.is_some() { -// return Err(RustPSQLDriverError::DatabasePoolError( -// "Database pool is already initialized".into(), -// )); -// } - -// if let Some(max_db_pool_size) = max_db_pool_size { -// if max_db_pool_size < 2 { -// return Err(RustPSQLDriverError::DataBasePoolConfigurationError( -// "Maximum database pool size must be more than 1".into(), -// )); -// } -// } - -// 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(db_host) = db_host { -// pg_config.host(&db_host); -// } - -// if let Some(db_port) = db_port { -// pg_config.port(db_port); -// } - -// if let Some(db_name) = db_name { -// pg_config.dbname(&db_name); -// } -// } - -// let mgr_config: ManagerConfig; -// if let Some(conn_recycling_method) = conn_recycling_method { -// mgr_config = ManagerConfig { -// recycling_method: conn_recycling_method.to_internal(), -// } -// } else { -// mgr_config = ManagerConfig { -// recycling_method: RecyclingMethod::Fast, -// }; -// } -// let mgr = Manager::from_config(pg_config, NoTls, mgr_config); - -// let mut db_pool_builder = Pool::builder(mgr); -// if let Some(max_db_pool_size) = max_db_pool_size { -// db_pool_builder = db_pool_builder.max_size(max_db_pool_size); -// } - -// self.db_pool = Some(db_pool_builder.build()?); -// Ok(self) -// } - -// /// Close connection pool. -// /// -// /// # Errors -// /// May return Err Result if connection pool isn't opened. -// pub fn inner_close(&self) -> RustPSQLDriverPyResult<()> { -// let db_pool_manager = -// self.db_pool -// .as_ref() -// .ok_or(RustPSQLDriverError::DatabasePoolError( -// "Database pool is not initialized".into(), -// ))?; - -// db_pool_manager.close(); - -// Ok(()) -// } -// } - -// #[pyclass()] -// pub struct PSQLPool { -// rust_psql_pool: Arc, -// } - -// #[pymethods] -// impl PSQLPool { -// #[new] -// #[allow(clippy::too_many_arguments)] -// #[allow(clippy::missing_errors_doc)] -// pub fn new( -// dsn: Option, -// username: Option, -// password: Option, -// host: Option, -// port: Option, -// db_name: Option, -// max_db_pool_size: Option, -// conn_recycling_method: Option, -// ) -> RustPSQLDriverPyResult { -// let inner_pool = RustPSQLPool { -// dsn, -// username, -// password, -// host, -// port, -// db_name, -// max_db_pool_size, -// conn_recycling_method, -// db_pool: None, -// } -// .inner_startup()?; -// Ok(PSQLPool { -// rust_psql_pool: Arc::new(inner_pool), -// }) -// } +/// # Errors +/// May return error if cannot build new connection pool. +#[pyfunction] +#[allow(clippy::too_many_arguments)] +pub fn connect( + dsn: Option, + username: Option, + password: Option, + host: Option, + port: Option, + db_name: Option, + max_db_pool_size: Option, + conn_recycling_method: Option, +) -> RustPSQLDriverPyResult { + if let Some(max_db_pool_size) = max_db_pool_size { + if max_db_pool_size < 2 { + return Err(RustPSQLDriverError::DataBasePoolConfigurationError( + "Maximum database pool size must be more than 1".into(), + )); + } + } -// /// Return single connection. -// /// -// /// # Errors -// /// May return Err Result if `inner_connection` returns error. -// pub fn connection<'a>(&'a self, py: Python<'a>) -> RustPSQLDriverPyResult<&'a PyAny> { -// let psql_pool_arc = self.rust_psql_pool.clone(); + 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); + } -// rustdriver_future(py, async move { psql_pool_arc.inner_connection().await }) -// } + if let Some(port) = port { + pg_config.port(port); + } -// /// Execute querystring with parameters. -// /// -// /// # Errors -// /// May return Err Result if cannot convert parameters -// /// or `inner_execute` returns Err. -// pub fn execute<'a>( -// &'a self, -// py: Python<'a>, -// querystring: String, -// parameters: Option<&'a PyAny>, -// prepared: Option, -// ) -> RustPSQLDriverPyResult<&'a PyAny> { -// let db_pool_arc = self.rust_psql_pool.clone(); -// let mut params: Vec = vec![]; -// if let Some(parameters) = parameters { -// params = convert_parameters(parameters)?; -// } + if let Some(db_name) = db_name { + pg_config.dbname(&db_name); + } + } -// rustdriver_future(py, async move { -// // let db_pool_guard = db_pool_arc.read().await; -// db_pool_arc -// .inner_execute(querystring, params, prepared.unwrap_or(true)) -// .await -// }) -// } + let mgr_config: ManagerConfig; + if let Some(conn_recycling_method) = conn_recycling_method { + mgr_config = ManagerConfig { + recycling_method: conn_recycling_method.to_internal(), + } + } else { + mgr_config = ManagerConfig { + recycling_method: RecyclingMethod::Fast, + }; + } + let mgr = Manager::from_config(pg_config, NoTls, mgr_config); -// Close connection pool. -// -// # Errors -//May return Err Result if connection pool isn't opened. -// pub fn close<'a>(&self, py: Python<'a>) -> RustPSQLDriverPyResult<&'a PyAny> { -// let db_pool_arc = self.rust_psql_pool.clone(); + let mut db_pool_builder = Pool::builder(mgr); + if let Some(max_db_pool_size) = max_db_pool_size { + db_pool_builder = db_pool_builder.max_size(max_db_pool_size); + } -// rustdriver_future(py, async move { -// // let db_pool_guard = db_pool_arc.read().await; + let db_pool = db_pool_builder.build()?; -// db_pool_arc.inner_close() -// }) -// } -// } + Ok(ConnectionPool(db_pool)) +} #[pyclass] pub struct ConnectionPool(Pool); @@ -320,56 +101,16 @@ impl ConnectionPool { max_db_pool_size: Option, conn_recycling_method: Option, ) -> RustPSQLDriverPyResult { - if let Some(max_db_pool_size) = max_db_pool_size { - if max_db_pool_size < 2 { - return Err(RustPSQLDriverError::DataBasePoolConfigurationError( - "Maximum database pool size must be more than 1".into(), - )); - } - } - - 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 mgr_config: ManagerConfig; - if let Some(conn_recycling_method) = conn_recycling_method { - mgr_config = ManagerConfig { - recycling_method: conn_recycling_method.to_internal(), - } - } else { - mgr_config = ManagerConfig { - recycling_method: RecyclingMethod::Fast, - }; - } - let mgr = Manager::from_config(pg_config, NoTls, mgr_config); - - let mut db_pool_builder = Pool::builder(mgr); - if let Some(max_db_pool_size) = max_db_pool_size { - db_pool_builder = db_pool_builder.max_size(max_db_pool_size); - } - - let db_pool = db_pool_builder.build()?; - - Ok(ConnectionPool(db_pool)) + connect( + dsn, + username, + password, + host, + port, + db_name, + max_db_pool_size, + conn_recycling_method, + ) } /// Execute querystring with parameters. diff --git a/src/lib.rs b/src/lib.rs index 614d9ab4..2392f098 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,12 +10,13 @@ pub mod value_converter; use common::add_module; use exceptions::python_errors::python_exceptions_module; use extra_types::extra_types_module; -use pyo3::{pymodule, types::PyModule, Bound, PyResult, Python}; +use pyo3::{pymodule, types::PyModule, wrap_pyfunction, Bound, PyResult, Python}; #[pymodule] #[pyo3(name = "_internal")] fn psqlpy(py: Python<'_>, pymod: &Bound<'_, PyModule>) -> PyResult<()> { pymod.add_class::()?; + pymod.add_function(wrap_pyfunction!(driver::connection_pool::connect, pymod)?)?; pymod.add_class::()?; pymod.add_class::()?; pymod.add_class::()?;