Improved ddb exception handling

coreplugins/dronedb/ddb.py

@@ -1,7 +1,6 @@
 import requests
 from os import path
 from app.plugins import logger
-import validators
 from urllib.parse import urlparse
 
 VALID_IMAGE_EXTENSIONS = ['.tiff', '.tif', '.png', '.jpeg', '.jpg']
@@ -10,7 +9,7 @@ class DroneDB:
   
     def __init__(self, registry_url, username, password):
 
-        if not validators.url(registry_url):
+        if not self.validate_url(registry_url):
             raise ValueError("Invalid registry URL.")
 
         self.username = username
@@ -27,109 +26,134 @@ class DroneDB:
         self.__get_files_list_url = self.__registry_url + "/orgs/{0}/ds/{1}/list"
         self.__download_file_url = self.__registry_url + "/orgs/{0}/ds/{1}/download?path={2}&inline=1"
  
+    # Validate url
+    def validate_url(self, url):
+        try:
+            result = urlparse(url)
+            return all([result.scheme, result.netloc])
+        except:
+            return False
+
     def login(self):
 
         if (self.public):
             logger.info("No need to login to DroneDB.")
             return True
 
-        # Authenticate
+        try:
-        payload = {'username': self.username, 'password': self.password}
-        response = requests.post(self.__authenticate_url, data=payload)
 
-        if response.status_code != 200:
+            # Authenticate
+            payload = {'username': self.username, 'password': self.password}
+            response = requests.post(self.__authenticate_url, data=payload)
+
+            if response.status_code != 200:
+                return False
+
+            # Get the token
+            self.token = response.json()['token']
+
+            logger.info("Logged in to DroneDB as user " + self.username + ".")    
+
+            return True
+        
+        except(Exception) as e:
+            logger.error(e)
             return False             
     
-        # Get the token
+    def wrapped_call(self, type, url, data=None, params=None, attempts=3):
-        self.token = response.json()['token']
         
-        logger.info("Logged in to DroneDB as user " + self.username + ".")    
+        headers = {}
+        
+        cnt = attempts
+
+        while True:
+                           
+            if not self.public and self.token is None and not self.login():
+                raise ValueError("Could not authenticate to DroneDB.")
+            
+            if self.token is not None:
+                headers = {'Authorization': 'Bearer ' + self.token }
+                
+            response = requests.request(type, url, data=data, params=params, headers=headers)
+                        
+            if response.status_code == 200:
+                return response
+            
+            if response.status_code == 401:            
+                if (self.public):
+                    raise Exception("Failed to call '" + url + "'.")
+                
+                if not self.login():
+                    raise Exception("Failed to re-authenticate to DroneDB, cannot call '" + url + "'.")
+                else:
+                    cnt -= 1
+                    if cnt == 0:
+                        raise Exception("Failed all attempts to re-authenticate to DroneDB, cannot call '" + url + "'.")
+            else:
+                raise Exception("Failed to call '" + url + "'.")        
         
-        return True            
 
     def get_organizations(self):
        
-        headers = {}
+        try:
         
-        if not self.public and self.token is None and not self.login():
+            response = self.wrapped_call('GET', self.__get_organizations_url)
-            raise ValueError("Could not authenticate to DroneDB.")
 
-        if self.token is not None:
+            return [{'slug': o['slug'], 'name': o['name']} for o in response.json()]
-            headers = {'Authorization': 'Bearer ' + self.token }
         
-        # Get the organizations
+        except Exception as e:
-        response = requests.get(self.__get_organizations_url, headers=headers)
+            raise Exception("Failed to get organizations.") from e
-        if response.status_code != 200:
-            raise Exception("Failed to get organizations.")
-
-        return [{'slug': o['slug'], 'name': o['name']} for o in response.json()]
 
     def get_datasets(self, orgSlug):
         
-        headers = {}
+        try:
             
-        if not self.public and self.token is None and not self.login():
+            response = self.wrapped_call('GET', self.__get_datasets_url.format(orgSlug))
-            raise ValueError("Could not authenticate to DroneDB.")
             
-        if self.token is not None:
+            return [
-            headers = {'Authorization': 'Bearer ' + self.token }
+                {'slug': o['slug'], 
+                'name': o['name'], 
+                'public': o['properties'].get('public'), 
+                'size': o['size'], 
+                'entries': o['properties'].get('entries')
+                } for o in response.json()]
         
-        # Get the datasets
+        except Exception as e:
-        response = requests.get(self.__get_datasets_url.format(orgSlug), headers=headers)
+            raise Exception("Failed to get datasets.") from e
-        if response.status_code != 200:
-            raise Exception("Failed to get datasets.")
                
-        return [
-            {'slug': o['slug'], 
-             'name': o['name'], 
-             'public': o['properties'].get('public'), 
-             'size': o['size'], 
-             'entries': o['properties'].get('entries')
-            } for o in response.json()]
     
     def get_folders(self, orgSlug, dsSlug):
         
-        headers = {}
+        try:
             
-        if not self.public and self.token is None and not self.login():
+            # Type 1 is folder
-            raise ValueError("Could not authenticate to DroneDB.")
+            payload = {'query': '*', 'recursive': True, 'type': 1}
 
-        if self.token is not None:
+            response = self.wrapped_call('POST', self.__get_folders_url.format(orgSlug, dsSlug), data=payload)
-            headers = {'Authorization': 'Bearer ' + self.token }
 
-        # Type 1 is folder
+            return [o['path'] for o in response.json()]        
-        payload = {'query': '*', 'recursive': True, 'type': 1}
                 
-        # Get the folders
+        except Exception as e:
-        response = requests.post(self.__get_folders_url.format(orgSlug, dsSlug), data=payload, headers=headers)
+            raise Exception("Failed to get folders.") from e
-        if response.status_code != 200:
-            raise Exception("Failed to get folders.")
-
-        return [o['path'] for o in response.json()]
 
     def get_files_list(self, orgSlug, dsSlug, folder=None):
         
-        headers = {}
+        try:
                     
-        if not self.public and self.token is None and not self.login():
+            # Type 1 is folder
-            raise ValueError("Could not authenticate to DroneDB.")
+            params = {'path': '' if folder is None else folder}
 
-        if self.token is not None:
+            # Get the folders
-            headers = {'Authorization': 'Bearer ' + self.token }
+            response = self.wrapped_call('GET', self.__get_files_list_url.format(orgSlug, dsSlug), params=params)
 
-        # Type 1 is folder
+            return [{'path': o['path'], 
-        params = {'path': '' if folder is None else folder}
+                'type': o['type'], 
+                'size': o['size'],
+                'url': self.__download_file_url.format(orgSlug, dsSlug, o['path'])
+                } for o in response.json()]
             
-        # Get the folders
+        except Exception as e:
-        response = requests.get(self.__get_files_list_url.format(orgSlug, dsSlug), params=params, headers=headers)
+            raise Exception("Failed to get files list.") from e
-        if response.status_code != 200:
-            raise Exception("Failed to get files list.")
-
-        return [{'path': o['path'], 
-             'type': o['type'], 
-             'size': o['size'],
-             'url': self.__download_file_url.format(orgSlug, dsSlug, o['path'])
-             } for o in response.json()]
 
 def verify_url(url, username=None, password=None):
     try: