diff --git a/ChangeLog.md b/ChangeLog.md index 729a47e3..233c54f3 100644 --- a/ChangeLog.md +++ b/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) diff --git a/src/graph_notebook/magics/graph_magic.py b/src/graph_notebook/magics/graph_magic.py index b7f94a15..c6cd10d8 100644 --- a/src/graph_notebook/magics/graph_magic.py +++ b/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 diff --git a/src/graph_notebook/neptune/client.py b/src/graph_notebook/neptune/client.py index dc021fab..91e24074 100644 --- a/src/graph_notebook/neptune/client.py +++ b/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,