def get_config_from_dict(data: dict) -> Configuration:
sparql_section = SparqlSection(
**data['sparql']) if 'sparql' in data else SparqlSection('')
gremlin_section = GremlinSection(
**data['gremlin']) if 'gremlin' in data else GremlinSection('')
if "amazonaws.com" in data['host']:
if gremlin_section.to_dict()['traversal_source'] != 'g':
print(
'Ignoring custom traversal source, Amazon Neptune does not support this functionality.\n'
)
config = Configuration(host=data['host'],
port=data['port'],
auth_mode=AuthModeEnum(data['auth_mode']),
ssl=data['ssl'],
load_from_s3_arn=data['load_from_s3_arn'],
aws_region=data['aws_region'],
sparql_section=sparql_section,
gremlin_section=gremlin_section)
else:
config = Configuration(host=data['host'],
port=data['port'],
ssl=data['ssl'],
sparql_section=sparql_section,
gremlin_section=gremlin_section)
return config
def test_generate_configuration_override_defaults_generic(self):
ssl = False
config = Configuration(self.generic_host, self.port, ssl=ssl)
c = generate_config(config.host, config.port, ssl=config.ssl)
c.write_to_file(self.test_file_path)
config_from_file = get_config(self.test_file_path)
self.assertEqual(config.to_dict(), config_from_file.to_dict())
def test_generate_configuration_with_defaults_neptune(self):
config = Configuration(self.neptune_host, self.port)
c = generate_config(config.host, config.port, auth_mode=config.auth_mode, ssl=config.ssl,
load_from_s3_arn=config.load_from_s3_arn, aws_region=config.aws_region)
c.write_to_file(self.test_file_path)
config_from_file = get_config(self.test_file_path)
self.assertEqual(config.to_dict(), config_from_file.to_dict())
def test_generate_configuration_with_defaults(self):
config = Configuration(self.host, self.port)
c = generate_config(config.host, config.port, config.auth_mode, config.ssl,
config.iam_credentials_provider_type,
config.load_from_s3_arn, config.aws_region)
c.write_to_file(self.test_file_path)
config_from_file = get_config(self.test_file_path)
self.assertEqual(config.to_dict(), config_from_file.to_dict())
def test_generate_configuration_main_empty_args_generic(self):
expected_config = Configuration(self.generic_host, self.port)
result = os.system(
f'{self.python_cmd} -m graph_notebook.configuration.generate_config '
f'--host "{expected_config.host}" --port "{expected_config.port}" --ssl "" '
f'--config_destination="{self.test_file_path}" ')
self.assertEqual(0, result)
config = get_config(self.test_file_path)
self.assertEqual(expected_config.to_dict(), config.to_dict())
def test_generate_configuration_override_defaults_neptune(self):
auth_mode = AuthModeEnum.IAM
ssl = False
loader_arn = 'foo'
aws_region = 'us-west-2'
config = Configuration(self.neptune_host, self.port, auth_mode=auth_mode, load_from_s3_arn=loader_arn, ssl=ssl,
aws_region=aws_region)
c = generate_config(config.host, config.port, auth_mode=config.auth_mode, ssl=config.ssl,
load_from_s3_arn=config.load_from_s3_arn, aws_region=config.aws_region)
c.write_to_file(self.test_file_path)
config_from_file = get_config(self.test_file_path)
self.assertEqual(config.to_dict(), config_from_file.to_dict())
def test_generate_configuration_override_defaults(self):
auth_mode = AuthModeEnum.IAM
credentials_provider = IAMAuthCredentialsProvider.ENV
ssl = False
loader_arn = 'foo'
aws_region = 'us-west-2'
config = Configuration(self.host, self.port, auth_mode, credentials_provider, loader_arn, ssl, aws_region)
c = generate_config(config.host, config.port, config.auth_mode, config.ssl,
config.iam_credentials_provider_type,
config.load_from_s3_arn, config.aws_region)
c.write_to_file(self.test_file_path)
config_from_file = get_config(self.test_file_path)
self.assertEqual(config.to_dict(), config_from_file.to_dict())
def test_configuration_default_auth_defaults_neptune(self):
config = Configuration(self.neptune_host, self.port)
self.assertEqual(self.neptune_host, config.host)
self.assertEqual(self.port, config.port)
self.assertEqual(DEFAULT_AUTH_MODE, config.auth_mode)
self.assertEqual(True, config.ssl)
self.assertEqual('', config.load_from_s3_arn)
def generate_config_from_stack(stack: dict, region: str,
iam: bool) -> Configuration:
logging.info('stack finished, retrieving stack details')
details = get_stack_details_to_run(stack, region)
logging.info(f'obtained stack details: {details}')
# because this stack puts Neptune behind a load balancer, we need to alias the ip address of the load
# balancer to the Neptune cluster endpoint to ensure that SSL certificates match the host that we are calling.
# To do this will we need to alter /etc/hosts
host_alias = f'{details["ip"]} {details["endpoint"]}'
logging.info(
f'adding {host_alias} to /etc/hosts and removing other aliased lines to cluster endpoint {details["endpoint"]}'
)
new_lines = []
with open('/etc/hosts', 'w+') as file:
lines = file.read().split('\n')
for line in lines:
if details['endpoint'] not in line:
new_lines.append(line)
new_lines.append(host_alias + '\n')
file.writelines(new_lines)
auth = AuthModeEnum.IAM if iam else AuthModeEnum.DEFAULT
conf = Configuration(details['endpoint'],
80,
auth,
details['loader_arn'],
ssl=True,
aws_region=region)
logging.info(f'generated configuration for test run: {conf.to_dict()}')
return conf
def test_configuration_override_defaults_neptune(self):
auth_mode = AuthModeEnum.IAM
ssl = False
loader_arn = 'foo'
config = Configuration(self.neptune_host, self.port, auth_mode=auth_mode, load_from_s3_arn=loader_arn, ssl=ssl)
self.assertEqual(auth_mode, config.auth_mode)
self.assertEqual(ssl, config.ssl)
self.assertEqual(loader_arn, config.load_from_s3_arn)
def get_config_from_dict(data: dict) -> Configuration:
sparql_section = SparqlSection(
**data['sparql']) if 'sparql' in data else SparqlSection('')
if ".neptune.amazonaws.com" in data['host']:
config = Configuration(host=data['host'],
port=data['port'],
auth_mode=AuthModeEnum(data['auth_mode']),
ssl=data['ssl'],
load_from_s3_arn=data['load_from_s3_arn'],
aws_region=data['aws_region'],
sparql_section=sparql_section)
else:
config = Configuration(host=data['host'],
port=data['port'],
ssl=data['ssl'],
sparql_section=sparql_section)
return config
def test_generate_configuration_main_empty_args(self):
expected_config = Configuration(self.host, self.port)
result = os.system(
f'{self.python_cmd} -m graph_notebook.configuration.generate_config --host "{expected_config.host}" --port "{expected_config.port}" --auth_mode "" --ssl "" --iam_credentials_provider "" --load_from_s3_arn "" --config_destination="{self.test_file_path}" '
)
self.assertEqual(0, result)
config = get_config(self.test_file_path)
self.assert_configs_are_equal(expected_config, config)
def generate_config_from_main_and_test(self, source_config: Configuration):
# This will run the main method that our install script runs on a Sagemaker notebook.
# The return code should be 0, but more importantly, we need to assert that the
# Configuration object we get from the resulting file is what we expect.
result = os.system(f'{self.python_cmd} -m graph_notebook.configuration.generate_config --host "{source_config.host}" --port "{source_config.port}" --auth_mode "{source_config.auth_mode.value}" --ssl "{source_config.ssl}" --iam_credentials_provider "{source_config.iam_credentials_provider_type.value}" --load_from_s3_arn "{source_config.load_from_s3_arn}" --config_destination="{self.test_file_path}" ')
self.assertEqual(result, 0)
config = get_config(self.test_file_path)
self.assertEqual(source_config.to_dict(), config.to_dict())
def test_configuration_default_auth_defaults(self):
config = Configuration(self.host, self.port)
self.assertEqual(self.host, config.host)
self.assertEqual(self.port, config.port)
self.assertEqual(DEFAULT_AUTH_MODE, config.auth_mode)
self.assertEqual(DEFAULT_IAM_CREDENTIALS_PROVIDER, config.iam_credentials_provider_type)
self.assertEqual(True, config.ssl)
self.assertEqual('', config.load_from_s3_arn)
def test_generate_configuration_main_override_defaults_neptune(self):
expected_config = Configuration(self.neptune_host,
self.port,
auth_mode=AuthModeEnum.IAM,
load_from_s3_arn='loader_arn',
ssl=False)
self.generate_config_from_main_and_test(expected_config,
host_type='neptune')
def test_generate_configuration_main_defaults_neptune(self):
expected_config = Configuration(self.neptune_host,
self.port,
auth_mode=AuthModeEnum.DEFAULT,
load_from_s3_arn='',
ssl=True)
self.generate_config_from_main_and_test(expected_config,
host_type='neptune')
def test_configuration_override_defaults(self):
auth_mode = AuthModeEnum.IAM
credentials_provider = IAMAuthCredentialsProvider.ENV
ssl = False
loader_arn = 'foo'
config = Configuration(self.host, self.port, auth_mode, credentials_provider, loader_arn, ssl)
self.assertEqual(auth_mode, config.auth_mode)
self.assertEqual(credentials_provider, config.iam_credentials_provider_type)
self.assertEqual(ssl, config.ssl)
self.assertEqual(loader_arn, config.load_from_s3_arn)
def get_config_from_dict(data: dict) -> Configuration:
config = Configuration(
host=data['host'],
port=data['port'],
auth_mode=AuthModeEnum(data['auth_mode']),
iam_credentials_provider_type=IAMAuthCredentialsProvider(
data['iam_credentials_provider_type']),
ssl=data['ssl'],
load_from_s3_arn=data['load_from_s3_arn'],
aws_region=data['aws_region'])
return config
def get_config_from_dict(data: dict) -> Configuration:
sparql_section = SparqlSection(
**data['sparql']) if 'sparql' in data else SparqlSection('')
config = Configuration(
host=data['host'],
port=data['port'],
auth_mode=AuthModeEnum(data['auth_mode']),
iam_credentials_provider_type=IAMAuthCredentialsProvider(
data['iam_credentials_provider_type']),
ssl=data['ssl'],
load_from_s3_arn=data['load_from_s3_arn'],
aws_region=data['aws_region'],
sparql_section=sparql_section)
return config
def test_configuration_override_defaults_generic(self):
ssl = False
config = Configuration(self.generic_host, self.port, ssl=ssl)
self.assertEqual(ssl, config.ssl)
def test_generate_configuration_main_override_defaults(self):
expected_config = Configuration(self.host, self.port, AuthModeEnum.IAM, IAMAuthCredentialsProvider.ROLE, 'loader_arn', False)
self.generate_config_from_main_and_test(expected_config)
def test_generate_configuration_main_defaults(self):
expected_config = Configuration(self.host, self.port, AuthModeEnum.DEFAULT, IAMAuthCredentialsProvider.ROLE, '', True)
self.generate_config_from_main_and_test(expected_config)
def test_generate_configuration_main_defaults(self):
expected_config = Configuration(self.host, self.port, AuthModeEnum.DEFAULT, '', True)
self.generate_config_from_main_and_test(expected_config)
def test_generate_configuration_main_override_defaults_generic(self):
expected_config = Configuration(self.generic_host,
self.port,
ssl=False)
self.generate_config_from_main_and_test(expected_config)
def test_configuration_default_auth_defaults_generic(self):
config = Configuration(self.generic_host, self.port)
self.assertEqual(self.generic_host, config.host)
self.assertEqual(self.port, config.port)
self.assertEqual(True, config.ssl)