Added context manager support to ConnectionPool

chandr-andr · chandr-andr · commit 4008c0dd9869 · 2024-08-09T16:59:48.000+02:00
Signed-off-by: chandr-andr (Kiselev Aleksandr) &lt;chandr@chandr.net&gt;
diff --git a/python/psqlpy/_internal/__init__.pyi b/python/psqlpy/_internal/__init__.pyi
@@ -1215,6 +1215,14 @@ class ConnectionPool:
         - `ca_file`: Loads trusted root certificates from a file.
             The file should contain a sequence of PEM-formatted CA certificates.
         """
+    def __iter__(self: Self) -> Self: ...
+    def __enter__(self: Self) -> Self: ...
+    def __exit__(
+        self: Self,
+        exception_type: type[BaseException] | None,
+        exception: BaseException | None,
+        traceback: types.TracebackType | None,
+    ) -> None: ...
     def status(self: Self) -> ConnectionPoolStatus:
         """Return information about connection pool.
 
diff --git a/python/tests/test_connection_pool.py b/python/tests/test_connection_pool.py
@@ -175,3 +175,15 @@ async def test_close_connection_pool() -> None:
 
     with pytest.raises(expected_exception=RustPSQLDriverPyBaseError):
         await pg_pool.execute("SELECT 1")
+
+
+async def test_connection_pool_as_context_manager() -> None:
+    """Test connection pool as context manager."""
+    with ConnectionPool(
+        dsn="postgres://postgres:postgres@localhost:5432/psqlpy_test",
+    ) as pg_pool:
+        res = await pg_pool.execute("SELECT 1")
+        assert res.result()
+
+    with pytest.raises(expected_exception=RustPSQLDriverPyBaseError):
+        await pg_pool.execute("SELECT 1")
diff --git a/src/driver/connection_pool.rs b/src/driver/connection_pool.rs
@@ -2,7 +2,7 @@ use crate::runtime::tokio_runtime;
 use deadpool_postgres::{Manager, ManagerConfig, Object, Pool, RecyclingMethod};
 use openssl::ssl::{SslConnector, SslMethod, SslVerifyMode};
 use postgres_openssl::MakeTlsConnector;
-use pyo3::{pyclass, pyfunction, pymethods, PyAny};
+use pyo3::{pyclass, pyfunction, pymethods, Py, PyAny};
 use std::{sync::Arc, vec};
 use tokio_postgres::NoTls;
 
@@ -253,6 +253,28 @@ impl ConnectionPool {
         )
     }
 
+    #[must_use]
+    pub fn __iter__(self_: Py<Self>) -> Py<Self> {
+        self_
+    }
+
+    #[allow(clippy::needless_pass_by_value)]
+    fn __enter__(self_: Py<Self>) -> Py<Self> {
+        self_
+    }
+
+    #[allow(clippy::needless_pass_by_value)]
+    fn __exit__(
+        self_: Py<Self>,
+        _exception_type: Py<PyAny>,
+        _exception: Py<PyAny>,
+        _traceback: Py<PyAny>,
+    ) {
+        pyo3::Python::with_gil(|gil| {
+            self_.borrow(gil).close();
+        });
+    }
+
     #[must_use]
     pub fn status(&self) -> ConnectionPoolStatus {
         let inner_status = self.0.status();
