Add %reset_graph line magic (#610)

* Add ResetGraph magic for Neptune Analytics

* Bump botocore and boto3 reqs

* Update SkipSnapshot behavior

* update changelog

Author: michaelnchin

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)
 

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

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',

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

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,