Skip to content

Commit 7c64ae5

Browse files
michaelnchinmichaelnchin
committed
Set service-based Gremlin protocol defaults for module-generated configs
1 parent 9c176c0 commit 7c64ae5

File tree

2 files changed

+101
-13
lines changed

2 files changed

+101
-13
lines changed

‎src/graph_notebook/configuration/generate_config.py

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,8 @@
1515
DEFAULT_NEO4J_USERNAME, DEFAULT_NEO4J_PASSWORD, DEFAULT_NEO4J_DATABASE,
1616
NEPTUNE_CONFIG_HOST_IDENTIFIERS, is_allowed_neptune_host, false_str_variants,
1717
GRAPHSONV3_VARIANTS, GRAPHSONV2_VARIANTS, GRAPHBINARYV1_VARIANTS,
18-
NEPTUNE_DB_SERVICE_NAME, normalize_service_name)
18+
NEPTUNE_DB_SERVICE_NAME, NEPTUNE_ANALYTICS_SERVICE_NAME,
19+
normalize_service_name)
1920

2021
DEFAULT_CONFIG_LOCATION = os.path.expanduser('~/graph_notebook_config.json')
2122

@@ -299,7 +300,7 @@ def generate_default_config():
299300
default=DEFAULT_GREMLIN_SERIALIZER)
300301
parser.add_argument("--gremlin_connection_protocol",
301302
help="the connection protocol to use for Gremlin connections",
302-
default=DEFAULT_GREMLIN_PROTOCOL)
303+
default='')
303304
parser.add_argument("--neo4j_username", help="the username to use for Neo4J connections",
304305
default=DEFAULT_NEO4J_USERNAME)
305306
parser.add_argument("--neo4j_password", help="the password to use for Neo4J connections",
@@ -311,6 +312,13 @@ def generate_default_config():
311312
args = parser.parse_args()
312313

313314
auth_mode_arg = args.auth_mode if args.auth_mode != '' else AuthModeEnum.DEFAULT.value
315+
protocol_arg = args.gremlin_connection_protocol
316+
include_protocol = False
317+
if is_allowed_neptune_host(args.host, args.neptune_hosts):
318+
include_protocol = True
319+
if not protocol_arg:
320+
protocol_arg = DEFAULT_HTTP_PROTOCOL \
321+
if args.neptune_service == NEPTUNE_ANALYTICS_SERVICE_NAME else DEFAULT_WS_PROTOCOL
314322
config = generate_config(args.host, int(args.port),
315323
AuthModeEnum(auth_mode_arg),
316324
args.ssl, args.ssl_verify,
@@ -319,7 +327,7 @@ def generate_default_config():
319327
SparqlSection(args.sparql_path, ''),
320328
GremlinSection(args.gremlin_traversal_source, args.gremlin_username,
321329
args.gremlin_password, args.gremlin_serializer,
322-
args.gremlin_connection_protocol),
330+
protocol_arg, include_protocol),
323331
Neo4JSection(args.neo4j_username, args.neo4j_password,
324332
args.neo4j_auth, args.neo4j_database),
325333
args.neptune_hosts)

‎test/unit/configuration/test_configuration_from_main.py

Lines changed: 90 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88

99
from graph_notebook.configuration.generate_config import AuthModeEnum, Configuration, GremlinSection
1010
from graph_notebook.configuration.get_config import get_config
11+
from graph_notebook.neptune.client import (NEPTUNE_DB_SERVICE_NAME, NEPTUNE_ANALYTICS_SERVICE_NAME,
12+
DEFAULT_HTTP_PROTOCOL, DEFAULT_WS_PROTOCOL)
1113

1214

1315
class TestGenerateConfigurationMain(unittest.TestCase):
@@ -49,25 +51,61 @@ def test_generate_configuration_main_defaults_generic(self):
4951
self.generate_config_from_main_and_test(expected_config)
5052

5153
def test_generate_configuration_main_override_defaults_neptune_reg(self):
52-
expected_config = Configuration(self.neptune_host_reg, self.port, neptune_service='neptune-graph',
53-
auth_mode=AuthModeEnum.IAM, load_from_s3_arn='loader_arn', ssl=False)
54+
expected_config = Configuration(self.neptune_host_reg,
55+
self.port,
56+
neptune_service='neptune-graph',
57+
auth_mode=AuthModeEnum.IAM,
58+
load_from_s3_arn='loader_arn',
59+
ssl=False,
60+
gremlin_section=GremlinSection(
61+
connection_protocol=DEFAULT_HTTP_PROTOCOL,
62+
include_protocol=True
63+
)
64+
)
5465
self.generate_config_from_main_and_test(expected_config, host_type='neptune')
5566

5667
def test_generate_configuration_main_override_defaults_neptune_no_verify(self):
57-
expected_config = Configuration(self.neptune_host_reg, self.port, neptune_service='neptune-graph',
58-
auth_mode=AuthModeEnum.IAM, load_from_s3_arn='loader_arn',
59-
ssl=True, ssl_verify=False)
68+
expected_config = Configuration(self.neptune_host_reg,
69+
self.port,
70+
neptune_service='neptune-graph',
71+
auth_mode=AuthModeEnum.IAM,
72+
load_from_s3_arn='loader_arn',
73+
ssl=True,
74+
ssl_verify=False,
75+
gremlin_section=GremlinSection(
76+
connection_protocol=DEFAULT_HTTP_PROTOCOL,
77+
include_protocol=True
78+
)
79+
)
6080
self.generate_config_from_main_and_test(expected_config, host_type='neptune')
6181

6282
def test_generate_configuration_main_override_defaults_neptune_with_serializer(self):
63-
expected_config = Configuration(self.neptune_host_reg, self.port, neptune_service='neptune-graph',
64-
auth_mode=AuthModeEnum.IAM, load_from_s3_arn='loader_arn', ssl=False,
65-
gremlin_section=GremlinSection(message_serializer='graphbinary'))
83+
expected_config = Configuration(self.neptune_host_reg,
84+
self.port,
85+
neptune_service='neptune-graph',
86+
auth_mode=AuthModeEnum.IAM,
87+
load_from_s3_arn='loader_arn',
88+
ssl=False,
89+
gremlin_section=GremlinSection(
90+
message_serializer='graphbinary',
91+
connection_protocol=DEFAULT_HTTP_PROTOCOL,
92+
include_protocol=True
93+
)
94+
)
6695
self.generate_config_from_main_and_test(expected_config, host_type='neptune')
6796

6897
def test_generate_configuration_main_override_defaults_neptune_cn(self):
69-
expected_config = Configuration(self.neptune_host_cn, self.port, neptune_service='neptune-graph',
70-
auth_mode=AuthModeEnum.IAM, load_from_s3_arn='loader_arn', ssl=False)
98+
expected_config = Configuration(self.neptune_host_cn,
99+
self.port,
100+
neptune_service='neptune-graph',
101+
auth_mode=AuthModeEnum.IAM,
102+
load_from_s3_arn='loader_arn',
103+
ssl=False,
104+
gremlin_section=GremlinSection(
105+
connection_protocol=DEFAULT_HTTP_PROTOCOL,
106+
include_protocol=True
107+
)
108+
)
71109
self.generate_config_from_main_and_test(expected_config, host_type='neptune')
72110

73111
def test_generate_configuration_main_override_defaults_generic(self):
@@ -85,6 +123,48 @@ def test_generate_configuration_main_empty_args_neptune(self):
85123
config = get_config(self.test_file_path)
86124
self.assertEqual(expected_config.to_dict(), config.to_dict())
87125

126+
def test_generate_configuration_main_gremlin_protocol_no_service(self):
127+
result = os.system(f'{self.python_cmd} -m graph_notebook.configuration.generate_config '
128+
f'--host "{self.neptune_host_reg}" '
129+
f'--port "{self.port}" '
130+
f'--neptune_service "" '
131+
f'--auth_mode "" '
132+
f'--ssl "" '
133+
f'--load_from_s3_arn "" '
134+
f'--config_destination="{self.test_file_path}" ')
135+
self.assertEqual(0, result)
136+
config = get_config(self.test_file_path)
137+
config_dict = config.to_dict()
138+
self.assertEqual(DEFAULT_WS_PROTOCOL, config_dict['gremlin']['connection_protocol'])
139+
140+
def test_generate_configuration_main_gremlin_protocol_db(self):
141+
result = os.system(f'{self.python_cmd} -m graph_notebook.configuration.generate_config '
142+
f'--host "{self.neptune_host_reg}" '
143+
f'--port "{self.port}" '
144+
f'--neptune_service "{NEPTUNE_DB_SERVICE_NAME}" '
145+
f'--auth_mode "" '
146+
f'--ssl "" '
147+
f'--load_from_s3_arn "" '
148+
f'--config_destination="{self.test_file_path}" ')
149+
self.assertEqual(0, result)
150+
config = get_config(self.test_file_path)
151+
config_dict = config.to_dict()
152+
self.assertEqual(DEFAULT_WS_PROTOCOL, config_dict['gremlin']['connection_protocol'])
153+
154+
def test_generate_configuration_main_gremlin_protocol_analytics(self):
155+
result = os.system(f'{self.python_cmd} -m graph_notebook.configuration.generate_config '
156+
f'--host "{self.neptune_host_reg}" '
157+
f'--port "{self.port}" '
158+
f'--neptune_service "{NEPTUNE_ANALYTICS_SERVICE_NAME}" '
159+
f'--auth_mode "" '
160+
f'--ssl "" '
161+
f'--load_from_s3_arn "" '
162+
f'--config_destination="{self.test_file_path}" ')
163+
self.assertEqual(0, result)
164+
config = get_config(self.test_file_path)
165+
config_dict = config.to_dict()
166+
self.assertEqual(DEFAULT_HTTP_PROTOCOL, config_dict['gremlin']['connection_protocol'])
167+
88168
def test_generate_configuration_main_empty_args_custom(self):
89169
expected_config = Configuration(self.neptune_host_custom, self.port, neptune_hosts=self.custom_hosts_list)
90170
result = os.system(f'{self.python_cmd} -m graph_notebook.configuration.generate_config '

0 commit comments

Comments
 (0)