Merge branch 'main' into feature/add-geo-types

Author: Mernus

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "psqlpy"
-version = "0.7.4"
+version = "0.7.5"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

Cargo.toml

@@ -13,8 +13,6 @@ You cannot set the minimum size for the connection pool, by it is 0.
 So, if you set `max_db_pool_size` to 100, pool will create new connection every time there aren't enough connections to handle the load.
 :::
 
-## Connection pool methods
-
 ### All available ConnectionPool parameters
 
 - `dsn`: Full dsn connection string.
@@ -105,9 +103,6 @@ db_pool: Final = ConnectionPool(
     db_name="postgres",
     max_db_pool_size=10,
 )
-
-async def main() -> None:
-
 ```
 
 ### Initialize Connection Pool with DSN
@@ -124,9 +119,6 @@ db_pool: Final = ConnectionPool(
     dsn="postgres://postgres:postgres@localhost:5432/postgres",
     max_db_pool_size=10,
 )
-
-async def main() -> None:
-
 ```
 
 ### Create Connection Pool with one function
@@ -143,6 +135,26 @@ db_pool: Final = connect(
 ```
 `connect` function has the same parameters as `ConnectionPool`.
 
+### Use Connection Pool as context manager
+```py
+from typing import Final
+
+from psqlpy import ConnectionPool
+
+
+async def main() -> None:
+    with ConnectionPool(
+        dsn="postgres://postgres:postgres@localhost:5432/postgres",
+        max_db_pool_size=10,
+    ) as db_pool:
+        # ConnectionPool is opened
+        await db_pool.execute("SOME_SQL")
+        # ConnectionPool is opened
+    # ConnectionPool is closed
+```
+
+## Connection pool methods
+
 ### Resize
 Resize connection pool capacity.
 

docs/components/connection_pool.md

@@ -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.
 

python/psqlpy/_internal/__init__.pyi

@@ -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")

python/tests/test_connection_pool.py

@@ -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();