[CIVP-10885] ENH Retry connection errors in CLI (#117)

* ENH Retry connection errors in CLI

When retrieving the API spec in the Civis client CLI, use the same code as we use in the client library. This has the advantage of including automatic retries for connection failures, and automatic waiting for rate limit refreshes. Also modify `invoke` so that we include the same retries as we would if we were using the client library.

* REF DRY API connection

We were creating connections to the Civis API in multiple places. Centralize everything in a single `open_session` function, and use that everywhere.

* ENH Change CLI's User-Agent to "civis-cli"

Make the CLI's User-Agent string distinct from the Python client's User-Agent. The CLI will send a User-Agent string which starts with "civis-cli".

Author: Stephen Hoover

CHANGELOG.md

@@ -7,6 +7,8 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 ### Changed
 - Edited example for safer null value handling
 - Make ``pubnub`` and ``joblib`` hard dependencies instead of optional dependencies (#110).
+- Retry network errors and wait for API rate limit refresh when using the CLI (#117).
+- The CLI now provides a User-Agent header which starts with "civis-cli" (#117)
 - Include ``pandas`` and ``sklearn``-dependent code in Travis CI tests.
 
 ### Added

civis/_utils.py

@@ -1,6 +1,15 @@
+from __future__ import absolute_import
+from builtins import super
+
 import re
 import uuid
 
+import requests
+from requests.adapters import HTTPAdapter
+from requests.packages.urllib3.util import Retry
+
+import civis
+
 
 UNDERSCORER1 = re.compile(r'(.)([A-Z][a-z]+)')
 UNDERSCORER2 = re.compile('([a-z0-9])([A-Z])')
@@ -20,3 +29,41 @@ def camel_to_snake(word):
 
 def to_camelcase(s):
     return re.sub(r'(^|_)([a-zA-Z])', lambda m: m.group(2).upper(), s)
+
+
+def open_session(api_key, max_retries=5, user_agent="civis-python"):
+    """Create a new Session which can connect with the Civis API"""
+    civis_version = civis.__version__
+    session = requests.Session()
+    session.auth = (api_key, '')
+    session_agent = session.headers.get('User-Agent', '')
+    user_agent = "{}/{} {}".format(user_agent, civis_version, session_agent)
+    session.headers.update({"User-Agent": user_agent.strip()})
+    max_retries = AggressiveRetry(max_retries, backoff_factor=.75,
+                                  status_forcelist=civis.civis.RETRY_CODES)
+    adapter = HTTPAdapter(max_retries=max_retries)
+    session.mount("https://", adapter)
+
+    return session
+
+
+class AggressiveRetry(Retry):
+    # Subclass Retry so that it retries more things. In particular,
+    # always retry API requests with a Retry-After header, regardless
+    # of the verb.
+    def is_retry(self, method, status_code, has_retry_after=False):
+        """ Is this method/status code retryable? (Based on whitelists and control
+        variables such as the number of total retries to allow, whether to
+        respect the Retry-After header, whether this header is present, and
+        whether the returned status code is on the list of status codes to
+        be retried upon on the presence of the aforementioned header)
+        """
+        if (self.total and
+                self.respect_retry_after_header and
+                has_retry_after and
+                (status_code in self.RETRY_AFTER_STATUS_CODES)):
+            return True
+
+        else:
+            return super().is_retry(method=method, status_code=status_code,
+                                    has_retry_after=has_retry_after)

civis/base.py

@@ -7,8 +7,6 @@
 import six
 import warnings
 
-from requests.packages.urllib3.util import Retry
-
 from civis.response import PaginatedResponse, convert_response_data_type
 
 FINISHED = ['success', 'succeeded']
@@ -76,28 +74,6 @@ def get_base_url():
     return base_url
 
 
-class AggressiveRetry(Retry):
-    # Subclass Retry so that it retries more things. In particular,
-    # always retry API requests with a Retry-After header, regardless
-    # of the verb.
-    def is_retry(self, method, status_code, has_retry_after=False):
-        """ Is this method/status code retryable? (Based on whitelists and control
-        variables such as the number of total retries to allow, whether to
-        respect the Retry-After header, whether this header is present, and
-        whether the returned status code is on the list of status codes to
-        be retried upon on the presence of the aforementioned header)
-        """
-        if (self.total and
-                self.respect_retry_after_header and
-                has_retry_after and
-                (status_code in self.RETRY_AFTER_STATUS_CODES)):
-            return True
-
-        else:
-            return super().is_retry(method=method, status_code=status_code,
-                                    has_retry_after=has_retry_after)
-
-
 class Endpoint(object):
 
     _lock = threading.Lock()

civis/civis.py

@@ -2,13 +2,9 @@
 import logging
 import os
 
-import requests
-from requests.adapters import HTTPAdapter
-
-import civis
-from civis.base import AggressiveRetry
 from civis.compat import lru_cache
 from civis.resources import generate_classes_maybe_cached
+from civis._utils import open_session
 
 
 log = logging.getLogger(__name__)
@@ -299,19 +295,7 @@ def __init__(self, api_key=None, return_type='snake',
                              "'pandas'")
         self._feature_flags = ()
         session_auth_key = _get_api_key(api_key)
-        self._session = session = requests.session()
-        session.auth = (session_auth_key, '')
-
-        civis_version = civis.__version__
-        session_agent = session.headers.get('User-Agent', '')
-        user_agent = "civis-python/{} {}".format(civis_version, session_agent)
-        session.headers.update({"User-Agent": user_agent.strip()})
-
-        max_retries = AggressiveRetry(retry_total, backoff_factor=.75,
-                                      status_forcelist=RETRY_CODES)
-        adapter = HTTPAdapter(max_retries=max_retries)
-
-        session.mount("https://", adapter)
+        self._session = session = open_session(session_auth_key, retry_total)
 
         classes = generate_classes_maybe_cached(local_api_spec,
                                                 session_auth_key,

civis/cli/__main__.py

@@ -22,10 +22,11 @@
 
 import click
 from jsonref import JsonRef
-import requests
 import yaml
 from civis.cli._cli_commands import \
     civis_ascii_art, files_download_cmd, files_upload_cmd
+from civis.resources import get_api_spec
+from civis._utils import open_session
 from civis.compat import FileNotFoundError
 
 
@@ -34,6 +35,7 @@
 _OPENAPI_SPEC_URL = "https://api.civisanalytics.com/endpoints"
 _CACHED_SPEC_PATH = \
     os.path.join(os.path.expanduser('~'), ".civis_api_spec.json")
+CLI_USER_AGENT = 'civis-cli'
 
 
 class YAMLParamType(click.ParamType):
@@ -71,18 +73,14 @@ def convert(self, value, param, ctx):
 }
 
 
-def make_api_request_headers():
+def get_api_key():
     try:
-        headers = {
-            'Authorization': "Bearer {}".format(os.environ["CIVIS_API_KEY"])
-        }
+        return os.environ["CIVIS_API_KEY"]
     except KeyError:
         print("You must set the CIVIS_API_KEY environment variable.",
               file=sys.stderr)
         sys.exit(1)
 
-    return headers
-
 
 def camel_to_snake(s):
     return re.sub(r'([A-Z]+)', r'_\1', s).lower()
@@ -160,13 +158,13 @@ def invoke(method, path, op, *args, **kwargs):
 
     # Make the request.
     request_info = dict(
-        headers=make_api_request_headers(),
         params=query,
         json=body,
         url=_API_URL + path.format(**kwargs),
         method=method
     )
-    response = requests.request(**request_info)
+    with open_session(get_api_key(), user_agent=CLI_USER_AGENT) as sess:
+        response = sess.request(**request_info)
 
     # Print the response to stderr if there was an error.
     output_file = sys.stdout
@@ -192,7 +190,7 @@ def invoke(method, path, op, *args, **kwargs):
             print("Error parsing response: {}".format(e), file=sys.stderr)
 
 
-def retrieve_spec_dict():
+def retrieve_spec_dict(api_version="1.0"):
     """Retrieve the API specification from a cached version or from Civis."""
 
     refresh_spec = True
@@ -211,15 +209,10 @@ def retrieve_spec_dict():
 
     # Download the spec and cache it in the user's home directory.
     if refresh_spec:
-        resp = requests.get(_OPENAPI_SPEC_URL,
-                            headers=make_api_request_headers())
-        assert resp.status_code == 200, \
-            "Failure downloading API specification: %d %s" % \
-            (resp.status_code, resp.reason)
+        spec_dict = get_api_spec(get_api_key(), api_version=api_version,
+                                 user_agent=CLI_USER_AGENT)
         with open(_CACHED_SPEC_PATH, "w") as f:
-            f.write(resp.text)
-        spec_dict = json.loads(resp.text, object_pairs_hook=OrderedDict)
-
+            json.dump(spec_dict, f)
     return spec_dict
 
 

civis/resources/_resources.py

@@ -9,12 +9,10 @@
 
 from jsonref import JsonRef
 import requests
-from requests.adapters import HTTPAdapter
 
-import civis
-from civis.base import AggressiveRetry, Endpoint, get_base_url
+from civis.base import Endpoint, get_base_url
 from civis.compat import lru_cache
-from civis._utils import camel_to_snake, to_camelcase
+from civis._utils import camel_to_snake, to_camelcase, open_session
 
 
 API_VERSIONS = ["1.0"]
@@ -435,7 +433,7 @@ def parse_api_spec(api_spec, api_version, resources):
 
 
 @lru_cache(maxsize=4)
-def get_api_spec(api_key, api_version="1.0"):
+def get_api_spec(api_key, api_version="1.0", user_agent="civis-python"):
     """Download the Civis API specification.
 
     Parameters
@@ -445,19 +443,13 @@ def get_api_spec(api_key, api_version="1.0"):
     api_version : string, optional
         The version of endpoints to call. May instantiate multiple client
         objects with different versions.  Currently only "1.0" is supported.
+    user_agent : string, optional
+        Provide this user agent to the the Civis API, along with an
+        API client version tag and ``requests`` version tag.
     """
-    civis_version = civis.__version__
-    session = requests.Session()
-    session.auth = (api_key, '')
-    session_agent = session.headers.get('User-Agent', '')
-    user_agent = "civis-python/{} {}".format(civis_version, session_agent)
-    session.headers.update({"User-Agent": user_agent.strip()})
-    max_retries = AggressiveRetry(MAX_RETRIES, backoff_factor=.75,
-                                  status_forcelist=civis.civis.RETRY_CODES)
-    adapter = HTTPAdapter(max_retries=max_retries)
-    session.mount("https://", adapter)
     if api_version == "1.0":
-        response = session.get("{}endpoints".format(get_base_url()))
+        with open_session(api_key, MAX_RETRIES, user_agent=user_agent) as sess:
+            response = sess.get("{}endpoints".format(get_base_url()))
     else:
         msg = "API specification for api version {} cannot be found"
         raise ValueError(msg.format(api_version))

civis/tests/test_cli.py

@@ -69,53 +69,38 @@ def test_generate_cli_civis(mock_retrieve_spec_dict):
             assert p.required
 
 
-@mock.patch("civis.cli.__main__.make_api_request_headers")
-@mock.patch("civis.cli.__main__.yaml")
-@mock.patch("civis.cli.__main__.requests.request")
-def test_blank_output(mock_request, mock_yaml, mock_make_api_request_headers):
+@mock.patch("civis.cli.__main__.open_session", autospec=True)
+def test_blank_output(mock_session):
     """
     Test that endpoints that return blank results don't cause exceptions.
-
-    We mock make_api_request_headers because the invoke function get the API
-    key from that, so this test will fail in environments where an API key
-    isn't in the env (e.g., travis).
-
-    We mock yaml because otherwise yaml will complain about being able to
-    serialize the mock response object.
     """
-
-    mock_make_api_request_headers.return_value = {}
-
     # The response object's json method will raise a ValueError when the output
     # is blank.
-    mock_request.return_value.json.side_effect = ValueError()
+    session_context = mock_session.return_value.__enter__.return_value
+    session_context.request.return_value.json.side_effect = ValueError()
 
     op = {"parameters": []}
     invoke("WIBBLE", "/wobble/wubble", op)
 
 
-@mock.patch("civis.cli.__main__.make_api_request_headers")
-@mock.patch("civis.cli.__main__.yaml")
-@mock.patch("civis.cli.__main__.requests.request")
-def test_parameter_case(mock_request, mock_yaml,
-                        mock_make_api_request_headers):
+@mock.patch("civis.cli.__main__.open_session", autospec=True)
+def test_parameter_case(mock_session):
     """
     Test that parameter names are sent in camelCase rather than snake_case.
-
-    We mock yaml because otherwise yaml will complain about being able to
-    serialize the mock response object.
     """
+    api_response = {'key': 'value'}
+    session_context = mock_session.return_value.__enter__.return_value
+    session_context.request.return_value.json.return_value = api_response
 
     # To avoid needing CIVIS_API_KEY set in the environment.
-    mock_make_api_request_headers.return_value = {}
     op = {"parameters": [{'name': 'firstParameter', 'in': 'query'},
                          {'name': 'secondParameter', 'in': 'query'}]}
     invoke("WIBBLE", "/wobble/wubble", op,
            first_parameter='a', second_parameter='b')
 
-    mock_request.assert_called_with(
+    mock_session.call_args[1]['user_agent'] = 'civis-cli'
+    session_context.request.assert_called_with(
         url='https://api.civisanalytics.com/wobble/wubble',
-        headers={},
         json={},
         params={'firstParameter': 'a', 'secondParameter': 'b'},
         method='WIBBLE')