Skip to content

Add --connection-protocol option to %%gremlin #617

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 2 commits into from
Jun 12, 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 @@ -4,6 +4,8 @@ Starting with v1.31.6, this file will contain a record of major features and upd

## Upcoming

- Added `--connection-protocol` option to `%%gremlin` ([Link to PR](https://github.com/aws/graph-notebook/pull/617))

## Release 4.4.0 (June 10, 2024)

- Added `%reset_graph` line magic ([Link to PR](https://github.com/aws/graph-notebook/pull/610))
Expand Down
26 changes: 18 additions & 8 deletions src/graph_notebook/magics/graph_magic.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,8 @@
NEPTUNE_CONFIG_HOST_IDENTIFIERS, is_allowed_neptune_host, \
STATISTICS_LANGUAGE_INPUTS, STATISTICS_LANGUAGE_INPUTS_SPARQL, STATISTICS_MODES, SUMMARY_MODES, \
SPARQL_EXPLAIN_MODES, OPENCYPHER_EXPLAIN_MODES, OPENCYPHER_PLAN_CACHE_MODES, OPENCYPHER_DEFAULT_TIMEOUT, \
OPENCYPHER_STATUS_STATE_MODES, normalize_service_name, GRAPH_PG_INFO_METRICS
OPENCYPHER_STATUS_STATE_MODES, normalize_service_name, GRAPH_PG_INFO_METRICS, \
DEFAULT_GREMLIN_PROTOCOL, GREMLIN_PROTOCOL_FORMATS, DEFAULT_HTTP_PROTOCOL, normalize_protocol_name
from graph_notebook.network import SPARQLNetwork
from graph_notebook.network.gremlin.GremlinNetwork import parse_pattern_list_str, GremlinNetwork
from graph_notebook.visualization.rows_and_columns import sparql_get_rows_and_columns, opencypher_get_rows_and_columns
Expand Down Expand Up @@ -1027,6 +1028,11 @@ def gremlin(self, line, cell, local_ns: dict = None):
parser = argparse.ArgumentParser()
parser.add_argument('query_mode', nargs='?', default='query',
help='query mode (default=query) [query|explain|profile]')
parser.add_argument('-cp', '--connection-protocol', type=str.lower, default='',
help=f'Neptune endpoints only. Connection protocol to use for connecting to the Gremlin '
f'database - either Websockets or HTTP. Valid inputs: {GREMLIN_PROTOCOL_FORMATS}. '
f'Default is {DEFAULT_GREMLIN_PROTOCOL}. Please note that this option has no effect '
f'on the Profile and Explain modes, which must use HTTP.')
parser.add_argument('--explain-type', type=str.lower, default='',
help='Explain mode to use when using the explain query mode.')
parser.add_argument('-p', '--path-pattern', default='', help='path pattern')
Expand Down Expand Up @@ -1167,13 +1173,17 @@ def gremlin(self, line, cell, local_ns: dict = None):
else:
using_http = False
query_start = time.time() * 1000 # time.time() returns time in seconds w/high precision; x1000 to get in ms
if self.graph_notebook_config.proxy_host != '' and self.client.is_neptune_domain():
using_http = True
query_res_http = self.client.gremlin_http_query(cell, headers={
'Accept': 'application/vnd.gremlin-v1.0+json;types=false'})
query_res_http.raise_for_status()
query_res_http_json = query_res_http.json()
query_res = query_res_http_json['result']['data']
if self.client.is_neptune_domain():
connection_protocol = normalize_protocol_name(args.connection_protocol)
if self.graph_notebook_config.proxy_host != '' or connection_protocol == DEFAULT_HTTP_PROTOCOL:
using_http = True
query_res_http = self.client.gremlin_http_query(cell, headers={
'Accept': 'application/vnd.gremlin-v1.0+json;types=false'})
query_res_http.raise_for_status()
query_res_http_json = query_res_http.json()
query_res = query_res_http_json['result']['data']
else:
query_res = self.client.gremlin_query(cell, transport_args=transport_args)
else:
query_res = self.client.gremlin_query(cell, transport_args=transport_args)
query_time = time.time() * 1000 - query_start
Expand Down
17 changes: 17 additions & 0 deletions src/graph_notebook/neptune/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,13 @@
GRAPHSONV2_VARIANTS = ['graphsonv2', 'graphsonv2d0', 'graphsonserializersv2d0', 'graphsonmessageserializerv2']
GRAPHBINARYV1_VARIANTS = ['graphbinaryv1', 'graphbinary', 'graphbinaryserializersv1', 'graphbinarymessageserializerv1']

DEFAULT_WS_PROTOCOL = "websockets"
DEFAULT_HTTP_PROTOCOL = "http"
WS_PROTOCOL_FORMATS = ["ws", "websocket", DEFAULT_WS_PROTOCOL]
HTTP_PROTOCOL_FORMATS = ["https", "rest", DEFAULT_HTTP_PROTOCOL]
GREMLIN_PROTOCOL_FORMATS = WS_PROTOCOL_FORMATS + HTTP_PROTOCOL_FORMATS
DEFAULT_GREMLIN_PROTOCOL = DEFAULT_WS_PROTOCOL

STATISTICS_MODES = ["", "status", "disableAutoCompute", "enableAutoCompute", "refresh", "delete"]
SUMMARY_MODES = ["", "basic", "detailed"]
STATISTICS_LANGUAGE_INPUTS_PG = ["propertygraph", "pg", "gremlin", "oc", "opencypher"]
Expand Down Expand Up @@ -154,6 +161,16 @@ def get_gremlin_serializer(serializer_str: str):
return serializer.GraphSONSerializersV3d0()


def normalize_protocol_name(protocol: str):
if protocol in WS_PROTOCOL_FORMATS:
return DEFAULT_WS_PROTOCOL
elif protocol in HTTP_PROTOCOL_FORMATS:
return DEFAULT_HTTP_PROTOCOL
else:
print(f"Provided connection protocol is invalid for Neptune, defaulting to {DEFAULT_GREMLIN_PROTOCOL}.")
return DEFAULT_GREMLIN_PROTOCOL


def normalize_service_name(neptune_service: str):
if neptune_service in NEPTUNE_ANALYTICS_CONFIG_NAMES:
return NEPTUNE_ANALYTICS_SERVICE_NAME
Expand Down