Add %get_import_task line magic (#668)

* Add %get_import_task line magic

* update changelog

ChangeLog.md

@@ -3,15 +3,17 @@
 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)
 
 - New Neptune Analytics notebooks - openCypher over RDF ([Link to PR](https://github.com/aws/graph-notebook/pull/672))
   - 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))

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

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,