Skip to content

Commit 536e98c

Browse files
committed
Solved conficts
2 parents 64281e9 + 1af0139 commit 536e98c

File tree

3 files changed

+57
-4
lines changed

3 files changed

+57
-4
lines changed

python/psqlpy/_internal/__init__.pyi

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -809,6 +809,12 @@ class PSQLPool:
809809
810810
You must call it before start making queries.
811811
"""
812+
async def close(self: Self) -> None:
813+
"""Close the connection pool.
814+
815+
By default it will be closed automatically,
816+
but you can call it manually.
817+
"""
812818
async def execute(
813819
self: Self,
814820
querystring: str,

python/tests/test_connection_pool.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import pytest
22

33
from psqlpy import Connection, ConnRecyclingMethod, PSQLPool, QueryResult
4+
from psqlpy.exceptions import RustPSQLDriverPyBaseError
45

56
pytestmark = pytest.mark.anyio
67

@@ -59,3 +60,19 @@ async def test_pool_conn_recycling_method(
5960
await pg_pool.startup()
6061

6162
await pg_pool.execute("SELECT 1")
63+
64+
65+
async def test_close_connection_pool() -> None:
66+
"""Test that `close` method closes connection pool."""
67+
pg_pool = PSQLPool(
68+
dsn="postgres://postgres:postgres@localhost:5432/psqlpy_test",
69+
)
70+
71+
await pg_pool.startup()
72+
73+
await pg_pool.execute("SELECT 1")
74+
75+
await pg_pool.close()
76+
77+
with pytest.raises(expected_exception=RustPSQLDriverPyBaseError):
78+
await pg_pool.execute("SELECT 1")

src/driver/connection_pool.rs

Lines changed: 34 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -190,6 +190,23 @@ impl RustPSQLPool {
190190
self.db_pool = Some(db_pool_builder.build()?);
191191
Ok(())
192192
}
193+
194+
/// Close connection pool.
195+
///
196+
/// # Errors
197+
/// May return Err Result if connection pool isn't opened.
198+
pub fn inner_close(&self) -> RustPSQLDriverPyResult<()> {
199+
let db_pool_manager =
200+
self.db_pool
201+
.as_ref()
202+
.ok_or(RustPSQLDriverError::DatabasePoolError(
203+
"Database pool is not initialized".into(),
204+
))?;
205+
206+
db_pool_manager.close();
207+
208+
Ok(())
209+
}
193210
}
194211

195212
#[pyclass()]
@@ -265,18 +282,31 @@ impl PSQLPool {
265282
parameters: Option<&'a PyAny>,
266283
prepared: Option<bool>,
267284
) -> RustPSQLDriverPyResult<&'a PyAny> {
268-
let engine_arc = self.rust_psql_pool.clone();
285+
let db_pool_arc = self.rust_psql_pool.clone();
269286
let mut params: Vec<PythonDTO> = vec![];
270287
if let Some(parameters) = parameters {
271288
params = convert_parameters(parameters)?;
272289
}
273290

274291
rustengine_future(py, async move {
275-
let engine_guard = engine_arc.read().await;
276-
277-
engine_guard
292+
let db_pool_guard = db_pool_arc.read().await;
293+
db_pool_guard
278294
.inner_execute(querystring, params, prepared.unwrap_or(true))
279295
.await
280296
})
281297
}
298+
299+
/// Close connection pool.
300+
///
301+
/// # Errors
302+
/// May return Err Result if connection pool isn't opened.
303+
pub fn close<'a>(&self, py: Python<'a>) -> RustPSQLDriverPyResult<&'a PyAny> {
304+
let db_pool_arc = self.rust_psql_pool.clone();
305+
306+
rustengine_future(py, async move {
307+
let db_pool_guard = db_pool_arc.read().await;
308+
309+
db_pool_guard.inner_close()
310+
})
311+
}
282312
}

0 commit comments

Comments
 (0)