diff --git a/ChangeLog.md b/ChangeLog.md index 850700e4..a6249e25 100644 --- a/ChangeLog.md +++ b/ChangeLog.md @@ -3,7 +3,9 @@ Starting with v1.31.6, this file will contain a record of major features and updates made in each release of graph-notebook. ## Upcoming + - Upgraded Neo4j Bolt driver to v5.x ([Link to PR](https://github.com/aws/graph-notebook/pull/682)) +- Added `%get_import_task` line magic ([Link to PR](https://github.com/aws/graph-notebook/pull/668)) ## Release 4.5.2 (August 15, 2024) @@ -11,7 +13,7 @@ Starting with v1.31.6, this file will contain a record of major features and upd - Path: 02-Neptune-Analytics > 04-OpenCypher-Over-RDF - Updated OC-RDF samples to use `%load` magic, and pull from regional S3 buckets ([Link to PR](https://github.com/aws/graph-notebook/pull/676)) - Added regional S3 bucket mappings to Neptune CloudFormation template ([Link to PR](https://github.com/aws/graph-notebook/pull/664)) -- Enabled n-triples data for `%load` with Neptune Analytics ([PR #1](https://github.com/aws/graph-notebook/pull/671)) ( ([PR #2](https://github.com/aws/graph-notebook/pull/675))) +- Enabled n-triples data for `%load` with Neptune Analytics ([PR #1](https://github.com/aws/graph-notebook/pull/671)) ([PR #2](https://github.com/aws/graph-notebook/pull/675)) - Removed unused options from `%load`([Link to PR](https://github.com/aws/graph-notebook/pull/662)) - Made EncryptionKey optional in Neptune CloudFormation template ([Link to PR](https://github.com/aws/graph-notebook/pull/663)) - Fixed unintended type coercion in results table with missing/null values ([Link to PR](https://github.com/aws/graph-notebook/pull/679)) diff --git a/src/graph_notebook/magics/graph_magic.py b/src/graph_notebook/magics/graph_magic.py index 41afc326..f14ec02a 100644 --- a/src/graph_notebook/magics/graph_magic.py +++ b/src/graph_notebook/magics/graph_magic.py @@ -1603,6 +1603,44 @@ 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 get_import_task(self, line='', local_ns: dict = None): + parser = argparse.ArgumentParser() + parser.add_argument('-ti', '--task-identifier', type=str, default='', + help="The unique identifier of an import task.") + 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()) + + task_id_regex = "t-[a-z0-9]{10}" + if args.task_identifier == '': + print("Please provide an import task ID using the -ti or --task-identifier parameter.") + return + if not re.match(fr"^{task_id_regex}$", args.task_identifier): + print(f"Import task ID must satisfy the regular expression pattern: {task_id_regex}") + return + + try: + res = self.client.get_import_task(task_id=args.task_identifier) + 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: + if "ResourceNotFoundException" in str(e): + print(f"Unable to import task with ID: {args.task_identifier}") + else: + print("Encountered an error when attempting to retrieve the specified import task:\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 3f6a929c..be9cef34 100644 --- a/src/graph_notebook/neptune/client.py +++ b/src/graph_notebook/neptune/client.py @@ -680,6 +680,16 @@ 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 get_import_task(self, task_id: str = '') -> dict: + try: + res = self.neptune_graph_client.get_import_task( + taskIdentifier=task_id + ) + 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,