Loader command bug fixes

* Fix load_status and cancel_load arg parsers * add more load status completion values Co-authored-by: Austin Kline <[email protected]>
aws · Nov 23, 2020 · e0a1be6 · e0a1be6
1 parent 3694ce0
commit e0a1be6
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 7 deletions.
diff --git a/src/graph_notebook/__init__.py b/src/graph_notebook/__init__.py
@@ -3,4 +3,4 @@
 SPDX-License-Identifier: Apache-2.0
 """
 
-__version__ = '2.0.0'
+__version__ = '2.0.1'
diff --git a/src/graph_notebook/loader/load.py b/src/graph_notebook/loader/load.py
@@ -20,8 +20,23 @@
 PARALLELISM_OPTIONS = [PARALLELISM_LOW, PARALLELISM_MEDIUM, PARALLELISM_HIGH, PARALLELISM_OVERSUBSCRIBE]
 LOADER_ACTION = 'loader'
 
+FINAL_LOAD_STATUSES = ['LOAD_COMPLETED',
+                       'LOAD_COMMITTED_W_WRITE_CONFLICTS',
+                       'LOAD_CANCELLED_BY_USER',
+                       'LOAD_CANCELLED_DUE_TO_ERRORS',
+                       'LOAD_FAILED',
+                       'LOAD_UNEXPECTED_ERROR',
+                       'LOAD_DATA_DEADLOCK',
+                       'LOAD_DATA_FAILED_DUE_TO_FEED_MODIFIED_OR_DELETED',
+                       'LOAD_S3_READ_ERROR',
+                       'LOAD_S3_ACCESS_DENIED_ERROR',
+                       'LOAD_IN_QUEUE',
+                       'LOAD_FAILED_BECAUSE_DEPENDENCY_NOT_SATISFIED',
+                       'LOAD_FAILED_INVALID_REQUEST', ]
 
-def do_load(host, port, load_format, use_ssl, source, region, arn, fail_on_error, parallelism, update_single_cardinality, request_param_generator):
+
+def do_load(host, port, load_format, use_ssl, source, region, arn, fail_on_error, parallelism,
+            update_single_cardinality, request_param_generator):
     payload = {
         'source': source,
         'format': load_format,

diff --git a/src/graph_notebook/magics/graph_magic.py b/src/graph_notebook/magics/graph_magic.py
@@ -35,7 +35,7 @@
 from graph_notebook.gremlin.client_provider.factory import create_client_provider
 from graph_notebook.request_param_generator.factory import create_request_generator
 from graph_notebook.loader.load import do_load, get_loader_jobs, get_load_status, cancel_load, VALID_FORMATS, \
-    PARALLELISM_OPTIONS, PARALLELISM_HIGH
+    PARALLELISM_OPTIONS, PARALLELISM_HIGH, FINAL_LOAD_STATUSES
 from graph_notebook.configuration.get_config import get_config, get_config_from_dict
 from graph_notebook.seed.load_query import get_data_sets, get_queries
 from graph_notebook.status.get_status import get_status
@@ -728,7 +728,7 @@ def on_button_clicked(b):
                         job_status_output.clear_output()
                         with job_status_output:
                             print(f'Overall Status: {interval_check_response["payload"]["overallStatus"]["status"]}')
-                            if interval_check_response["payload"]["overallStatus"]["status"] == 'LOAD_COMPLETED':
+                            if interval_check_response["payload"]["overallStatus"]["status"] in FINAL_LOAD_STATUSES:
                                 interval_output.close()
                                 print('Done.')
                                 return
@@ -778,13 +778,14 @@ def load_ids(self, line, local_ns: dict = None):
     @needs_local_scope
     def load_status(self, line, local_ns: dict = None):
         parser = argparse.ArgumentParser()
+        parser.add_argument('load_id', default='', help='loader id to check status for')
         parser.add_argument('--store-to', type=str, default='')
         args = parser.parse_args(line.split())
 
         credentials_provider_mode = self.graph_notebook_config.iam_credentials_provider_type
         request_generator = create_request_generator(self.graph_notebook_config.auth_mode, credentials_provider_mode)
         res = get_load_status(self.graph_notebook_config.host, self.graph_notebook_config.port,
-                              self.graph_notebook_config.ssl, request_generator, line)
+                              self.graph_notebook_config.ssl, request_generator, args.load_id)
         print(json.dumps(res, indent=2))
 
         if args.store_to != '' and local_ns is not None:
@@ -795,13 +796,14 @@ def load_status(self, line, local_ns: dict = None):
     @needs_local_scope
     def cancel_load(self, line, local_ns: dict = None):
         parser = argparse.ArgumentParser()
+        parser.add_argument('load_id', default='', help='loader id to check status for')
         parser.add_argument('--store-to', type=str, default='')
         args = parser.parse_args(line.split())
 
         credentials_provider_mode = self.graph_notebook_config.iam_credentials_provider_type
         request_generator = create_request_generator(self.graph_notebook_config.auth_mode, credentials_provider_mode)
         res = cancel_load(self.graph_notebook_config.host, self.graph_notebook_config.port,
-                          self.graph_notebook_config.ssl, request_generator, line)
+                          self.graph_notebook_config.ssl, request_generator, args.load_id)
         if res:
             print('Cancelled successfully.')
         else:

diff --git a/src/graph_notebook/widgets/package.json b/src/graph_notebook/widgets/package.json
@@ -1,6 +1,6 @@
 {
   "name": "graph_notebook_widgets",
-  "version": "2.0.0",
+  "version": "2.0.1",
   "author": "amazon",
   "description": "A Custom Jupyter Library for rendering NetworkX MultiDiGraphs using vis-network",
   "dependencies": {