Sync with master

Author: akharit

HISTORY.rst

@@ -3,6 +3,14 @@
 Release History
 ===============
 
+0.0.37 (2018-11-02)
++++++++++++++++++++
+* Reverted some changes introduced in 0.0.35 that didn't work with other tokens
+
+0.0.36 (2018-10-31)
++++++++++++++++++++
+* Fixed typo in refresh_token call
+
 0.0.35 (2018-10-29)
 +++++++++++++++++++
 * Added retry for getting tokens

azure/datalake/store/__init__.py

@@ -6,7 +6,8 @@
 # license information.
 # --------------------------------------------------------------------------
 
-__version__ = "0.0.35"
+__version__ = "0.0.37"
+
 
 
 from .core import AzureDLFileSystem

azure/datalake/store/lib.py

@@ -155,7 +155,9 @@ def get_token_internal():
 
     return DataLakeCredential(out)
 
+
 class DataLakeCredential:
+    # Be careful modifying this. DataLakeCredential is a general class in azure, and we have to maintain parity.
     def __init__(self, token):
         self.token = token
 
@@ -191,21 +193,15 @@ def refresh_token(self, authority=None, retry_policy=None):
         context = adal.AuthenticationContext(authority +
                                              self.token['tenant'])
 
-        @retry_decorator_for_auth(retry_policy=retry_policy)
-        def get_token_internal():
-            # Internal function used so as to use retry decorator
-            if self.token.get('secret') and self.token.get('client'):
-                out = context.acquire_token_with_client_credentials(self.token['resource'],
-                                                                    self.token['client'],
-                                                                    self.token['secret'])
-                out.update({'secret': self.token['secret']})
-            else:
-                out = context.acquire_token_with_refresh_token(self.token['refresh'],
-                                                               client_id=self.token['client'],
-                                                               resource=self.token['resource'])
-            return out
-
-        out = get_token_internal()
+        if self.token.get('secret') and self.token.get('client'):
+            out = context.acquire_token_with_client_credentials(self.token['resource'],
+                                                                self.token['client'],
+                                                                self.token['secret'])
+            out.update({'secret': self.token['secret']})
+        else:
+            out = context.acquire_token_with_refresh_token(self.token['refresh'],
+                                                           client_id=self.token['client'],
+                                                           resource=self.token['resource'])
         # common items to update
         out.update({'access': out['accessToken'],
                     'time': time.time(), 'tenant': self.token['tenant'],
@@ -271,7 +267,9 @@ def __init__(self, store_name=default_store, token=None,
         # There is a case where the user can opt to exclude an API version, in which case
         # the service itself decides on the API version to use (it's default).
         self.api_version = api_version or None
-        self.head = {'Authorization': token.signed_session(retry_policy=None).headers['Authorization']}
+        self.head = None
+        self._check_token()  # Retryable method. Will ensure that signed_session token is current when we set it on next line
+        self.head = {'Authorization': token.signed_session().headers['Authorization']}
         self.url = 'https://%s.%s/' % (store_name, url_suffix)
         self.webhdfs = 'webhdfs/v1/'
         self.extended_operations = 'webhdfsext/'
@@ -296,11 +294,15 @@ def session(self):
             self.local.session = s
         return s
 
-    def _check_token(self, retry_policy=None):
-        cur_session = self.token.signed_session(retry_policy=retry_policy)
-        if not self.head or self.head.get('Authorization') != cur_session.headers['Authorization']:
-            self.head = {'Authorization': cur_session.headers['Authorization']}
-            self.local.session = None
+
+    def _check_token(self, retry_policy= None):
+        @retry_decorator_for_auth(retry_policy=retry_policy)
+        def check_token_internal():
+            cur_session = self.token.signed_session()
+            if not self.head or self.head.get('Authorization') != cur_session.headers['Authorization']:
+                self.head = {'Authorization': cur_session.headers['Authorization']}
+                self.local.session = None
+        check_token_internal()
 
     def _log_request(self, method, url, op, path, params, headers, retry_count):
         msg = "HTTP Request\n{} {}\n".format(method.upper(), url)

azure/datalake/store/retry.py

@@ -86,9 +86,9 @@ def deco_retry(func):
         @wraps(func)
         def f_retry(*args, **kwargs):
             retry_count = -1
-            last_exception = None
-            out = None
+
             while True:
+                last_exception = None
                 retry_count += 1
                 try:
                     out = func(*args, **kwargs)
@@ -105,7 +105,7 @@ def f_retry(*args, **kwargs):
                 request_successful = last_exception is None or response.status_code == 401  # 401 = Invalid credentials
                 if request_successful or not retry_policy.should_retry(response, last_exception, retry_count):
                     break
-            if out is None:
+            if last_exception is not None:
                 raise last_exception
             return out