Prometheus: add timeout parameter to query method

Change-Id: I4edf2745a5ba6a416d6bee12da573ced270d4023

spicerack/prometheus.py

@@ -2,7 +2,7 @@
 
 import logging
 
-from typing import Dict, List
+from typing import Dict, List, Optional, Union
 
 import requests
 
@@ -22,16 +22,20 @@ class Prometheus:
 
     _prometheus_api = 'http://prometheus.svc.{site}.wmnet/ops/api/v1/query'
 
-    def query(self, query: str, site: str) -> List[Dict]:
+    def query(self, query: str, site: str, *, timeout: Optional[Union[float, int]] = 10) -> List[Dict]:
         """Preform a generic query
 
-        Arguments
+        Arguments:
             query (str): a prometheus query
-            site (str): The site to use for queries. Must be one of :py:const:`spicerack.constants.ALL_DATACENTERS`
+            site (str): The site to use for queries. Must be one of
+                :py:const:`spicerack.constants.ALL_DATACENTERS`
+            timeout (float, int, None, optional): How many seconds to wait for the prometheus to
+                send data before giving up, as a float or int. Alternatively None to indicate an
+                infinite timeout.
 
-        Returns
+        Returns:
             list: returns an empty list if there are no results otherwise return a list of
-                  results of the form: {'metric': {}, 'value': [$timestamp, $value]}
+            results of the form: {'metric': {}, 'value': [$timestamp, $value]}
 
         Raises:
             spicerack.prometheus.PrometheusError: on error
@@ -43,7 +47,7 @@ def query(self, query: str, site: str) -> List[Dict]:
             raise PrometheusError(msg)
 
         url = self._prometheus_api.format(site=site)
-        response = requests.get(url, params={'query': query})
+        response = requests.get(url, params={'query': query}, timeout=timeout)
         if response.status_code != requests.codes['ok']:
             msg = 'Unable to get metric: HTTP {code}: {text}'.format(
                 code=response.status_code, text=response.text)