[BTS-1515] Pass PoolTimeout as argument (#257)

apetenchea · web-flow · commit 350c36b00308 · 2023-07-14T11:01:25.000+03:00
* Adding extra parameters to the DefaultHTTPClient

* Keeping default behavior the same as before

* Corrections

* Updating test

* Adding request_timeout together with pool_timeout
diff --git a/arango/client.py b/arango/client.py
@@ -13,7 +13,7 @@
 )
 from arango.database import StandardDatabase
 from arango.exceptions import ServerConnectionError
-from arango.http import DefaultHTTPClient, HTTPClient
+from arango.http import DEFAULT_REQUEST_TIMEOUT, DefaultHTTPClient, HTTPClient
 from arango.resolver import (
     HostResolver,
     RandomHostResolver,
@@ -55,10 +55,10 @@ class ArangoClient:
     :type verify_override: Union[bool, str, None]
     :param request_timeout: This is the default request timeout (in seconds)
        for http requests issued by the client if the parameter http_client is
-       not secified. The default value is 60.
+       not specified. The default value is 60.
        None: No timeout.
        int: Timeout value in seconds.
-    :type request_timeout: Any
+    :type request_timeout: int | float
     """
 
     def __init__(
@@ -70,7 +70,7 @@ def __init__(
         serializer: Callable[..., str] = lambda x: dumps(x),
         deserializer: Callable[[str], Any] = lambda x: loads(x),
         verify_override: Union[bool, str, None] = None,
-        request_timeout: Any = 60,
+        request_timeout: Union[int, float] = DEFAULT_REQUEST_TIMEOUT,
     ) -> None:
         if isinstance(hosts, str):
             self._hosts = [host.strip("/") for host in hosts.split(",")]
@@ -88,11 +88,7 @@ def __init__(
             self._host_resolver = RoundRobinHostResolver(host_count, resolver_max_tries)
 
         # Initializes the http client
-        self._http = http_client or DefaultHTTPClient()
-        # Sets the request timeout.
-        # This call can only happen AFTER initializing the http client.
-        if http_client is None:
-            self.request_timeout = request_timeout
+        self._http = http_client or DefaultHTTPClient(request_timeout=request_timeout)
 
         self._serializer = serializer
         self._deserializer = deserializer
@@ -137,12 +133,12 @@ def request_timeout(self) -> Any:
         :return: Request timeout.
         :rtype: Any
         """
-        return self._http.REQUEST_TIMEOUT  # type: ignore
+        return self._http.request_timeout  # type: ignore
 
     # Setter for request_timeout
     @request_timeout.setter
     def request_timeout(self, value: Any) -> None:
-        self._http.REQUEST_TIMEOUT = value  # type: ignore
+        self._http.request_timeout = value  # type: ignore
 
     def db(
         self,
diff --git a/arango/http.py b/arango/http.py
@@ -1,16 +1,20 @@
-__all__ = ["HTTPClient", "DefaultHTTPClient"]
+__all__ = ["HTTPClient", "DefaultHTTPClient", "DEFAULT_REQUEST_TIMEOUT"]
 
+import typing
 from abc import ABC, abstractmethod
-from typing import MutableMapping, Optional, Tuple, Union
+from typing import Any, MutableMapping, Optional, Tuple, Union
 
 from requests import Session
-from requests.adapters import HTTPAdapter
+from requests.adapters import DEFAULT_POOLBLOCK, DEFAULT_POOLSIZE, HTTPAdapter
 from requests_toolbelt import MultipartEncoder
+from urllib3.poolmanager import PoolManager
 from urllib3.util.retry import Retry
 
 from arango.response import Response
 from arango.typings import Headers
 
+DEFAULT_REQUEST_TIMEOUT = 60
+
 
 class HTTPClient(ABC):  # pragma: no cover
     """Abstract base class for HTTP clients."""
@@ -63,12 +67,92 @@ def send_request(
         raise NotImplementedError
 
 
-class DefaultHTTPClient(HTTPClient):
-    """Default HTTP client implementation."""
+class DefaultHTTPAdapter(HTTPAdapter):
+    """Default transport adapter implementation
+
+    :param connection_timeout: Socket timeout in seconds for each individual connection.
+    :type connection_timeout: int | float
+    :param pool_connections: The number of urllib3 connection pools to cache.
+    :type pool_connections: int
+    :param pool_maxsize: The maximum number of connections to save in the pool.
+    :type pool_maxsize: int
+    :param pool_timeout: If set, then the pool will be set to block=True,
+        and requests will block for pool_timeout seconds and raise
+        EmptyPoolError if no connection is available within the time period.
+    :type pool_timeout: int | float | None
+    :param kwargs: Additional keyword arguments passed to the HTTPAdapter constructor.
+    :type kwargs: Any
+    """
+
+    def __init__(
+        self,
+        connection_timeout: Union[int, float] = DEFAULT_REQUEST_TIMEOUT,
+        pool_connections: int = DEFAULT_POOLSIZE,
+        pool_maxsize: int = DEFAULT_POOLSIZE,
+        pool_timeout: Union[int, float, None] = None,
+        **kwargs: Any
+    ) -> None:
+        self._connection_timeout = connection_timeout
+        self._pool_timeout = pool_timeout
+        super().__init__(
+            pool_connections=pool_connections, pool_maxsize=pool_maxsize, **kwargs
+        )
+
+    @typing.no_type_check
+    def init_poolmanager(
+        self, connections, maxsize, block=DEFAULT_POOLBLOCK, **pool_kwargs
+    ) -> None:
+        kwargs = pool_kwargs
+        kwargs.update(
+            dict(
+                num_pools=connections,
+                maxsize=maxsize,
+                strict=True,
+                timeout=self._connection_timeout,
+            )
+        )
+        if self._pool_timeout is not None:
+            kwargs["block"] = True
+            kwargs["timeout"] = self._pool_timeout
+        else:
+            kwargs["block"] = False
+        self.poolmanager = PoolManager(**kwargs)
 
-    REQUEST_TIMEOUT = 60
-    RETRY_ATTEMPTS = 3
-    BACKOFF_FACTOR = 1
+
+class DefaultHTTPClient(HTTPClient):
+    """Default HTTP client implementation.
+
+    :param request_timeout: Timeout in seconds for each individual connection.
+    :type request_timeout: int | float
+    :param retry_attempts: Number of retry attempts.
+    :type retry_attempts: int
+    :param backoff_factor: Backoff factor for retry attempts.
+    :type backoff_factor: float
+    :param pool_connections: The number of urllib3 connection pools to cache.
+    :type pool_connections: int
+    :param pool_maxsize: The maximum number of connections to save in the pool.
+    :type pool_maxsize: int
+    :param pool_timeout: If set, then the pool will be set to block=True,
+        and requests will block for pool_timeout seconds and raise
+        EmptyPoolError if no connection is available within the time period.
+    :type pool_timeout: int | float | None
+    """
+
+    def __init__(
+        self,
+        request_timeout: Union[int, float] = DEFAULT_REQUEST_TIMEOUT,
+        retry_attempts: int = 3,
+        backoff_factor: float = 1.0,
+        pool_connections: int = 10,
+        pool_maxsize: int = 10,
+        pool_timeout: Union[int, float, None] = None,
+    ) -> None:
+        self.request_timeout = request_timeout
+        self._retry_attempts = retry_attempts
+        self._backoff_factor = backoff_factor
+        self._pool_connections = pool_connections
+        self._pool_maxsize = pool_maxsize
+        self._pool_timeout = pool_timeout
 
     def create_session(self, host: str) -> Session:
         """Create and return a new session/connection.
@@ -79,12 +163,18 @@ def create_session(self, host: str) -> Session:
         :rtype: requests.Session
         """
         retry_strategy = Retry(
-            total=self.RETRY_ATTEMPTS,
-            backoff_factor=self.BACKOFF_FACTOR,
+            total=self._retry_attempts,
+            backoff_factor=self._backoff_factor,
             status_forcelist=[429, 500, 502, 503, 504],
             allowed_methods=["HEAD", "GET", "OPTIONS"],
         )
-        http_adapter = HTTPAdapter(max_retries=retry_strategy)
+        http_adapter = DefaultHTTPAdapter(
+            connection_timeout=self.request_timeout,
+            pool_connections=self._pool_connections,
+            pool_maxsize=self._pool_maxsize,
+            pool_timeout=self._pool_timeout,
+            max_retries=retry_strategy,
+        )
 
         session = Session()
         session.mount("https://", http_adapter)
@@ -128,7 +218,7 @@ def send_request(
             data=data,
             headers=headers,
             auth=auth,
-            timeout=self.REQUEST_TIMEOUT,
+            timeout=self.request_timeout,
         )
         return Response(
             method=method,
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -54,7 +54,7 @@ def test_client_attributes():
     assert isinstance(client._host_resolver, RandomHostResolver)
 
     client = ArangoClient(hosts=client_hosts, request_timeout=120)
-    assert client.request_timeout == client._http.REQUEST_TIMEOUT == 120
+    assert client.request_timeout == client._http.request_timeout == 120
 
 
 def test_client_good_connection(db, username, password):
@@ -92,6 +92,23 @@ def test_client_bad_connection(db, username, password, cluster):
     assert "bad connection" in str(err.value)
 
 
+def test_client_http_client_attributes(db, username, password):
+    http_client = DefaultHTTPClient(
+        request_timeout=80,
+        retry_attempts=5,
+        backoff_factor=1.0,
+        pool_connections=16,
+        pool_maxsize=12,
+        pool_timeout=120,
+    )
+    client = ArangoClient(
+        hosts="http://127.0.0.1:8529", http_client=http_client, request_timeout=30
+    )
+    client.db(db.name, username, password, verify=True)
+    assert http_client.request_timeout == 80
+    assert client.request_timeout == http_client.request_timeout
+
+
 def test_client_custom_http_client(db, username, password):
     # Define custom HTTP client which increments the counter on any API call.
     class MyHTTPClient(DefaultHTTPClient):
