Merge pull request #3 from dapper91/dev

- Contrib bugs fixed

Author: dapper91

CHANGELOG.rst

@@ -2,7 +2,13 @@ Changelog
 =========
 
 
-0.1.0 (2021-03-15)
+0.1.1 (2023-03-17)
+------------------
+
+- Contrib bugs fixed
+
+
+0.1.0 (2023-03-15)
 ------------------
 
 - Initial release

examples/connection_manager.py

@@ -22,8 +22,13 @@ def create(self, endpoint: SslEndpoint, timeout: Optional[float] = None) -> Conn
         hostname, port = endpoint
 
         sock = self._ssl.wrap_socket(socket.socket(type=socket.SOCK_STREAM), server_hostname=hostname)
+
+        orig_timeout = sock.gettimeout()
         sock.settimeout(timeout)
-        sock.connect((hostname, port))
+        try:
+            sock.connect((hostname, port))
+        finally:
+            sock.settimeout(orig_timeout)
 
         return sock
 

generic_connection_pool/asyncio.py

@@ -119,7 +119,7 @@ async def connection(
             timeout: Optional[float] = None,
     ) -> AsyncGenerator[ConnectionT, None]:
         """
-        Acquires a connection form the pool.
+        Acquires a connection from the pool.
 
         :param endpoint: connection endpoint
         :param timeout: number of seconds to wait. If timeout is reached :py:class:`asyncio.TimeoutError` is raised.
@@ -134,7 +134,7 @@ async def connection(
 
     async def acquire(self, endpoint: EndpointT, timeout: Optional[float] = None) -> ConnectionT:
         """
-        Acquires a connection form the pool.
+        Acquires a connection from the pool.
 
         :param endpoint: connection endpoint
         :param timeout: number of seconds to wait. If timeout is reached :py:class:`asyncio.TimeoutError` is raised.
@@ -262,7 +262,7 @@ async def _dispose_connection(
             logger.error("connection disposal timed-out: %s", conn_info.endpoint)
             raise
         except Exception as e:
-            logger.exception("connection disposal failed: %s", e)
+            logger.error("connection disposal failed: %s", e)
             return False
 
         logger.debug("connection disposed: %s", conn_info.endpoint)

generic_connection_pool/contrib/psycopg2.py

@@ -1,12 +1,12 @@
-from typing import Mapping, Optional
+from typing import List, Mapping, Optional
 
 import psycopg2.extensions
 
 from generic_connection_pool.threding import BaseConnectionManager
 
 DbEndpoint = str
 Connection = psycopg2.extensions.connection
-DsnParameters = Mapping[str, Mapping[str, str]]
+DsnParameters = Mapping[DbEndpoint, Mapping[str, str]]
 
 
 class DbConnectionManager(BaseConnectionManager[DbEndpoint, Connection]):
@@ -28,8 +28,14 @@ def dispose(self, endpoint: DbEndpoint, conn: Connection, timeout: Optional[floa
     def check_aliveness(self, endpoint: DbEndpoint, conn: Connection, timeout: Optional[float] = None) -> bool:
         try:
             with conn.cursor() as cur:
-                cur.execute("SELECT 1;")
+                query: List[str] = []
+                if timeout is not None:
+                    query.append(f"SET statement_timeout = '{int(timeout)}s';")
+
+                query.append("SELECT 1;")
+                cur.execute(''.join(query))
                 cur.fetchone()
+
         except (psycopg2.Error, OSError):
             return False
 

generic_connection_pool/contrib/socket.py

@@ -27,8 +27,13 @@ def create(self, endpoint: TcpEndpoint, timeout: Optional[float] = None) -> sock
             raise RuntimeError("unsupported address version type: %s", addr.version)
 
         sock = socket.socket(family=family, type=socket.SOCK_STREAM)
+
+        orig_timeout = sock.gettimeout()
         sock.settimeout(timeout)
-        sock.connect((str(addr), port))
+        try:
+            sock.connect((str(addr), port))
+        finally:
+            sock.settimeout(orig_timeout)
 
         return sock
 
@@ -57,8 +62,12 @@ def create(self, endpoint: SslEndpoint, timeout: Optional[float] = None) -> SSLS
         hostname, port = endpoint
 
         sock = self._ssl.wrap_socket(socket.socket(type=socket.SOCK_STREAM), server_hostname=hostname)
+        orig_timeout = sock.gettimeout()
         sock.settimeout(timeout)
-        sock.connect((hostname, port))
+        try:
+            sock.connect((hostname, port))
+        finally:
+            sock.settimeout(orig_timeout)
 
         return sock
 

generic_connection_pool/contrib/socket_async.py

@@ -54,18 +54,22 @@ def __init__(self, ssl: Union[None, bool, SSLContext]):
 
     async def create(self, endpoint: TcpStreamEndpoint) -> TcpStream:
         hostname, port = endpoint
+        server_hostname = hostname if self._ssl is not None else None
+
         reader, writer = await asyncio.open_connection(
             hostname,
             port,
-            server_hostname=hostname,
+            server_hostname=server_hostname,
             ssl=self._ssl,
         )
 
         return reader, writer
 
     async def dispose(self, endpoint: TcpStreamEndpoint, conn: TcpStream) -> None:
         reader, writer = conn
-        writer.write_eof()
+        if writer.can_write_eof():
+            writer.write_eof()
+
         writer.close()
         await writer.wait_closed()
 

generic_connection_pool/threding.py

@@ -123,7 +123,7 @@ def __init__(
     @contextlib.contextmanager
     def connection(self, endpoint: EndpointT, timeout: Optional[float] = None) -> Generator[ConnectionT, None, None]:
         """
-        Acquires a connection form the pool.
+        Acquires a connection from the pool.
 
         :param endpoint: connection endpoint
         :param timeout: number of seconds to wait. If timeout is reached :py:class:`TimeoutError` is raised.
@@ -138,7 +138,7 @@ def connection(self, endpoint: EndpointT, timeout: Optional[float] = None) -> Ge
 
     def acquire(self, endpoint: EndpointT, timeout: Optional[float] = None) -> ConnectionT:
         """
-        Acquires a connection form the pool.
+        Acquires a connection from the pool.
 
         :param endpoint: connection endpoint
         :param timeout: number of seconds to wait. If timeout is reached :py:class:`TimeoutError` is raised.
@@ -261,7 +261,7 @@ def _dispose_connection(self, conn_info: ConnectionInfo[EndpointT, ConnectionT],
             logger.error("connection disposal timed-out: %s", conn_info.endpoint)
             raise
         except Exception as e:
-            logger.exception("connection disposal failed: %s", e)
+            logger.error("connection disposal failed: %s", e)
             return False
 
         logger.debug("connection disposed: %s", conn_info.endpoint)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "generic-connection-pool"
-version = "0.1.0"
+version = "0.1.1"
 description = "generic connection pool"
 authors = ["Dmitry Pershin <[email protected]>"]
 license = "Unlicense"

tests/conftest.py

@@ -1,8 +1,19 @@
 import random
+from pathlib import Path
 
 import pytest
 
 
+@pytest.fixture
+def test_dir() -> Path:
+    return Path(__file__).parent
+
+
+@pytest.fixture
+def resource_dir(test_dir) -> Path:
+    return test_dir / 'resources'
+
+
 @pytest.fixture(autouse=True)
 def init_random():
     random.seed(0)

tests/contrib/test_async_socket_manager.py

@@ -0,0 +1,81 @@
+import asyncio
+import ssl
+from ipaddress import IPv4Address
+from pathlib import Path
+from typing import Tuple
+
+import pytest
+import pytest_asyncio.plugin
+
+from generic_connection_pool.asyncio import ConnectionPool
+from generic_connection_pool.contrib.socket_async import TcpSocketConnectionManager, TcpStreamConnectionManager
+
+
+@pytest.fixture
+async def tcp_server(request: pytest_asyncio.plugin.SubRequest, resource_dir: Path):
+    params = getattr(request, 'param', {})
+    user_ssl = params.get('use_ssl', False)
+
+    if user_ssl:
+        context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
+        context.load_cert_chain(resource_dir / 'ssl.cert', resource_dir / 'ssl.key')
+    else:
+        context = None
+
+    hostname, addr, port = 'localhost', IPv4Address('127.0.0.1'),  10000
+
+    async def echo(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
+        data = await reader.read(1024)
+        writer.write(data)
+        await writer.drain()
+        writer.close()
+        await writer.wait_closed()
+
+    async with (server := await asyncio.start_server(echo, host=hostname, port=port, ssl=context, reuse_port=True)):
+        server_task = asyncio.create_task(server.serve_forever())
+        yield hostname, addr, port
+        server_task.cancel()
+        try:
+            await server_task
+        except asyncio.CancelledError:
+            pass
+
+
+async def test_tcp_socket_manager(tcp_server: Tuple[str, IPv4Address, int]):
+    loop = asyncio.get_running_loop()
+    hostname, addr, port = tcp_server
+
+    pool = ConnectionPool(TcpSocketConnectionManager())
+    async with pool.connection((addr, port)) as sock:
+        request = b'test'
+        await loop.sock_sendall(sock, request)
+        response = await loop.sock_recv(sock, len(request))
+        assert response == request
+
+    await pool.close()
+
+
+@pytest.mark.parametrize(
+    'use_ssl, tcp_server', [
+        (True, {'use_ssl': True}),
+        (False, {'use_ssl': False}),
+    ],
+    indirect=['tcp_server'],
+)
+async def test_ssl_stream_manager(resource_dir: Path, use_ssl: bool, tcp_server: Tuple[str, IPv4Address, int]):
+    hostname, addr, port = tcp_server
+
+    if use_ssl:
+        ssl_context = ssl.create_default_context(cafile=resource_dir / 'ssl.cert')
+    else:
+        ssl_context = None
+
+    pool = ConnectionPool(TcpStreamConnectionManager(ssl_context))
+    async with pool.connection((hostname, port)) as (reader, writer):
+        request = b'test'
+        writer.write(request)
+        await writer.drain()
+        response = await reader.read()
+        assert response == request
+
+    await pool.close()