Merge pull request #58 from qaspen-python/feature/better_connection_managment

Better connection management

Author: chandr-andr

docs/components/connection.md

@@ -5,20 +5,37 @@ title: Connection
 `Connection` object represents single connection to the `PostgreSQL`. You must work with database within it.
 `Connection` get be made with `ConnectionPool().connection()` method.
 
+## Usage
+::: tabs
+
+@tab default
 ```python
 from psqlpy import ConnectionPool
 
-
 db_pool: Final = ConnectionPool(
     dsn="postgres://postgres:postgres@localhost:5432/postgres",
 )
 
-
 async def main() -> None:
-    ...
     connection = await db_pool.connection()
 ```
 
+@tab async context manager
+```python
+from psqlpy import ConnectionPool
+
+db_pool: Final = ConnectionPool(
+    dsn="postgres://postgres:postgres@localhost:5432/postgres",
+)
+
+async def main() -> None:
+    async with db_pool.acquire() as connection:
+        # connection is valid here
+        ...
+    # connection is invalid here
+```
+:::
+
 ## Connection methods
 
 ### Execute
@@ -158,3 +175,19 @@ async def main() -> None:
         deferrable=True,
     )
 ```
+
+### Back To Pool
+Returns connection to the pool.
+It's crucial to commit all transactions and close all cursor which are made from the connection.
+Otherwise, this method won't do anything useful.
+
+::: tip
+There is no need in this method if you use async context manager.
+:::
+
+```python
+async def main() -> None:
+    ...
+    connection = await db_pool.connection()
+    connection.back_to_pool()
+```

docs/components/connection_pool.md

@@ -141,6 +141,29 @@ db_pool: Final = connect(
 ```
 `connect` function has the same parameters as `ConnectionPool`.
 
+### Resize
+Resize connection pool capacity.
+
+This change the max_size of the pool dropping excess objects and/or making space for new ones.
+
+#### Parameters:
+- `new_max_size`: new size of the pool.
+
+```python
+async def main() -> None:
+    ...
+    db_pool.resize(15)
+```
+
+### Status
+Retrieve status of the connection pool.
+
+It has 4 parameters:
+- `max_size` - maximum possible size of the connection pool.
+- `size` - current size of the connection pool.
+- `available` - available connection in the connection pool.
+- `waiting` - waiting requests to retrieve connection from connection pool.
+
 ### Execute
 
 #### Parameters:
@@ -190,6 +213,18 @@ async def main() -> None:
     dict_results: list[dict[str, Any]] = results.result()
 ```
 
+### Acquire
+
+Get single connection for async context manager.
+Must be used only in async context manager.
+
+```python
+async def main() -> None:
+    ...
+    async with db_pool.acquire() as connection:
+        ...
+```
+
 ### Connection
 
 To get single connection from the `ConnectionPool` there is method named `connection()`.

docs/components/exceptions.md

@@ -28,18 +28,21 @@ stateDiagram-v2
     }
     state BaseConnectionError {
         [*] --> ConnectionExecuteError
+        [*] --> ConnectionClosedError
     }
     state BaseTransactionError {
         [*] --> TransactionBeginError
         [*] --> TransactionCommitError
         [*] --> TransactionRollbackError
         [*] --> TransactionSavepointError
         [*] --> TransactionExecuteError
+        [*] --> TransactionClosedError
     }
     state BaseCursorError {
         [*] --> CursorStartError
         [*] --> CursorCloseError
         [*] --> CursorFetchError
+        [*] --> CursorClosedError
     }
     state RustException {
         [*] --> DriverError
@@ -86,6 +89,9 @@ Base error for Connection errors.
 #### ConnectionExecuteError
 Error in connection execution.
 
+#### ConnectionClosedError
+Error if underlying connection is closed.
+
 ### BaseTransactionError
 Base error for all transaction errors.
 
@@ -104,6 +110,9 @@ Error in transaction savepoint.
 #### TransactionExecuteError
 Error in transaction execution.
 
+#### TransactionClosedError
+Error if underlying connection is closed.
+
 ### BaseCursorError
 Base error for Cursor errors.
 
@@ -115,3 +124,6 @@ Error in cursor close.
 
 #### CursorFetchError
 Error in cursor fetch (any fetch).
+
+#### CursorClosedError
+Error if underlying connection is closed.

python/psqlpy/_internal/__init__.pyi

@@ -816,6 +816,13 @@ class Connection:
     It can be created only from connection pool.
     """
 
+    async def __aenter__(self: Self) -> Self: ...
+    async def __aexit__(
+        self: Self,
+        exception_type: type[BaseException] | None,
+        exception: BaseException | None,
+        traceback: types.TracebackType | None,
+    ) -> None: ...
     async def execute(
         self: Self,
         querystring: str,
@@ -1040,6 +1047,18 @@ class Connection:
                         ... # do something with this result.
         ```
         """
+    def back_to_pool(self: Self) -> None:
+        """Return connection back to the pool.
+
+        It necessary to commit all transactions and close all cursor
+        made by this connection. Otherwise, it won't have any practical usage.
+        """
+
+class ConnectionPoolStatus:
+    max_size: int
+    size: int
+    available: int
+    waiting: int
 
 class ConnectionPool:
     """Connection pool for executing queries.
@@ -1142,6 +1161,21 @@ class ConnectionPool:
         - `max_db_pool_size`: maximum size of the connection pool.
         - `conn_recycling_method`: how a connection is recycled.
         """
+    def status(self: Self) -> ConnectionPoolStatus:
+        """Return information about connection pool.
+
+        ### Returns
+        `ConnectionPoolStatus`
+        """
+    def resize(self: Self, new_max_size: int) -> None:
+        """Resize the connection pool.
+
+        This change the max_size of the pool dropping
+        excess objects and/or making space for new ones.
+
+        ### Parameters:
+        - `new_max_size`: new size for the connection pool.
+        """
     async def execute(
         self: Self,
         querystring: str,
@@ -1202,6 +1236,24 @@ class ConnectionPool:
 
         It acquires new connection from the database pool.
         """
+    def acquire(self: Self) -> Connection:
+        """Create new connection for async context manager.
+
+        Must be used only in async context manager.
+
+        ### Example:
+        ```python
+        import asyncio
+
+        from psqlpy import PSQLPool, QueryResult
+
+
+        async def main() -> None:
+            db_pool = PSQLPool()
+            async with db_pool.acquire() as connection:
+                res = await connection.execute(...)
+        ```
+        """
     def close(self: Self) -> None:
         """Close the connection pool."""
 

python/psqlpy/_internal/exceptions.pyi

@@ -19,6 +19,9 @@ class BaseConnectionError(RustPSQLDriverPyBaseError):
 class ConnectionExecuteError(BaseConnectionError):
     """Error in connection execution."""
 
+class ConnectionClosedError(BaseConnectionError):
+    """Error if underlying connection is already closed."""
+
 class BaseTransactionError(RustPSQLDriverPyBaseError):
     """Base error for all transaction errors."""
 
@@ -37,6 +40,9 @@ class TransactionSavepointError(BaseTransactionError):
 class TransactionExecuteError(BaseTransactionError):
     """Error in transaction execution."""
 
+class TransactionClosedError(BaseTransactionError):
+    """Error if underlying connection is already closed."""
+
 class BaseCursorError(RustPSQLDriverPyBaseError):
     """Base error for Cursor errors."""
 
@@ -49,6 +55,9 @@ class CursorCloseError(BaseCursorError):
 class CursorFetchError(BaseCursorError):
     """Error in cursor fetch (any fetch)."""
 
+class CursorClosedError(BaseCursorError):
+    """Error if underlying connection is already closed."""
+
 class UUIDValueConvertError(RustPSQLDriverPyBaseError):
     """Error if it's impossible to convert py string UUID into rust UUID."""
 

python/psqlpy/exceptions.py

@@ -3,10 +3,12 @@
     BaseConnectionPoolError,
     BaseCursorError,
     BaseTransactionError,
+    ConnectionClosedError,
     ConnectionExecuteError,
     ConnectionPoolBuildError,
     ConnectionPoolConfigurationError,
     ConnectionPoolExecuteError,
+    CursorClosedError,
     CursorCloseError,
     CursorFetchError,
     CursorStartError,
@@ -15,6 +17,7 @@
     RustPSQLDriverPyBaseError,
     RustToPyValueMappingError,
     TransactionBeginError,
+    TransactionClosedError,
     TransactionCommitError,
     TransactionExecuteError,
     TransactionRollbackError,
@@ -29,16 +32,19 @@
     "ConnectionPoolExecuteError",
     "BaseConnectionError",
     "ConnectionExecuteError",
+    "ConnectionClosedError",
     "BaseTransactionError",
     "TransactionBeginError",
     "TransactionCommitError",
     "TransactionRollbackError",
     "TransactionSavepointError",
     "TransactionExecuteError",
+    "TransactionClosedError",
     "BaseCursorError",
     "CursorStartError",
     "CursorCloseError",
     "CursorFetchError",
+    "CursorClosedError",
     "RustPSQLDriverPyBaseError",
     "RustToPyValueMappingError",
     "PyToRustValueMappingError",

python/tests/test_connection.py

@@ -6,7 +6,11 @@
 from tests.helpers import count_rows_in_test_table
 
 from psqlpy import ConnectionPool, Cursor, QueryResult, Transaction
-from psqlpy.exceptions import ConnectionExecuteError, TransactionExecuteError
+from psqlpy.exceptions import (
+    ConnectionClosedError,
+    ConnectionExecuteError,
+    TransactionExecuteError,
+)
 
 pytestmark = pytest.mark.anyio
 
@@ -147,3 +151,32 @@ async def test_connection_cursor(
             all_results.extend(cur_res.result())
 
     assert len(all_results) == number_database_records
+
+
+async def test_connection_async_context_manager(
+    psql_pool: ConnectionPool,
+    table_name: str,
+    number_database_records: int,
+) -> None:
+    """Test connection as a async context manager."""
+    async with psql_pool.acquire() as connection:
+        conn_result = await connection.execute(
+            querystring=f"SELECT * FROM {table_name}",
+        )
+        assert not psql_pool.status().available
+
+    assert psql_pool.status().available == 1
+
+    assert isinstance(conn_result, QueryResult)
+    assert len(conn_result.result()) == number_database_records
+
+
+async def test_closed_connection_error(
+    psql_pool: ConnectionPool,
+) -> None:
+    """Test exception when connection is closed."""
+    connection = await psql_pool.connection()
+    connection.back_to_pool()
+
+    with pytest.raises(expected_exception=ConnectionClosedError):
+        await connection.execute("SELECT 1")

python/tests/test_cursor.py

@@ -170,3 +170,37 @@ async def test_cursor_as_async_manager(
             all_results.append(result)
 
     assert len(all_results) == expected_num_results
+
+
+async def test_cursor_send_underlying_connection_to_pool(
+    psql_pool: ConnectionPool,
+    table_name: str,
+) -> None:
+    """Test send underlying connection to the pool."""
+    async with psql_pool.acquire() as connection:
+        async with connection.transaction() as transaction:
+            async with transaction.cursor(
+                querystring=f"SELECT * FROM {table_name}",
+            ) as cursor:
+                await cursor.fetch(10)
+                assert not psql_pool.status().available
+            assert not psql_pool.status().available
+        assert not psql_pool.status().available
+    assert psql_pool.status().available == 1
+
+
+async def test_cursor_send_underlying_connection_to_pool_manually(
+    psql_pool: ConnectionPool,
+    table_name: str,
+) -> None:
+    """Test send underlying connection to the pool."""
+    async with psql_pool.acquire() as connection:
+        async with connection.transaction() as transaction:
+            cursor = transaction.cursor(querystring=f"SELECT * FROM {table_name}")
+            await cursor.start()
+            await cursor.fetch(10)
+            assert not psql_pool.status().available
+            await cursor.close()
+            assert not psql_pool.status().available
+        assert not psql_pool.status().available
+    assert psql_pool.status().available == 1