4n4nd#277 allow setting a default timeout for all requests

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: (int) 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.
 
@@ -403,7 +411,8 @@ def custom_query(self, query: str, params: dict = None):
             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"]
@@ -447,7 +456,8 @@ def custom_query_range(
             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"]