diff --git a/ChangeLog.md b/ChangeLog.md index 026c8be4..729a47e3 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 `%reset_graph` line magic ([Link to PR](https://github.com/aws/graph-notebook/pull/610)) ## Release 4.3.1 (June 3, 2024) diff --git a/requirements.txt b/requirements.txt index 12339484..a7a4529b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -25,7 +25,7 @@ numpy<1.24.0 nest_asyncio>=1.5.5,<=1.5.6 # requirements for testing -botocore~=1.21.49 -boto3~=1.18.49 +botocore~=1.34.74 +boto3~=1.34.74 pytest==6.2.5 parameterized==0.8.1 \ No newline at end of file diff --git a/setup.py b/setup.py index 8bcdacf5..31a3c73d 100644 --- a/setup.py +++ b/setup.py @@ -78,8 +78,8 @@ def get_version(): 'jupyter-contrib-nbextensions<=0.7.0', 'widgetsnbextension<=3.6.1', 'jupyter==1.0.0', - 'botocore>=1.21.49', - 'boto3>=1.18.49', + 'botocore>=1.34.74', + 'boto3>=1.34.74', 'ipython>=7.16.1,<=8.10.0', 'neo4j>=4.4.9,<5.0.0', 'rdflib==7.0.0', diff --git a/src/graph_notebook/magics/graph_magic.py b/src/graph_notebook/magics/graph_magic.py index 2d535df8..b7f94a15 100644 --- a/src/graph_notebook/magics/graph_magic.py +++ b/src/graph_notebook/magics/graph_magic.py @@ -1601,6 +1601,39 @@ def on_button_cancel_clicked(b): logger.info(f'got the response {res}') return res + @line_magic + @needs_local_scope + @display_exceptions + @neptune_graph_only + def graph_reset(self, line, local_ns: dict = None): + self.reset_graph(line, local_ns) + + @line_magic + @needs_local_scope + @display_exceptions + @neptune_graph_only + def reset_graph(self, line, local_ns: dict = None): + parser = argparse.ArgumentParser() + parser.add_argument('-ns', '--no-skip-snapshot', action='store_true', default=False, + help='Creates a final graph snapshot before the graph data is deleted.') + 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.reset_graph(graph_id=graph_id, no_skip_snapshot=args.no_skip_snapshot) + if not args.silent: + print(f"ResetGraph call submitted successfully for graph ID [{graph_id}]. Please note that the graph " + f"may take several minutes to become available again.\n") + 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("Received an error when attempting graph reset:") + print(e) + store_to_ns(args.store_to, e, local_ns) + @magic_variables @line_magic @needs_local_scope diff --git a/src/graph_notebook/neptune/client.py b/src/graph_notebook/neptune/client.py index de5ec4aa..dc021fab 100644 --- a/src/graph_notebook/neptune/client.py +++ b/src/graph_notebook/neptune/client.py @@ -12,9 +12,11 @@ from urllib.parse import urlparse, urlunparse from SPARQLWrapper import SPARQLWrapper from boto3 import Session +from boto3 import client as boto3_client from botocore.session import Session as botocoreSession from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest +from botocore.exceptions import ClientError from gremlin_python.driver import client, serializer from gremlin_python.driver.protocol import GremlinServerError from gremlin_python.driver.aiohttp.transport import AiohttpTransport @@ -202,6 +204,8 @@ def __init__(self, host: str, port: int = DEFAULT_PORT, self._http_session = None + self.neptune_graph_client = boto3_client(service_name='neptune-graph', region_name=self.region) + @property def host(self): if self.proxy_host != '': @@ -242,6 +246,11 @@ def get_uri_with_port(self, use_websocket=False, use_proxy=False): uri = f'{protocol}://{uri_host}:{uri_port}' return uri + def get_graph_id(self): + graph_host = self.host + graph_id = graph_host.split('.')[0] + return graph_id + def sparql_query(self, query: str, headers=None, explain: str = '', path: str = '') -> requests.Response: if headers is None: headers = {} @@ -569,6 +578,17 @@ def perform_reset(self, token: str) -> requests.Response: res = self._http_session.send(req, verify=self.ssl_verify) return res + def reset_graph(self, graph_id: str = '', no_skip_snapshot: bool = False) -> dict: + try: + res = self.neptune_graph_client.reset_graph( + graphIdentifier=graph_id, + skipSnapshot=(not no_skip_snapshot) + ) + return res + except ClientError as e: + logger.debug(f"Reset Graph 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,