diff --git a/README.md b/README.md index 4c203b7a..587a2a86 100644 --- a/README.md +++ b/README.md @@ -80,6 +80,31 @@ async def main() -> None: # rust does it instead. ``` +### Control connection recycling +There are 3 available options to control how a connection is recycled - `Fast`, `Verified` and `Clean`. +As connection can be closed in different situations on various sides you can select preferable behavior of how a connection is recycled. + +- `Fast`: Only run `is_closed()` when recycling existing connections. +- `Verified`: Run `is_closed()` and execute a test query. This is slower, but guarantees that the database connection is ready to + be used. Normally, `is_closed()` should be enough to filter + out bad connections, but under some circumstances (i.e. hard-closed + network connections) it's possible that `is_closed()` + returns `false` while the connection is dead. You will receive an error + on your first query then. +- `Clean`: Like [`Verified`] query method, but instead use the following sequence of statements which guarantees a pristine connection: + ```sql + CLOSE ALL; + SET SESSION AUTHORIZATION DEFAULT; + RESET ALL; + UNLISTEN *; + SELECT pg_advisory_unlock_all(); + DISCARD TEMP; + DISCARD SEQUENCES; + ``` + This is similar to calling `DISCARD ALL`. but doesn't call + `DEALLOCATE ALL` and `DISCARD PLAN`, so that the statement cache is not + rendered ineffective. + ## Query parameters You can pass parameters into queries. Parameters can be passed in any `execute` method as the second parameter, it must be a list. diff --git a/python/psqlpy/__init__.py b/python/psqlpy/__init__.py index 737eeb49..42477791 100644 --- a/python/psqlpy/__init__.py +++ b/python/psqlpy/__init__.py @@ -1,5 +1,6 @@ from ._internal import ( Connection, + ConnRecyclingMethod, Cursor, IsolationLevel, PSQLPool, @@ -16,4 +17,5 @@ "ReadVariant", "Connection", "Cursor", + "ConnRecyclingMethod", ] diff --git a/python/psqlpy/_internal/__init__.pyi b/python/psqlpy/_internal/__init__.pyi index b2721168..4e5254a0 100644 --- a/python/psqlpy/_internal/__init__.pyi +++ b/python/psqlpy/_internal/__init__.pyi @@ -24,6 +24,49 @@ class ReadVariant(Enum): ReadOnly = 1 ReadWrite = 2 +class ConnRecyclingMethod(Enum): + """Possible methods of how a connection is recycled. + + The default is [`Fast`] which does not check the connection health or + perform any clean-up queries. + + # Description: + ## Fast: + Only run [`is_closed()`] when recycling existing connections. + + Unless you have special needs this is a safe choice. + + ## Verified: + Run [`is_closed()`] and execute a test query. + + This is slower, but guarantees that the database connection is ready to + be used. Normally, [`is_closed()`] should be enough to filter + out bad connections, but under some circumstances (i.e. hard-closed + network connections) it's possible that [`is_closed()`] + returns `false` while the connection is dead. You will receive an error + on your first query then. + + ## Clean: + Like [`Verified`] query method, but instead use the following sequence + of statements which guarantees a pristine connection: + ```sql + CLOSE ALL; + SET SESSION AUTHORIZATION DEFAULT; + RESET ALL; + UNLISTEN *; + SELECT pg_advisory_unlock_all(); + DISCARD TEMP; + DISCARD SEQUENCES; + ``` + This is similar to calling `DISCARD ALL`. but doesn't call + `DEALLOCATE ALL` and `DISCARD PLAN`, so that the statement cache is not + rendered ineffective. + """ + + Fast = 1 + Verified = 2 + Clean = 3 + class Cursor: """Represent opened cursor in a transaction. @@ -446,6 +489,7 @@ class PSQLPool: port: Optional[int] = None, db_name: Optional[str] = None, max_db_pool_size: Optional[str] = None, + conn_recycling_method: Optional[ConnRecyclingMethod] = None, ) -> None: """Create new PostgreSQL connection pool. @@ -468,6 +512,7 @@ class PSQLPool: - `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. """ async def startup(self: Self) -> None: """Startup the connection pool. diff --git a/python/tests/test_connection_pool.py b/python/tests/test_connection_pool.py index 5ea6ffdb..70613899 100644 --- a/python/tests/test_connection_pool.py +++ b/python/tests/test_connection_pool.py @@ -1,6 +1,6 @@ import pytest -from psqlpy import Connection, PSQLPool, QueryResult +from psqlpy import Connection, ConnRecyclingMethod, PSQLPool, QueryResult @pytest.mark.anyio @@ -39,3 +39,25 @@ async def test_pool_connection( """Test that PSQLPool can return single connection from the pool.""" connection = await psql_pool.connection() assert isinstance(connection, Connection) + + +@pytest.mark.anyio +@pytest.mark.parametrize( + "conn_recycling_method", + [ + ConnRecyclingMethod.Fast, + ConnRecyclingMethod.Verified, + ConnRecyclingMethod.Clean, + ], +) +async def test_pool_conn_recycling_method( + conn_recycling_method: ConnRecyclingMethod, +) -> None: + pg_pool = PSQLPool( + dsn="postgres://postgres:postgres@localhost:5432/psqlpy_test", + conn_recycling_method=conn_recycling_method, + ) + + await pg_pool.startup() + + await pg_pool.execute("SELECT 1") diff --git a/src/driver/common_options.rs b/src/driver/common_options.rs new file mode 100644 index 00000000..cd8cb362 --- /dev/null +++ b/src/driver/common_options.rs @@ -0,0 +1,21 @@ +use deadpool_postgres::RecyclingMethod; +use pyo3::pyclass; + +#[pyclass] +#[derive(Clone, Copy)] +pub enum ConnRecyclingMethod { + Fast, + Verified, + Clean, +} + +impl ConnRecyclingMethod { + #[must_use] + pub fn to_internal(&self) -> RecyclingMethod { + match self { + ConnRecyclingMethod::Fast => RecyclingMethod::Fast, + ConnRecyclingMethod::Verified => RecyclingMethod::Verified, + ConnRecyclingMethod::Clean => RecyclingMethod::Clean, + } + } +} diff --git a/src/driver/connection_pool.rs b/src/driver/connection_pool.rs index b5b481c5..cba6fac6 100644 --- a/src/driver/connection_pool.rs +++ b/src/driver/connection_pool.rs @@ -10,7 +10,7 @@ use crate::{ value_converter::{convert_parameters, PythonDTO}, }; -use super::connection::Connection; +use super::{common_options::ConnRecyclingMethod, connection::Connection}; /// `PSQLPool` for internal use only. /// @@ -23,12 +23,14 @@ pub struct RustPSQLPool { port: Option, db_name: Option, max_db_pool_size: Option, + conn_recycling_method: Option, db_pool: Arc>>, } impl RustPSQLPool { /// Create new `RustPSQLPool`. #[must_use] + #[allow(clippy::too_many_arguments)] pub fn new( dsn: Option, username: Option, @@ -37,6 +39,7 @@ impl RustPSQLPool { port: Option, db_name: Option, max_db_pool_size: Option, + conn_recycling_method: Option, ) -> Self { RustPSQLPool { dsn, @@ -46,6 +49,7 @@ impl RustPSQLPool { port, db_name, max_db_pool_size, + conn_recycling_method, db_pool: Arc::new(tokio::sync::RwLock::new(None)), } } @@ -124,6 +128,7 @@ impl RustPSQLPool { 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; let mut db_pool_guard = db_pool_arc.write().await; @@ -163,9 +168,16 @@ impl RustPSQLPool { } } - let mgr_config = ManagerConfig { - recycling_method: RecyclingMethod::Fast, - }; + 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); @@ -186,6 +198,7 @@ pub struct PSQLPool { #[pymethods] impl PSQLPool { #[new] + #[allow(clippy::too_many_arguments)] #[must_use] pub fn new( dsn: Option, @@ -195,6 +208,7 @@ impl PSQLPool { port: Option, db_name: Option, max_db_pool_size: Option, + conn_recycling_method: Option, ) -> Self { PSQLPool { rust_psql_pool: Arc::new(tokio::sync::RwLock::new(RustPSQLPool { @@ -205,6 +219,7 @@ impl PSQLPool { port, db_name, max_db_pool_size, + conn_recycling_method, db_pool: Arc::new(tokio::sync::RwLock::new(None)), })), } diff --git a/src/driver/mod.rs b/src/driver/mod.rs index 6b0b1a7f..aec33d5b 100644 --- a/src/driver/mod.rs +++ b/src/driver/mod.rs @@ -1,3 +1,4 @@ +pub mod common_options; pub mod connection; pub mod connection_pool; pub mod cursor; diff --git a/src/lib.rs b/src/lib.rs index 7008e375..5db06244 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -19,6 +19,7 @@ fn psqlpy(py: Python<'_>, pymod: &PyModule) -> PyResult<()> { pymod.add_class::()?; pymod.add_class::()?; pymod.add_class::()?; + pymod.add_class::()?; pymod.add_class::()?; add_module(py, pymod, "extra_types", extra_types_module)?; add_module(py, pymod, "exceptions", python_exceptions_module)?;