class GremlinUtils:
@classmethod
def init_statics(cls, globals):
statics.load_statics(globals)
del globals['range']
del globals['map']
del globals['min']
del globals['sum']
del globals['property']
del globals['max']
def __init__(self, endpoints=None):
if endpoints is None:
self.endpoints = Endpoints()
else:
self.endpoints = endpoints
def remote_connection(self, show_endpoint=True):
neptune_gremlin_endpoint = self.endpoints.gremlin_endpoint()
if show_endpoint:
print('gremlin: {}'.format(neptune_gremlin_endpoint))
retry_count = 0
while True:
try:
return DriverRemoteConnection(neptune_gremlin_endpoint, 'g')
except HTTPError as e:
exc_info = sys.exc_info()
if retry_count < 3:
retry_count += 1
print('Connection timeout. Retrying...')
else:
raise exc_info[0].with_traceback(exc_info[1], exc_info[2])
def traversal_source(self, show_endpoint=True, connection=None):
if connection is None:
connection = self.remote_connection(show_endpoint)
return traversal().withRemote(connection)
def client(self, pool_size=None, max_workers=None):
return Client(self.endpoints.gremlin_endpoint(),
'g',
pool_size=pool_size,
max_workers=max_workers)
def sessioned_client(self,
session_id=None,
pool_size=None,
max_workers=None):
return SessionedClient(
self.endpoints.gremlin_endpoint(),
'g',
uuid.uuid4().hex if session_id is None else session_id,
pool_size=pool_size,
max_workers=max_workers)
class GremlinUtils:
@classmethod
def init_statics(cls, globals):
statics.load_statics(globals)
del globals['range']
del globals['map']
del globals['min']
del globals['sum']
del globals['property']
del globals['max']
def __init__(self, endpoints=None):
if endpoints is None:
self.endpoints = Endpoints()
else:
self.endpoints = endpoints
self.connections = []
def close(self):
for connection in self.connections:
connection.close()
def remote_connection(self,
show_endpoint=False,
protocol_factory=None,
transport_factory=lambda: TornadoTransportProxy(),
pool_size=None,
max_workers=None,
message_serializer=None,
graphson_reader=None,
graphson_writer=None):
gremlin_endpoint = self.endpoints.gremlin_endpoint()
if show_endpoint:
print('gremlin: {}'.format(gremlin_endpoint))
retry_count = 0
while True:
try:
request_parameters = gremlin_endpoint.prepare_request()
signed_ws_request = httpclient.HTTPRequest(
request_parameters.uri, headers=request_parameters.headers)
connection = DriverRemoteConnection(
signed_ws_request,
'g',
protocol_factory=protocol_factory,
transport_factory=transport_factory,
pool_size=pool_size,
max_workers=max_workers,
message_serializer=message_serializer,
graphson_reader=graphson_reader,
graphson_writer=graphson_writer)
self.connections.append(connection)
return connection
except HTTPError as e:
exc_info = sys.exc_info()
if retry_count < 3:
retry_count += 1
print('Connection timeout. Retrying...')
else:
raise exc_info[0].with_traceback(exc_info[1], exc_info[2])
def traversal_source(self, show_endpoint=True, connection=None):
if connection is None:
connection = self.remote_connection(show_endpoint)
return traversal().withRemote(connection)
def client(self, pool_size=None, max_workers=None):
gremlin_endpoint = self.endpoints.gremlin_endpoint()
request_parameters = gremlin_endpoint.prepare_request()
signed_ws_request = httpclient.HTTPRequest(
request_parameters.uri, headers=request_parameters.headers)
return Client(signed_ws_request,
'g',
pool_size=pool_size,
max_workers=max_workers)
def sessioned_client(self,
session_id=None,
pool_size=None,
max_workers=None):
gremlin_endpoint = self.endpoints.gremlin_endpoint()
request_parameters = gremlin_endpoint.prepare_request()
signed_ws_request = httpclient.HTTPRequest(
request_parameters.uri, headers=request_parameters.headers)
return SessionedClient(
signed_ws_request,
'g',
uuid.uuid4().hex if session_id is None else session_id,
pool_size=pool_size,
max_workers=max_workers)
class NeptuneGremlinProxy(AbstractGremlinProxy):
"""
A proxy to a Neptune using the Gremlin protocol.
See also https://docs.aws.amazon.com/neptune/latest/userguide/access-graph-gremlin-differences.html
See also https://docs.aws.amazon.com/neptune/latest/userguide/access-graph-gremlin-sessions.html
TODO: HTTP proxy support. This does *NOT* support HTTP proxies as-is. Why? The default transport factory in
gremlin_python is tornado.websocket, which is hardcoded to use simple_httpclient (look at
WebSocketClientConnection). But, even if that could be made to use curl_httpclient, curl_httpclient requires pycurl
which requires libcurl and other native libraries which is a pain to install.
"""
def __init__(self, *, host: str, port: Optional[int] = None, user: str = None,
password: Optional[Union[str, boto3.session.Session]] = None,
driver_remote_connection_options: Mapping[str, Any] = {},
client_kwargs: Dict = dict(),
**kwargs: dict) -> None:
driver_remote_connection_options = dict(driver_remote_connection_options)
# port should be part of that url
if port is not None:
raise NotImplementedError(f'port is not allowed! port={port}')
# for IAM auth, we need the triplet or a Session which is more general
if isinstance(password, boto3.session.Session):
session = password
self.aws_auth = AssumeRoleAWS4Auth(session.get_credentials(), session.region_name, 'neptune-db')
else:
raise NotImplementedError(f'to use authentication, pass a boto3.session.Session!)')
if isinstance(host, str):
# usually a wss URI
url = urlsplit(host)
assert url.scheme in ('wss', 'ws') and url.path == '/gremlin' and not url.query and not url.fragment, \
f'url is not a Neptune ws url?: {host}'
self.endpoints = Endpoints(
neptune_endpoint=url.hostname, neptune_port=url.port,
region_name=session.region_name, credentials=session.get_credentials())
self.override_uri = None
elif isinstance(host, Mapping):
# ...but development is a little complicated
assert all(k in host for k in ('neptune_endpoint', 'neptune_port', 'uri')), \
f'pass a dict with neptune_endpoint, neptune_port, and uri not: {host}'
self.endpoints = Endpoints(
neptune_endpoint=host['neptune_endpoint'], neptune_port=int(host['neptune_port']),
region_name=session.region_name, credentials=session.get_credentials())
uri = urlsplit(host['uri'])
assert uri.scheme in ('wss', 'ws') and uri.path == '/gremlin' and not uri.query and not uri.fragment, \
f'''url is not a Neptune ws url?: {host['uri']}'''
self.override_uri = uri
else:
raise NotImplementedError(f'to use authentication, pass a Mapping with aws_access_key_id, '
f'aws_secret_access_key, service_region!')
# always g for Neptune
driver_remote_connection_options.update(traversal_source='g')
try:
s3_bucket_name = client_kwargs['neptune_bulk_loader_s3_bucket_name'] # noqa: E731
except Exception:
raise NotImplementedError(f'Cannot find s3 bucket name!')
# Instantiate bulk loader and graph traversal factory
bulk_loader_config: Dict[str, Any] = dict(NEPTUNE_SESSION=password, NEPTUNE_URL=host,
NEPTUNE_BULK_LOADER_S3_BUCKET_NAME=s3_bucket_name)
self.neptune_bulk_loader_api = NeptuneBulkLoaderApi.create_from_config(bulk_loader_config)
self.neptune_graph_traversal_source_factory = get_neptune_graph_traversal_source_factory(session=password,
neptune_url=host)
AbstractGremlinProxy.__init__(self, key_property_name='key',
driver_remote_connection_options=driver_remote_connection_options)
@classmethod
@overrides
def script_translator(cls) -> Type[ScriptTranslatorTargetNeptune]:
return ScriptTranslatorTargetNeptune
def override_prepared_request_parameters(
self, request_parameters: RequestParameters, *, method: Optional[str] = None,
data: Optional[str] = None) -> httpclient.HTTPRequest:
http_request_param: Dict[str, Any] = dict(url=request_parameters.uri, headers=request_parameters.headers)
if method is not None:
http_request_param['method'] = method
if data is not None:
http_request_param['body'] = data
if self.override_uri:
# we override the URI slightly (because the instance thinks it's a different host than we're connecting to)
uri = urlsplit(request_parameters.uri)
http_request_param['headers'] = dict(request_parameters.headers)
http_request_param['headers']['Host'] = uri.netloc
http_request_param['ssl_options'] = OverrideServerHostnameSSLContext(server_hostname=uri.hostname)
http_request_param['url'] = urlunsplit(
(uri.scheme, self.override_uri.netloc, uri.path, uri.query, uri.fragment))
return httpclient.HTTPRequest(**http_request_param)
@overrides
def possibly_signed_ws_client_request_or_url(self) -> Union[httpclient.HTTPRequest, str]:
return self.override_prepared_request_parameters(self.endpoints.gremlin_endpoint().prepare_request())
@classmethod
@overrides
def _is_retryable_exception(cls, *, method_name: str, exception: Exception) -> bool:
# any method
return _is_neptune_retryable_exception(exception) or isinstance(exception, ConnectionError)
def is_healthy(self) -> None:
signed_request = self.override_prepared_request_parameters(self.endpoints.status_endpoint().prepare_request())
http_client = httpclient.HTTPClient()
# this will throw if the instance is really borked or we can't connect or we're not allowed (see
# https://docs.aws.amazon.com/neptune/latest/userguide/access-graph-status.html )
response = http_client.fetch(signed_request)
status = json.loads(response.body, encoding='utf-8')
if status.get('status') == 'healthy' and status.get('role') == 'writer':
LOGGER.debug(f'status is healthy: {status}')
else:
# we'll log in healthcheck
raise RuntimeError(f'status is unhealthy: {status}')
def _non_standard_endpoint(self, scheme: str, path: str) -> Endpoint:
return self.endpoints._Endpoints__endpoint(
scheme, self.endpoints.neptune_endpoint, self.endpoints.neptune_port, path)
def _gremlin_status(self, query_id: Optional[str] = None, include_waiting: bool = False) -> str:
"""
https://docs.aws.amazon.com/neptune/latest/userguide/gremlin-api-status.html
"""
endpoint = self._non_standard_endpoint('https', 'gremlin/status')
query_parameters = {}
if query_id is not None:
query_parameters['queryId'] = query_id
if include_waiting:
query_parameters['includeWaiting'] = 'true'
signed_request = self.override_prepared_request_parameters(
endpoint.prepare_request(querystring=query_parameters))
http_client = httpclient.HTTPClient()
response = http_client.fetch(signed_request)
return json.loads(response.body, encoding='utf-8')
def _sparql_status(self, query_id: Optional[str] = None) -> str:
"""
https://docs.aws.amazon.com/neptune/latest/userguide/sparql-api-status.html
"""
endpoint = self._non_standard_endpoint('https', 'sparql/status')
query_parameters = {}
if query_id is not None:
query_parameters['queryId'] = query_id
signed_request = self.override_prepared_request_parameters(
endpoint.prepare_request(querystring=query_parameters))
http_client = httpclient.HTTPClient()
response = http_client.fetch(signed_request)
return json.loads(response.body, encoding='utf-8')
def _explain(self, gremlin_query: str) -> str:
"""
return the Neptune specific explaination of the query
see https://docs.aws.amazon.com/neptune/latest/userguide/gremlin-explain-api.html
see https://docs.aws.amazon.com/neptune/latest/userguide/gremlin-explain-background.html
"""
# why not use endpoints? Despite the fact that it accepts a method and payload, it doesn't *actually* generate
# sufficient headers so we'll use requests for these since we can
url = urlsplit(self.endpoints.gremlin_endpoint().prepare_request().uri)
assert url.scheme in ('wss', 'ws') and url.path == '/gremlin' and not url.query and not url.fragment, \
f'url is not a Neptune ws url?: {url}'
_explain_url = urlunsplit(
('https' if url.scheme == 'wss' else 'http', url.netloc, url.path + '/explain', '', ''))
host = to_aws4_request_compatible_host(_explain_url)
if self.override_uri:
_explain_url = urlunsplit(
('https' if url.scheme == 'wss' else 'http', self.override_uri.netloc, url.path + '/explain', '', ''))
s = requests.Session()
s.mount('https://', HostHeaderSSLAdapter())
response = s.post(_explain_url, auth=self.aws_auth,
data=json.dumps(dict(gremlin=gremlin_query)).encode('utf-8'),
# include Host header
headers=dict(Host=host))
return response.content.decode('utf-8')
def _profile(self, gremlin_query: str) -> str:
"""
return the Neptune specific explaination of the RUNNING query. Now it can't return the result set, so the
utility is limited to cases where you can re-run this, or are running as a one off from console, or as a last
resort
see https://docs.aws.amazon.com/neptune/latest/userguide/gremlin-profile-api.htlm
see https://docs.aws.amazon.com/neptune/latest/userguide/gremlin-explain-background.html
"""
# why not use endpoints? Despite the fact that it accepts a method and payload, it doesn't *actually* generate
# sufficient headers so we'll use requests for these since we can
url = urlsplit(self.endpoints.gremlin_endpoint().prepare_request().uri)
assert url.scheme in ('wss', 'ws') and url.path == '/gremlin' and not url.query and not url.fragment, \
f'url is not a Neptune ws url?: {url}'
_profile_url = urlunsplit(
('https' if url.scheme == 'wss' else 'http', url.netloc, url.path + '/profile', '', ''))
host = to_aws4_request_compatible_host(_profile_url)
if self.override_uri:
_profile_url = urlunsplit(
('https' if url.scheme == 'wss' else 'http', self.override_uri.netloc, url.path + '/profile', '', ''))
s = requests.Session()
s.mount('https://', HostHeaderSSLAdapter())
response = s.post(_profile_url, auth=self.aws_auth,
data=json.dumps(dict(gremlin=gremlin_query)).encode('utf-8'),
# include Host header
headers=dict(Host=host))
return response.content.decode('utf-8')
@overrides
def drop(self) -> None:
test_shard = get_shard()
g = self.g.V()
if test_shard:
g = g.has(WellKnownProperties.TestShard.value.name, test_shard)
g = g.drop()
LOGGER.warning('DROPPING ALL NODES')
self.query_executor()(query=g, get=FromResultSet.iterate)
# we seem to mess this up easily
leftover = self.query_executor()(query=self.g.V().hasId(TextP.startingWith(test_shard)).id(),
get=FromResultSet.toList)
self.query_executor()(query=self.g.V().hasId(TextP.startingWith(test_shard)).drop(),
get=FromResultSet.iterate)
assert not leftover, f'we have some leftover: {leftover}'
LOGGER.warning('COMPLETED DROP OF ALL NODES')