From 0df06babeeeebdb5be67d47a914ecae7a75401c5 Mon Sep 17 00:00:00 2001
From: choeh <75442468+choeh@users.noreply.github.com>
Date: Mon, 3 May 2021 20:00:01 +0200
Subject: [PATCH] Add timeout option in connection wrapper (#119)

---
 bin/qds.py            | 11 ++++++++++-
 qds_sdk/connection.py | 13 +++++++------
 qds_sdk/qubole.py     |  9 ++++++---
 3 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/bin/qds.py b/bin/qds.py
index 01cde904..2e991207 100755
--- a/bin/qds.py
+++ b/bin/qds.py
@@ -602,6 +602,11 @@ def main():
                          help="Number of re-attempts for an api-call in case of "
                               " retryable exceptions. Defaults to 7.")
 
+    optparser.add_option("--timeout", dest="timeout",
+                         type=int,
+                         default=os.getenv('QDS_TIMEOUT'),
+                         help="number of seconds to await response from QDS. defaults to 300s")
+
     optparser.add_option("-v", dest="verbose", action="store_true",
                          default=False,
                          help="verbose mode - info level logging")
@@ -639,6 +644,9 @@ def main():
     if options.base_retry_delay is None:
         options.base_retry_delay = 10
 
+    if options.timeout is None:
+        options.timeout = 300
+
     if options.cloud_name is None:
         options.cloud_name = "AWS"
 
@@ -654,7 +662,8 @@ def main():
                      skip_ssl_cert_check=options.skip_ssl_cert_check,
                      cloud_name=options.cloud_name,
                      base_retry_delay=options.base_retry_delay,
-                     max_retries=options.max_retries
+                     max_retries=options.max_retries,
+                     timeout=options.timeout
                      )
 
     if len(args) < 1:
diff --git a/qds_sdk/connection.py b/qds_sdk/connection.py
index 324c27f5..4dd90f97 100644
--- a/qds_sdk/connection.py
+++ b/qds_sdk/connection.py
@@ -37,7 +37,7 @@ class Connection:
 
     def __init__(self, auth, rest_url, skip_ssl_cert_check,
                  reuse=True, max_retries=7,
-                 base_retry_delay=10):
+                 base_retry_delay=10, timeout=300):
         self.auth = auth
         self.rest_url = rest_url
         self.skip_ssl_cert_check = skip_ssl_cert_check
@@ -47,6 +47,7 @@ def __init__(self, auth, rest_url, skip_ssl_cert_check,
         self.reuse = reuse
         self.max_retries = max_retries
         self.base_retry_delay = base_retry_delay
+        self.timeout = timeout
         if reuse:
             self.session = requests.Session()
             self.session.mount('https://', RequestAdapter())
@@ -109,7 +110,7 @@ def _api_call_raw(self, req_type, path, data=None, params=None):
             x_with_retries = requests.Session()
             x_with_retries.mount('https://', RequestAdapter(max_retries=3))
 
-        kwargs = {'headers': self._headers, 'auth': self.auth, 'verify': not self.skip_ssl_cert_check}
+        kwargs = {'headers': self._headers, 'auth': self.auth, 'verify': not self.skip_ssl_cert_check, 'timeout': self.timeout}
 
         if data:
             kwargs['data'] = json.dumps(data)
@@ -121,13 +122,13 @@ def _api_call_raw(self, req_type, path, data=None, params=None):
         log.info("Params: %s" % params)
 
         if req_type == 'GET':
-            r = x_with_retries.get(url, timeout=300, **kwargs)
+            r = x_with_retries.get(url, **kwargs)
         elif req_type == 'POST':
-            r = x.post(url, timeout=300, **kwargs)
+            r = x.post(url, **kwargs)
         elif req_type == 'PUT':
-            r = x.put(url, timeout=300, **kwargs)
+            r = x.put(url, **kwargs)
         elif req_type == 'DELETE':
-            r = x.delete(url, timeout=300, **kwargs)
+            r = x.delete(url, **kwargs)
         else:
             raise NotImplemented
 
diff --git a/qds_sdk/qubole.py b/qds_sdk/qubole.py
index 58b40304..29adc9d9 100644
--- a/qds_sdk/qubole.py
+++ b/qds_sdk/qubole.py
@@ -23,6 +23,7 @@ class Qubole:
     MIN_POLL_INTERVAL = 1
     RETRIES_CAP = 7
     MAX_RETRY_DELAY = 10
+    TIMEOUT = 300
 
     _auth = None
     api_token = None
@@ -40,7 +41,7 @@ class Qubole:
     def configure(cls, api_token,
                   api_url="https://api.qubole.com/api/", version="v1.2",
                   poll_interval=5, skip_ssl_cert_check=False, cloud_name="AWS",
-                  base_retry_delay=10, max_retries=7):
+                  base_retry_delay=10, max_retries=7, timeout=300):
         """
         Set parameters governing interaction with QDS
         Args:
@@ -51,6 +52,7 @@ def configure(cls, api_token,
             `delay` : interval in secs to sleep in between successive retries
             `retries` : maximum number of time to retry an api call in case
                         of retryable exception.
+            `timeout` : maximum time in secs to wait for response from QDS
         """
 
         cls._auth = QuboleAuth(api_token)
@@ -81,6 +83,7 @@ def configure(cls, api_token,
             cls.max_retries = Qubole.RETRIES_CAP
         else:
             cls.max_retries = max_retries
+        cls.timeout = timeout
 
     @classmethod
     def agent(cls, version=None):
@@ -105,12 +108,12 @@ def agent(cls, version=None):
         if not reuse_cached_agent:
             uncached_agent = Connection(cls._auth, cls.rest_url,
                                         cls.skip_ssl_cert_check,
-                                        True, cls.max_retries, cls.base_retry_delay)
+                                        True, cls.max_retries, cls.base_retry_delay, cls.timeout)
             return uncached_agent
         if cls.cached_agent is None:
             cls.cached_agent = Connection(cls._auth, cls.rest_url,
                                           cls.skip_ssl_cert_check,
-                                          True, cls.max_retries, cls.base_retry_delay)
+                                          True, cls.max_retries, cls.base_retry_delay, cls.timeout)
 
         return cls.cached_agent