Skip to content

Add %create_graph_snapshot line magic #653

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jul 19, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions ChangeLog.md
Original file line number Diff line number Diff line change
Expand Up @@ -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))

Expand Down
54 changes: 54 additions & 0 deletions src/graph_notebook/magics/graph_magic.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
14 changes: 14 additions & 0 deletions src/graph_notebook/neptune/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Loading