#277 Add the ability to specify request timeouts (#280)

* #277 allow setting a default timeout for all requests

* #277 allow custom timeout for custom_query and custom_query_range

* Apply suggestions from code review

Author: FRosner

prometheus_api_client/prometheus_connect.py

@@ -40,6 +40,7 @@ class PrometheusConnect:
     :param proxy: (Optional) Proxies dictionary to enable connection through proxy.
         Example: {"http_proxy": "<ip_address/hostname:port>", "https_proxy": "<ip_address/hostname:port>"}
     :param session (Optional) Custom requests.Session to enable complex HTTP configuration
+    :param timeout: (Optional) A timeout (in seconds) applied to all requests
     """
 
     def __init__(
@@ -51,6 +52,7 @@ def __init__(
         auth: tuple = None,
         proxy: dict = None,
         session: Session = None,
+        timeout: int = None,
     ):
         """Functions as a Constructor for the class PrometheusConnect."""
         if url is None:
@@ -60,6 +62,7 @@ def __init__(
         self.url = url
         self.prometheus_host = urlparse(self.url).netloc
         self._all_metrics = None
+        self._timeout = timeout
 
         if retry is None:
             retry = Retry(
@@ -94,7 +97,8 @@ def check_prometheus_connection(self, params: dict = None) -> bool:
             headers=self.headers,
             params=params,
             auth=self.auth,
-            cert=self._session.cert
+            cert=self._session.cert,
+            timeout=self._timeout,
         )
         return response.ok
 
@@ -131,7 +135,8 @@ def get_label_names(self, params: dict = None):
             headers=self.headers,
             params=params,
             auth=self.auth,
-            cert=self._session.cert
+            cert=self._session.cert,
+            timeout=self._timeout,
         )
 
         if response.status_code == 200:
@@ -161,7 +166,8 @@ def get_label_values(self, label_name: str, params: dict = None):
             headers=self.headers,
             params=params,
             auth=self.auth,
-            cert=self._session.cert
+            cert=self._session.cert,
+            timeout=self._timeout,
         )
 
         if response.status_code == 200:
@@ -212,7 +218,8 @@ def get_current_metric_value(
             verify=self._session.verify,
             headers=self.headers,
             auth=self.auth,
-            cert=self._session.cert
+            cert=self._session.cert,
+            timeout=self._timeout,
         )
 
         if response.status_code == 200:
@@ -304,7 +311,8 @@ def get_metric_range_data(
                 verify=self._session.verify,
                 headers=self.headers,
                 auth=self.auth,
-                cert=self._session.cert
+                cert=self._session.cert,
+                timeout=self._timeout,
             )
             if response.status_code == 200:
                 data += response.json()["data"]["result"]
@@ -377,7 +385,7 @@ def _metric_filename(self, metric_name: str, end_timestamp: int):
         )
         return object_path
 
-    def custom_query(self, query: str, params: dict = None):
+    def custom_query(self, query: str, params: dict = None, timeout: int = None):
         """
         Send a custom query to a Prometheus Host.
 
@@ -388,6 +396,7 @@ def custom_query(self, query: str, params: dict = None):
             at https://prometheus.io/docs/prometheus/latest/querying/examples/
         :param params: (dict) Optional dictionary containing GET parameters to be
             sent along with the API request, such as "time"
+        :param timeout: (Optional) A timeout (in seconds) applied to the request
         :returns: (list) A list of metric data received in response of the query sent
         :raises:
             (RequestException) Raises an exception in case of a connection error
@@ -396,14 +405,16 @@ def custom_query(self, query: str, params: dict = None):
         params = params or {}
         data = None
         query = str(query)
+        timeout = self._timeout if timeout is None else timeout
         # using the query API to get raw data
         response = self._session.get(
             "{0}/api/v1/query".format(self.url),
             params={**{"query": query}, **params},
             verify=self._session.verify,
             headers=self.headers,
             auth=self.auth,
-            cert=self._session.cert
+            cert=self._session.cert,
+            timeout=timeout,
         )
         if response.status_code == 200:
             data = response.json()["data"]["result"]
@@ -415,7 +426,7 @@ def custom_query(self, query: str, params: dict = None):
         return data
 
     def custom_query_range(
-        self, query: str, start_time: datetime, end_time: datetime, step: str, params: dict = None
+        self, query: str, start_time: datetime, end_time: datetime, step: str, params: dict = None, timeout: int = None
     ):
         """
         Send a query_range to a Prometheus Host.
@@ -430,6 +441,7 @@ def custom_query_range(
         :param step: (str) Query resolution step width in duration format or float number of seconds
         :param params: (dict) Optional dictionary containing GET parameters to be
             sent along with the API request, such as "timeout"
+        :param timeout: (Optional) A timeout (in seconds) applied to the request
         :returns: (dict) A dict of metric data received in response of the query sent
         :raises:
             (RequestException) Raises an exception in case of a connection error
@@ -440,14 +452,16 @@ def custom_query_range(
         params = params or {}
         data = None
         query = str(query)
+        timeout = self._timeout if timeout is None else timeout
         # using the query_range API to get raw data
         response = self._session.get(
             "{0}/api/v1/query_range".format(self.url),
             params={**{"query": query, "start": start, "end": end, "step": step}, **params},
             verify=self._session.verify,
             headers=self.headers,
             auth=self.auth,
-            cert=self._session.cert
+            cert=self._session.cert,
+            timeout=timeout,
         )
         if response.status_code == 200:
             data = response.json()["data"]["result"]