From cf9a5e01e1b9e2a3d2e81841a540a81eefc6fd92 Mon Sep 17 00:00:00 2001
From: Michael Chin <chnmch@amazon.com>
Date: Wed, 17 Jul 2024 16:03:55 -0700
Subject: [PATCH 1/2] 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 <chnmch@amazon.com>
Date: Wed, 17 Jul 2024 17:28:32 -0700
Subject: [PATCH 2/2] 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)