validate user inputs

Signed-off-by: Ching Yi, Chan <qrtt1@infuseai.io>

primehub/__init__.py

@@ -214,8 +214,8 @@ def _ensure_config_details(self, config: PrimeHubConfig):
         except BaseException:
             pass
 
-    def request(self, variables: dict, query: str):
-        return Client(self.primehub_config).request(variables, query)
+    def request(self, variables: dict, query: str, error_handler: Callable = None):
+        return Client(self.primehub_config).request(variables, query, error_handler)
 
     def request_logs(self, endpint: str, follow: bool, tail: int):
         return Client(self.primehub_config).request_logs(endpint, follow, tail)

primehub/cli.py

@@ -3,7 +3,7 @@
 
 import primehub as ph
 from primehub.config import CliConfig
-from primehub.utils import create_logger, PrimeHubException, PrimeHubReturnsRequiredException
+from primehub.utils import create_logger, PrimeHubException, PrimeHubReturnsRequiredException, ResourceNotFoundException
 from primehub.utils.argparser import create_command_parser, create_action_parser
 from primehub.utils.decorators import find_actions, find_action_method, find_action_info
 from primehub.utils.permission import has_permission_flag, enable_ask_for_permission_feature
@@ -246,6 +246,13 @@ def main(sdk=None):
         hide_help = True
         exit_normally = False
         sys.exit(1)
+    except ResourceNotFoundException as e:
+        logger.debug('got PrimeHubReturnsRequiredException')
+        hide_help = True
+        exit_normally = False
+        type_name, key, key_type = e.args
+        print(f'Cannot find resource type "{type_name}" with [{key_type}: {key}]', file=sdk.stderr)
+        sys.exit(1)
     except PrimeHubException as e:
         hide_help = True
         exit_normally = False

primehub/groups.py

@@ -1,6 +1,7 @@
 from typing import Optional
 
 from primehub import Helpful, cmd, Module
+from primehub.utils import resource_not_found
 
 
 class Groups(Helpful, Module):
@@ -53,6 +54,8 @@ def get(self, group_name: str) -> Optional[dict]:
         group = [x for x in groups if x['name'] == group_name]
         if group:
             return group[0]
+
+        resource_not_found('group', group_name, 'name')
         return None
 
     def help_description(self):

primehub/jobs.py

@@ -6,9 +6,29 @@
 import json
 import sys
 
+from primehub.utils import resource_not_found, PrimeHubException
 from primehub.utils.optionals import toggle_flag
 
 
+def _error_handler(response):
+    import re
+
+    if 'errors' in response:
+        message = [x for x in response['errors'] if 'message' in x]
+        if message:
+            message = message[0]['message']
+            result = re.findall(r'phjobs.primehub.io "([^"]+)" not found', message)
+            if result:
+                resource_not_found('job', result[0], 'id')
+
+
+def invalid_config(message: str):
+    example = """
+    {"instanceType":"cpu-1","image":"base-notebook","displayName":"job-example","command":"echo 'good job'"}
+    """.strip()
+    raise PrimeHubException(message + "\n\nExample:\n" + json.dumps(json.loads(example), indent=2))
+
+
 class Jobs(Helpful, Module):
     """
     The jobs module provides functions to manage PrimeHub Jobs
@@ -18,7 +38,7 @@ class Jobs(Helpful, Module):
         "instanceType": "cpu-1",
         "image": "base-notebook",
         "displayName": "test",
-        "command": "echo \"test\"",
+        "command": "echo \"test\""
     }
     """
 
@@ -133,11 +153,10 @@ def get(self, id):
           }
         }
         """
-        results = self.request({'where': {'id': id}}, query)
+        results = self.request({'where': {'id': id}}, query, _error_handler)
         return results['data']['phJob']
 
     # TODO: add -f
-    # TODO: handel invalid config
     @cmd(name='submit', description='Submit a job', optionals=[('file', str), ('from', str)])
     def submit_cmd(self, **kwargs):
         """
@@ -164,10 +183,12 @@ def submit_cmd(self, **kwargs):
         if has_data_from_stdin():
             config = json.loads("".join(sys.stdin.readlines()))
 
+        if not config:
+            invalid_config('Job description is required.')
+
         config['groupId'] = self.group_id
         return self.submit(config)
 
-    # TODO: Add validation for config
     def submit(self, config):
         """
         Submit a job with config
@@ -208,7 +229,23 @@ def submit(self, config):
           }
         }
         """
+
+        if not config or (len(config) == 1):
+            raise PrimeHubException('config is required')
+        print(config, len(config))
+
         config['groupId'] = self.group_id
+
+        # verify required fields in the config
+        if 'instanceType' not in config:
+            invalid_config('instanceType is required')
+        if 'image' not in config:
+            invalid_config('image is required')
+        if 'displayName' not in config:
+            invalid_config('displayName is required')
+        if 'command' not in config:
+            invalid_config('command is required')
+
         results = self.request({'data': config}, query)
         return results['data']['createPhJob']
 
@@ -231,10 +268,9 @@ def submit_from_schedule(self, id):
           }
         }
         """
-        results = self.request({'where': {'id': id}}, query)
+        results = self.request({'where': {'id': id}}, query, _error_handler)
         return results['data']['runPhSchedule']
 
-    # TODO: handel id does not exist
     @cmd(name='rerun', description='Rerun a job by id')
     def rerun(self, id):
         """
@@ -276,7 +312,7 @@ def rerun(self, id):
           }
         }
         """
-        results = self.request({'where': {'id': id}}, query)
+        results = self.request({'where': {'id': id}}, query, _error_handler)
         return results['data']['rerunPhJob']
 
     # TODO: handel id does not exist
@@ -298,7 +334,7 @@ def cancel(self, id):
           }
         }
         """
-        self.request({'where': {'id': id}}, query)
+        self.request({'where': {'id': id}}, query, _error_handler)
         return self.get(id)
 
     @cmd(name='wait', description='Wait a job by id', optionals=[('timeout', int)])
@@ -325,7 +361,7 @@ def wait(self, id, **kwargs):
         timeout = kwargs.get('timeout', 0)
         start_time = time.time()
         while True:
-            results = self.request({'where': {'id': id}}, query)
+            results = self.request({'where': {'id': id}}, query, _error_handler)
             phase = results['data']['phJob']['phase']
             if phase in ['Succeeded', 'Failed', 'Cancelled']:
                 break
@@ -362,7 +398,7 @@ def logs(self, id, **kwargs) -> Iterator[str]:
         follow = kwargs.get('follow', False)
         tail = kwargs.get('tail', 10)
 
-        results = self.primehub.request({'where': {'id': id}}, query)
+        results = self.primehub.request({'where': {'id': id}}, query, _error_handler)
         endpoint = results['data']['phJob']['logEndpoint']
         return self.primehub.request_logs(endpoint, follow, tail)
 
@@ -390,10 +426,9 @@ def list_artifacts(self, id):
           }
         }
         """
-        results = self.request({'where': {'id': id}}, query)
+        results = self.request({'where': {'id': id}}, query, _error_handler)
         return results['data']['phJob']['artifact']['items']
 
-    # TODO: handel path or dest does not exist
     @cmd(name='download-artifacts', description='Download artifacts', optionals=[('recursive', toggle_flag)])
     def download_artifacts(self, id, path, dest, **kwargs):
         """
@@ -415,6 +450,9 @@ def download_artifacts(self, id, path, dest, **kwargs):
         if path in ['.', '', './']:
             path = '/'
         path = os.path.join('/jobArtifacts', id, path.lstrip('/'))
+
+        # get id to verify existing
+        self.get(id)
         self.primehub.files.download(path, dest, **kwargs)
         return
 

primehub/resource_operations.py

@@ -1,3 +1,6 @@
+from primehub.utils import resource_not_found
+
+
 class GroupResourceOperation(object):
 
     def do_list(self, query, resource_key):
@@ -14,4 +17,4 @@ def do_get(self, query, resource_key, resource_name):
         data = [x for x in resources if x['name'] == resource_name]
         if data:
             return data[0]
-        return None
+        resource_not_found(resource_key, resource_name, 'name')

primehub/schedules.py

@@ -1,12 +1,46 @@
 from typing import Iterator
 
 from primehub import Helpful, cmd, Module, has_data_from_stdin
+from primehub.utils import resource_not_found, PrimeHubException
 from primehub.utils.permission import ask_for_permission
 import os
 import json
 import sys
 
 
+def _error_handler(response):
+    import re
+
+    if 'errors' in response:
+        message = [x for x in response['errors'] if 'message' in x]
+        if message:
+            message = message[0]['message']
+            result = re.findall(r'phschedules.primehub.io "([^"]+)" not found', message)
+            if result:
+                resource_not_found('schedule', result[0], 'id')
+
+
+def invalid_config(message: str):
+    example = """
+    {"instanceType":"cpu-1","image":"base-notebook","displayName":"schedule-example","command":"echo 'good job'","recurrence":{"type":"daily","cron":"0 4 * * *"}}
+    """.strip()
+    raise PrimeHubException(message + "\n\nExample:\n" + json.dumps(json.loads(example), indent=2))
+
+
+def verify_config(config):
+    # verify required fields in the config
+    if 'instanceType' not in config:
+        invalid_config('instanceType is required')
+    if 'image' not in config:
+        invalid_config('image is required')
+    if 'displayName' not in config:
+        invalid_config('displayName is required')
+    if 'command' not in config:
+        invalid_config('command is required')
+    if 'recurrence' not in config:
+        invalid_config('recurrence is required')
+
+
 class Schedules(Helpful, Module):
     """
     The schedules module provides functions to manage PrimeHub Schedules
@@ -133,11 +167,10 @@ def get(self, id):
           }
         }
         """
-        results = self.request({'where': {'id': id}}, query)
+        results = self.request({'where': {'id': id}}, query, _error_handler)
         return results['data']['phSchedule']
 
     # TODO: add -f
-    # TODO: handel invalid config
     @cmd(name='create', description='Create a schedule', optionals=[('file', str)])
     def create_cmd(self, **kwargs):
         """
@@ -158,9 +191,11 @@ def create_cmd(self, **kwargs):
         if has_data_from_stdin():
             config = json.loads("".join(sys.stdin.readlines()))
 
+        if not config:
+            invalid_config('Schedule description is required.')
+
         return self.create(config)
 
-    # TODO: add validation for config
     def create(self, config):
         """
         Create a schedules with config
@@ -201,11 +236,12 @@ def create(self, config):
         }
         """
         config['groupId'] = self.group_id
+
+        verify_config(config)
         results = self.request({'data': config}, query)
         return results['data']['createPhSchedule']
 
     # TODO: add -f
-    # TODO: handel invalid config
     @cmd(name='update', description='Update a schedule by id', optionals=[('file', str)])
     def update_cmd(self, id, **kwargs):
         """
@@ -226,9 +262,11 @@ def update_cmd(self, id, **kwargs):
         if has_data_from_stdin():
             config = json.loads("".join(sys.stdin.readlines()))
 
+        if not config:
+            invalid_config('Schedule description is required.')
+
         return self.update(id, config)
 
-    # TODO: add validation for config
     def update(self, id, config):
         """
         Update a schedule with config
@@ -272,7 +310,8 @@ def update(self, id, config):
         }
         """
         config['groupId'] = self.group_id
-        results = self.request({'data': config, 'where': {'id': id}}, query)
+        verify_config(config)
+        results = self.request({'data': config, 'where': {'id': id}}, query, _error_handler)
         return results['data']['updatePhSchedule']
 
     @ask_for_permission
@@ -294,7 +333,7 @@ def delete(self, id):
           }
         }
         """
-        results = self.request({'where': {'id': id}}, query)
+        results = self.request({'where': {'id': id}}, query, _error_handler)
         return results['data']['deletePhSchedule']
 
     def help_description(self):

primehub/utils/__init__.py

@@ -31,6 +31,10 @@ class PrimeHubReturnsRequiredException(PrimeHubException):
     pass
 
 
+class ResourceNotFoundException(PrimeHubException):
+    pass
+
+
 def reject_action(action):
     raise UserRejectAction(
         'User rejects action [%s], please use the flag "--yes-i-really-mean-it" to allow the action.' % action)
@@ -44,6 +48,10 @@ def group_not_found(group):
     raise GroupIsRequiredException('No group information for [%s], please check the configuration again.' % group)
 
 
+def resource_not_found(resource_type: str, key: str, key_type: str):
+    raise ResourceNotFoundException(resource_type, key, key_type)
+
+
 def create_logger(name) -> logging.Logger:
     log_level = logging.WARNING
     if os.environ.get('PRIMEHUB_SDK_LOG_LEVEL') == 'DEBUG':