From cf9a5e01e1b9e2a3d2e81841a540a81eefc6fd92 Mon Sep 17 00:00:00 2001 From: Michael Chin Date: Wed, 17 Jul 2024 16:03:55 -0700 Subject: [PATCH 1/4] Add %create_graph_snapshot line magic --- src/graph_notebook/magics/graph_magic.py | 54 ++++++++++++++++++++++++ src/graph_notebook/neptune/client.py | 14 ++++++ 2 files changed, 68 insertions(+) diff --git a/src/graph_notebook/magics/graph_magic.py b/src/graph_notebook/magics/graph_magic.py index ec312d1e..11b4cec8 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, From b80e4f2c729bdd273dead5a4c8d6834683e83696 Mon Sep 17 00:00:00 2001 From: Michael Chin Date: Wed, 17 Jul 2024 17:28:32 -0700 Subject: [PATCH 2/4] update changelog --- ChangeLog.md | 1 + 1 file changed, 1 insertion(+) diff --git a/ChangeLog.md b/ChangeLog.md index 1c8e11cb..f914ec6a 100644 --- a/ChangeLog.md +++ b/ChangeLog.md @@ -3,6 +3,7 @@ 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)) - Upgraded `setuptools` dependency to 70.x ([Link to PR](https://github.com/aws/graph-notebook/pull/649)) ## Release 4.5.0 (July 15, 2024) From 1a855336aa47b55cd4c0ecf3912df8cae7cbb771 Mon Sep 17 00:00:00 2001 From: Michael Chin Date: Wed, 17 Jul 2024 20:51:17 -0700 Subject: [PATCH 3/4] Initiate and poll snapshot creation independently in %reset --- src/graph_notebook/magics/graph_magic.py | 66 +++++++++++++++++++++--- src/graph_notebook/neptune/client.py | 8 +++ 2 files changed, 66 insertions(+), 8 deletions(-) diff --git a/src/graph_notebook/magics/graph_magic.py b/src/graph_notebook/magics/graph_magic.py index 30e074dd..82cda5c3 100644 --- a/src/graph_notebook/magics/graph_magic.py +++ b/src/graph_notebook/magics/graph_magic.py @@ -53,7 +53,8 @@ SPARQL_EXPLAIN_MODES, OPENCYPHER_EXPLAIN_MODES, GREMLIN_EXPLAIN_MODES, \ OPENCYPHER_PLAN_CACHE_MODES, OPENCYPHER_DEFAULT_TIMEOUT, OPENCYPHER_STATUS_STATE_MODES, normalize_service_name, NEPTUNE_DB_SERVICE_NAME, NEPTUNE_ANALYTICS_SERVICE_NAME, GRAPH_PG_INFO_METRICS, \ - DEFAULT_GREMLIN_PROTOCOL, GREMLIN_PROTOCOL_FORMATS, DEFAULT_HTTP_PROTOCOL, normalize_protocol_name) + DEFAULT_GREMLIN_PROTOCOL, GREMLIN_PROTOCOL_FORMATS, DEFAULT_HTTP_PROTOCOL, normalize_protocol_name, + generate_snapshot_name) from graph_notebook.network import SPARQLNetwork from graph_notebook.network.gremlin.GremlinNetwork import parse_pattern_list_str, GremlinNetwork from graph_notebook.visualization.rows_and_columns import sparql_get_rows_and_columns, opencypher_get_rows_and_columns @@ -1554,9 +1555,7 @@ def create_graph_snapshot(self, line='', local_ns: dict = None): 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}" + snapshot_name = generate_snapshot_name(graph_id) if args.tags: try: @@ -1628,8 +1627,12 @@ def reset(self, line, local_ns: dict = None, service: str = None): res = perform_reset_res.json() return res else: + if snapshot: + print(f"Snapshot creation is currently unsupported for prompt skip mode. Please use " + f"%create_graph_snapshot to take a snapshot prior to attempting graph reset.") + return try: - res = self.client.reset_graph(graph_id=graph_id, snapshot=snapshot) + res = self.client.reset_graph(graph_id=graph_id, snapshot=False) print( f"ResetGraph call submitted successfully for graph ID [{graph_id}]. " f"Please note that the graph may take several minutes to become available again, " @@ -1712,8 +1715,55 @@ def on_button_delete_clicked(b): logger.error(result) return else: + if snapshot_check_box.value: + snapshot_name = generate_snapshot_name(graph_id) + try: + self.client.create_graph_snapshot(graph_id=graph_id, snapshot_name=snapshot_name) + except Exception as e: + with output: + print("Graph snapshot creation request failed, please see the exception below.") + print(f"\n{e}") + logger.error(e) + return + + snapshot_static_output = widgets.Output() + snapshot_progress_output = widgets.Output() + snapshot_hbox = widgets.HBox([snapshot_static_output, snapshot_progress_output]) + with snapshot_static_output: + print("Creating graph snapshot, this may take several minutes") + with output: + display(snapshot_hbox) + + poll_interval = 5 + poll_index = 0 + status_ellipses = "" + while True: + snapshot_progress_output.clear_output() + poll_index += 1 + if poll_index > poll_interval: + snapshot_progress_output.clear_output() + status_ellipses = "" + interval_check_response = self.client.get_graph(graph_id=graph_id) + current_status = interval_check_response["status"] + if current_status == 'AVAILABLE': + snapshot_static_output.clear_output() + with snapshot_static_output: + print(f'Snapshot creation complete, starting reset.') + break + elif current_status != 'SNAPSHOTTING': + snapshot_static_output.clear_output() + with snapshot_static_output: + print(f'Something went wrong with the snapshot creation.') + return + poll_index = 0 + else: + status_ellipses += "." + with snapshot_progress_output: + print(status_ellipses) + time.sleep(1) + snapshot_progress_output.close() try: - result = self.client.reset_graph(graph_id=graph_id, snapshot=snapshot_check_box.value) + result = self.client.reset_graph(graph_id=graph_id, snapshot=False) except Exception as e: with output: print("Failed to initiate graph reset, please see the exception below.") @@ -1739,7 +1789,7 @@ def on_button_delete_clicked(b): interval_output.clear_output() if time_elapsed > poll_interval: with interval_output: - print('checking status...') + print('Checking status...') job_status_output.clear_output() new_interval = True try: @@ -1776,7 +1826,7 @@ def on_button_delete_clicked(b): display_html(HTML(loading_wheel_html)) new_interval = False with interval_output: - print(f'checking status in {time_remaining} seconds') + print(f'Checking status in {time_remaining} seconds') time.sleep(1) with (output): job_status_output.clear_output() diff --git a/src/graph_notebook/neptune/client.py b/src/graph_notebook/neptune/client.py index 63d8bcdd..4c262731 100644 --- a/src/graph_notebook/neptune/client.py +++ b/src/graph_notebook/neptune/client.py @@ -6,6 +6,7 @@ import json import logging import re +import datetime import requests import urllib3 @@ -191,6 +192,13 @@ def set_plan_cache_hint(query: str, plan_cache_value: str): return query_with_hint +def generate_snapshot_name(graph_id: str): + datetime_iso = datetime.datetime.utcnow().isoformat() + timestamp = re.sub(r'\D', '', datetime_iso) + snapshot_name = f"snapshot-{graph_id}-{timestamp}" + return snapshot_name + + class Client(object): def __init__(self, host: str, port: int = DEFAULT_PORT, neptune_service: str = NEPTUNE_DB_SERVICE_NAME, From 0ba04bc3955a3936cc8c224cde7b728924548175 Mon Sep 17 00:00:00 2001 From: Michael Chin Date: Wed, 17 Jul 2024 21:03:01 -0700 Subject: [PATCH 4/4] update changelog --- ChangeLog.md | 1 + 1 file changed, 1 insertion(+) diff --git a/ChangeLog.md b/ChangeLog.md index d0ff6655..9b9710fc 100644 --- a/ChangeLog.md +++ b/ChangeLog.md @@ -6,6 +6,7 @@ Starting with v1.31.6, this file will contain a record of major features and upd - 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)) +- Modified the `%reset --snapshot` option to use the CreateGraphSnapshot API ([Link to PR](https://github.com/aws/graph-notebook/pull/654)) - Upgraded `setuptools` dependency to 70.x ([Link to PR](https://github.com/aws/graph-notebook/pull/649)) ## Release 4.5.0 (July 15, 2024)