Merge pull request #22 from qaspen-python/feature/add_pool_close_method

Added close method for Connection Pool

Author: chandr-andr

python/psqlpy/_internal/__init__.pyi

@@ -784,6 +784,12 @@ class PSQLPool:
 
         You must call it before start making queries.
         """
+    async def close(self: Self) -> None:
+        """Close the connection pool.
+
+        By default it will be closed automatically,
+        but you can call it manually.
+        """
     async def execute(
         self: Self,
         querystring: str,

python/tests/test_connection_pool.py

@@ -1,6 +1,7 @@
 import pytest
 
 from psqlpy import Connection, ConnRecyclingMethod, PSQLPool, QueryResult
+from psqlpy.exceptions import RustPSQLDriverPyBaseError
 
 pytestmark = pytest.mark.anyio
 
@@ -59,3 +60,19 @@ async def test_pool_conn_recycling_method(
     await pg_pool.startup()
 
     await pg_pool.execute("SELECT 1")
+
+
+async def test_close_connection_pool() -> None:
+    """Test that `close` method closes connection pool."""
+    pg_pool = PSQLPool(
+        dsn="postgres://postgres:postgres@localhost:5432/psqlpy_test",
+    )
+
+    await pg_pool.startup()
+
+    await pg_pool.execute("SELECT 1")
+
+    await pg_pool.close()
+
+    with pytest.raises(expected_exception=RustPSQLDriverPyBaseError):
+        await pg_pool.execute("SELECT 1")

src/driver/connection_pool.rs

@@ -183,6 +183,23 @@ impl RustPSQLPool {
         self.db_pool = Some(db_pool_builder.build()?);
         Ok(())
     }
+
+    /// 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()]
@@ -257,16 +274,30 @@ impl PSQLPool {
         querystring: String,
         parameters: Option<&'a PyAny>,
     ) -> RustPSQLDriverPyResult<&'a PyAny> {
-        let engine_arc = self.rust_psql_pool.clone();
+        let db_pool_arc = self.rust_psql_pool.clone();
         let mut params: Vec<PythonDTO> = vec![];
         if let Some(parameters) = parameters {
             params = convert_parameters(parameters)?;
         }
 
         rustengine_future(py, async move {
-            let engine_guard = engine_arc.read().await;
+            let db_pool_guard = db_pool_arc.read().await;
+
+            db_pool_guard.inner_execute(querystring, params).await
+        })
+    }
+
+    /// 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();
+
+        rustengine_future(py, async move {
+            let db_pool_guard = db_pool_arc.read().await;
 
-            engine_guard.inner_execute(querystring, params).await
+            db_pool_guard.inner_close()
         })
     }
 }