diff --git a/ChangeLog.md b/ChangeLog.md index d0fb1340..d0ff6655 100644 --- a/ChangeLog.md +++ b/ChangeLog.md @@ -3,6 +3,8 @@ Starting with v1.31.6, this file will contain a record of major features and updates made in each release of graph-notebook. ## Upcoming + +- Added `%create_graph_snapshot` line magic ([Link to PR](https://github.com/aws/graph-notebook/pull/653)) - Added better `%reset` user messaging on status check timeout ([Link to PR](https://github.com/aws/graph-notebook/pull/652)) - Upgraded `setuptools` dependency to 70.x ([Link to PR](https://github.com/aws/graph-notebook/pull/649)) diff --git a/src/graph_notebook/magics/graph_magic.py b/src/graph_notebook/magics/graph_magic.py index 7410d8af..30e074dd 100644 --- a/src/graph_notebook/magics/graph_magic.py +++ b/src/graph_notebook/magics/graph_magic.py @@ -1528,6 +1528,60 @@ def get_graph(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 create_graph_snapshot(self, line='', local_ns: dict = None): + parser = argparse.ArgumentParser() + parser.add_argument('-s', '--snapshot-name', type=str, default='', + help="The name for the snapshot. Must start with a letter, contain only alphanumeric " + "characters or hyphens, and not end with or contain two consecutive hyphens. If not " + "supplied, this will default to the format 'snapshot-[graph_id]-[timestamp]'.") + parser.add_argument('-t', '--tags', type=str, default='', + help='Metadata tags to attach to the graph snapshot. Pass a dict in string format, ' + 'ex. {"tag1":"foo","tag2":"bar"}') + 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()) + + graph_id = self.client.get_graph_id() + + if args.snapshot_name: + snapshot_name = args.snapshot_name + else: + datetime_iso = datetime.datetime.utcnow().isoformat() + timestamp = re.sub(r'\D', '', datetime_iso) + snapshot_name = f"snapshot-{graph_id}-{timestamp}" + + if args.tags: + try: + tags = json.loads(args.tags) + except JSONDecodeError as e: + print("Tags map is improperly formatted, skipping.") + tags = None + logger.error(e) + else: + tags = None + + try: + res = self.client.create_graph_snapshot(graph_id=graph_id, snapshot_name=snapshot_name, tags=tags) + if not args.include_metadata: + res.pop('ResponseMetadata', None) + if not args.silent: + print("Successfully submitted snapshot request:") + 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 create the graph snapshot:\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 510ba193..63d8bcdd 100644 --- a/src/graph_notebook/neptune/client.py +++ b/src/graph_notebook/neptune/client.py @@ -654,6 +654,20 @@ def get_graph(self, graph_id: str = '') -> dict: logger.debug(f"GetGraph call failed with service exception: {e}") raise e + def create_graph_snapshot(self, graph_id: str = '', snapshot_name: str = '', tags: dict = None) -> dict: + if not tags: + tags = {} + try: + res = self.neptune_graph_client.create_graph_snapshot( + graphIdentifier=graph_id, + snapshotName=snapshot_name, + tags=tags + ) + return res + except ClientError as e: + logger.debug(f"CreateGraphSnapshot 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,