diff --git a/src/graph_notebook/magics/graph_magic.py b/src/graph_notebook/magics/graph_magic.py index 93c6265e..7379c5b6 100644 --- a/src/graph_notebook/magics/graph_magic.py +++ b/src/graph_notebook/magics/graph_magic.py @@ -1603,6 +1603,45 @@ def create_graph_snapshot(self, line='', local_ns: dict = None): print(e) store_to_ns(args.store_to, e, local_ns) + @line_magic + @needs_local_scope + @display_exceptions + @neptune_graph_only + def list_import_tasks(self, line='', local_ns: dict = None): + parser = argparse.ArgumentParser() + parser.add_argument('-m', '--max-results', type=int, default=None, + help="The total number of records to return in the command's output, valid range is " + "[1,100]. If there are more available import tasks than --max-results," + " then you will receive a nextToken in the output, which you can resubmit with the " + "--next-token argument to paginate on the next subset of results.") + parser.add_argument('-nt', '--next-token', type=str, default='', + help="Pagination token used to paginate output. If applicable, tokens can be found as the " + "nextToken field in the output of a previous %get_import_task execution. When this " + "argument is provided as input, the service returns results from where the previous " + "response left off.") + parser.add_argument('--include-metadata', action='store_true', default=False, + help="Display the response metadata if it is available.") + parser.add_argument('--silent', action='store_true', default=False, help="Display no output.") + parser.add_argument('--store-to', type=str, default='', help='store query result to this variable') + args = parser.parse_args(line.split()) + + if args.max_results and not (1 <= args.max_results <= 100): + print("--max-results must be in range [1,100].") + return + + try: + res = self.client.list_import_tasks(max_results=args.max_results, next_token=args.next_token) + if not args.include_metadata: + res.pop('ResponseMetadata', None) + if not args.silent: + print(json.dumps(res, indent=2, default=str)) + store_to_ns(args.store_to, res, local_ns) + except Exception as e: + if not args.silent: + print("Encountered an error when attempting to list import tasks:\n") + print(e) + store_to_ns(args.store_to, e, local_ns) + @line_magic @needs_local_scope @display_exceptions diff --git a/src/graph_notebook/neptune/client.py b/src/graph_notebook/neptune/client.py index 4099147d..467e31ab 100644 --- a/src/graph_notebook/neptune/client.py +++ b/src/graph_notebook/neptune/client.py @@ -680,6 +680,19 @@ def create_graph_snapshot(self, graph_id: str = '', snapshot_name: str = '', tag logger.debug(f"CreateGraphSnapshot call failed with service exception: {e}") raise e + def list_import_tasks(self, max_results: int = None, next_token: str = '') -> dict: + kwargs = {} + if max_results is not None: + kwargs['maxResults'] = max_results + if next_token != '': + kwargs['nextToken'] = next_token + try: + res = self.neptune_graph_client.list_import_tasks(**kwargs) + return res + except ClientError as e: + logger.debug(f"GetGraph call failed with service exception: {e}") + raise e + def dataprocessing_start(self, s3_input_uri: str, s3_output_uri: str, **kwargs) -> requests.Response: data = { 'inputDataS3Location': s3_input_uri,