Add %get_graph line magic, enable %status for Neptune Analytics

ChangeLog.md

@@ -4,6 +4,7 @@ Starting with v1.31.6, this file will contain a record of major features and upd
 
 ## Upcoming
 - Added `%reset_graph` line magic ([Link to PR](https://github.com/aws/graph-notebook/pull/610))
+- Added `%get_graph` line magic and enabled `%status` for Neptune Analytics ([Link to PR](https://github.com/aws/graph-notebook/pull/611))
 
 ## Release 4.3.1 (June 3, 2024)
 

src/graph_notebook/magics/graph_magic.py

@@ -1411,8 +1411,11 @@ def opencypher_status(self, line='', local_ns: dict = None):
     @line_magic
     @needs_local_scope
     @display_exceptions
-    @neptune_db_only
     def status(self, line='', local_ns: dict = None):
+        if self.client.is_analytics_domain():
+            logger.info(f'Redirected %status call to %get_graph.')
+            self.get_graph(line, local_ns)
+            return
         logger.info(f'calling for status on endpoint {self.graph_notebook_config.host}')
         parser = argparse.ArgumentParser()
         parser.add_argument('--silent', action='store_true', default=False, help="Display no output.")
@@ -1439,6 +1442,33 @@ def status(self, line='', local_ns: dict = None):
                     print()
                 return status_res
 
+    @line_magic
+    @needs_local_scope
+    @display_exceptions
+    @neptune_graph_only
+    def get_graph(self, line='', local_ns: dict = None):
+        logger.info(f'calling for status on endpoint {self.graph_notebook_config.host}')
+        parser = argparse.ArgumentParser()
+        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())
+
+        try:
+            graph_id = self.client.get_graph_id()
+            res = self.client.get_graph(graph_id=graph_id)
+            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 retrieve graph status:\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

@@ -589,6 +589,16 @@ def reset_graph(self, graph_id: str = '', no_skip_snapshot: bool = False) -> dic
             logger.debug(f"Reset Graph call failed with service exception: {e}")
             raise e
 
+    def get_graph(self, graph_id: str = '') -> dict:
+        try:
+            res = self.neptune_graph_client.get_graph(
+                graphIdentifier=graph_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,