Add %create_graph_snapshot line magic

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))
 

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

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,