Skip to content

Commit 572ef94

Browse files
michaelnchinmichaelnchin
authored
Fix %%graph_notebook_config error when excluding optional Gremlin section (#633)
* Fix %%graph_notebook_config exception when excluding optional Gremlin section * Linter fix * update changelog
1 parent f723b3d commit 572ef94

File tree

3 files changed

+228
-5
lines changed

3 files changed

+228
-5
lines changed

‎ChangeLog.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ Starting with v1.31.6, this file will contain a record of major features and upd
55
## Upcoming
66
- Fixed broken `--help` option for `%%gremlin` ([Link to PR](https://github.com/aws/graph-notebook/pull/630))
77
- Fixed openCypher query bug regression in the [`01-About-the-Neptune-Notebook`](https://github.com/aws/graph-notebook/blob/main/src/graph_notebook/notebooks/01-Getting-Started/01-About-the-Neptune-Notebook.ipynb) sample ([Link to PR](https://github.com/aws/graph-notebook/pull/631))
8+
- Fixed `%%graph_notebook_config` error when excluding optional Gremlin section ([Link to PR](https://github.com/aws/graph-notebook/pull/633))
89

910
## Release 4.4.2 (June 18, 2024)
1011
- Set Gremlin `connection_protocol` defaults based on Neptune service when generating configuration via arguments ([Link to PR](https://github.com/aws/graph-notebook/pull/626))

‎src/graph_notebook/configuration/get_config.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -59,9 +59,10 @@ def get_config_from_dict(data: dict, neptune_hosts: list = NEPTUNE_CONFIG_HOST_I
5959
for p in neptune_params:
6060
if p in data:
6161
excluded_params.append(p)
62-
for gp in neptune_gremlin_params:
63-
if gp in data['gremlin']:
64-
excluded_params.append(gp)
62+
if 'gremlin' in data:
63+
for gp in neptune_gremlin_params:
64+
if gp in data['gremlin']:
65+
excluded_params.append(gp)
6566
if excluded_params:
6667
print(f"The provided configuration contains the following parameters that are incompatible with the "
6768
f"specified host: {str(excluded_params)}. These parameters have not been saved.\n")

‎test/unit/configuration/test_configuration.py‎

Lines changed: 223 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,11 @@
66
import os
77
import unittest
88

9-
from graph_notebook.configuration.get_config import get_config
9+
from graph_notebook.configuration.get_config import get_config, get_config_from_dict
1010
from graph_notebook.configuration.generate_config import Configuration, DEFAULT_AUTH_MODE, AuthModeEnum, \
1111
generate_config, generate_default_config, GremlinSection
1212
from graph_notebook.neptune.client import NEPTUNE_DB_SERVICE_NAME, NEPTUNE_ANALYTICS_SERVICE_NAME, \
13-
DEFAULT_GREMLIN_PROTOCOL, DEFAULT_HTTP_PROTOCOL
13+
DEFAULT_GREMLIN_PROTOCOL, DEFAULT_HTTP_PROTOCOL, NEPTUNE_CONFIG_HOST_IDENTIFIERS
1414

1515

1616
class TestGenerateConfiguration(unittest.TestCase):
@@ -121,6 +121,227 @@ def test_configuration_override_defaults_generic(self):
121121
config = Configuration(self.generic_host, self.port, ssl=ssl)
122122
self.assertEqual(ssl, config.ssl)
123123

124+
def test_get_configuration_empty_input(self):
125+
input_config = {}
126+
with self.assertRaises(KeyError):
127+
get_config_from_dict(input_config, neptune_hosts=NEPTUNE_CONFIG_HOST_IDENTIFIERS)
128+
129+
def test_get_configuration_no_host(self):
130+
input_config = {
131+
"port": 8182,
132+
"ssl": True
133+
}
134+
with self.assertRaises(KeyError):
135+
get_config_from_dict(input_config, neptune_hosts=NEPTUNE_CONFIG_HOST_IDENTIFIERS)
136+
137+
def test_get_configuration_generic_no_port(self):
138+
input_config = {
139+
"host": "localhost",
140+
"ssl": True
141+
}
142+
with self.assertRaises(KeyError):
143+
get_config_from_dict(input_config, neptune_hosts=NEPTUNE_CONFIG_HOST_IDENTIFIERS)
144+
145+
def test_get_configuration_generic_no_ssl(self):
146+
input_config = {
147+
"host": "localhost",
148+
"port": 8182
149+
}
150+
with self.assertRaises(KeyError):
151+
get_config_from_dict(input_config, neptune_hosts=NEPTUNE_CONFIG_HOST_IDENTIFIERS)
152+
153+
def test_get_configuration_generic_required_input(self):
154+
input_config = {
155+
"host": "localhost",
156+
"port": 8182,
157+
"ssl": True
158+
}
159+
expected_config = {
160+
'host': 'localhost',
161+
'port': 8182,
162+
'proxy_host': '',
163+
'proxy_port': 8182,
164+
'ssl': True,
165+
'ssl_verify': True,
166+
'sparql': {
167+
'path': ''
168+
},
169+
'gremlin': {
170+
'traversal_source': 'g',
171+
'username': '',
172+
'password': '',
173+
'message_serializer': 'graphsonv3'
174+
},
175+
'neo4j': {
176+
'username': 'neo4j',
177+
'password': 'password',
178+
'auth': True,
179+
'database': None
180+
}
181+
}
182+
config = get_config_from_dict(input_config, neptune_hosts=NEPTUNE_CONFIG_HOST_IDENTIFIERS)
183+
self.assertEqual(config.to_dict(), expected_config)
184+
185+
def test_get_configuration_generic_all_input(self):
186+
input_and_expected_config = {
187+
'host': 'a_host',
188+
'port': 9999,
189+
'proxy_host': 'a_proxy_host',
190+
'proxy_port': 9999,
191+
'ssl': False,
192+
'ssl_verify': False,
193+
'sparql': {
194+
'path': 'a_path'
195+
},
196+
'gremlin': {
197+
'traversal_source': 'a',
198+
'username': 'user',
199+
'password': 'pass',
200+
'message_serializer': 'graphbinaryv1'
201+
},
202+
'neo4j': {
203+
'username': 'neo_user',
204+
'password': 'neo_pass',
205+
'auth': False,
206+
'database': 'neo_db'
207+
}
208+
}
209+
config = get_config_from_dict(input_and_expected_config, neptune_hosts=NEPTUNE_CONFIG_HOST_IDENTIFIERS)
210+
self.assertEqual(config.to_dict(), input_and_expected_config)
211+
212+
def test_get_configuration_neptune_no_auth_mode(self):
213+
input_config = {
214+
"host": "db.cluster-xxxxxxxxx.us-west-2.neptune.amazonaws.com",
215+
"port": 8182,
216+
"ssl": True,
217+
"load_from_s3_arn": "",
218+
"aws_region": "us-west-2"
219+
}
220+
with self.assertRaises(KeyError):
221+
get_config_from_dict(input_config, neptune_hosts=NEPTUNE_CONFIG_HOST_IDENTIFIERS)
222+
223+
def test_get_configuration_neptune_no_load_arn(self):
224+
input_config = {
225+
"host": "db.cluster-xxxxxxxxx.us-west-2.neptune.amazonaws.com",
226+
"port": 8182,
227+
"ssl": True,
228+
"aws_region": "us-west-2"
229+
}
230+
with self.assertRaises(KeyError):
231+
get_config_from_dict(input_config, neptune_hosts=NEPTUNE_CONFIG_HOST_IDENTIFIERS)
232+
233+
def test_get_configuration_neptune_no_region(self):
234+
input_config = {
235+
"host": "db.cluster-xxxxxxxxx.us-west-2.neptune.amazonaws.com",
236+
"port": 8182,
237+
"ssl": True,
238+
"load_from_s3_arn": ""
239+
}
240+
with self.assertRaises(KeyError):
241+
get_config_from_dict(input_config, neptune_hosts=NEPTUNE_CONFIG_HOST_IDENTIFIERS)
242+
243+
def test_get_configuration_neptune_required_input(self):
244+
input_config = {
245+
"host": "db.cluster-xxxxxxxxx.us-west-2.neptune.amazonaws.com",
246+
"port": 8182,
247+
"auth_mode": "IAM",
248+
"load_from_s3_arn": "",
249+
"ssl": True,
250+
"aws_region": "us-west-2"
251+
}
252+
expected_config = {
253+
'host': 'db.cluster-xxxxxxxxx.us-west-2.neptune.amazonaws.com',
254+
'neptune_service': 'neptune-db',
255+
'port': 8182,
256+
'proxy_host': '',
257+
'proxy_port': 8182,
258+
'auth_mode': 'IAM',
259+
'load_from_s3_arn': '',
260+
'ssl': True,
261+
'ssl_verify': True,
262+
'aws_region': 'us-west-2',
263+
'sparql': {
264+
'path': ''
265+
},
266+
'gremlin': {
267+
'traversal_source': 'g',
268+
'username': '',
269+
'password': '',
270+
'message_serializer': 'graphsonv3',
271+
'connection_protocol': 'websockets'
272+
},
273+
'neo4j': {
274+
'username': 'neo4j',
275+
'password': 'password',
276+
'auth': True,
277+
'database': None
278+
}
279+
}
280+
281+
config = get_config_from_dict(input_config, neptune_hosts=NEPTUNE_CONFIG_HOST_IDENTIFIERS)
282+
self.assertEqual(config.to_dict(), expected_config)
283+
284+
def test_get_configuration_neptune_all_input(self):
285+
input_config = {
286+
'host': 'db.cluster-xxxxxxxxx.us-west-2.neptune.amazonaws.com',
287+
'neptune_service': 'neptune-graph',
288+
'port': 9999,
289+
'proxy_host': 'a_proxy+port',
290+
'proxy_port': 9999,
291+
'auth_mode': 'DEFAULT',
292+
'load_from_s3_arn': 'a_role',
293+
'ssl': False,
294+
'ssl_verify': False,
295+
'aws_region': 'us-west-2',
296+
'sparql': {
297+
'path': 'a_path'
298+
},
299+
'gremlin': {
300+
'traversal_source': 'a',
301+
'username': 'a_user',
302+
'password': 'a_pass',
303+
'message_serializer': 'graphbinaryv1',
304+
'connection_protocol': 'http'
305+
},
306+
'neo4j': {
307+
'username': 'a_user',
308+
'password': 'a_pass',
309+
'auth': False,
310+
'database': 'a_db'
311+
}
312+
}
313+
expected_config = {
314+
'host': 'db.cluster-xxxxxxxxx.us-west-2.neptune.amazonaws.com',
315+
'neptune_service': 'neptune-graph',
316+
'port': 9999,
317+
'proxy_host': 'a_proxy+port',
318+
'proxy_port': 9999,
319+
'auth_mode': 'DEFAULT',
320+
'load_from_s3_arn': 'a_role',
321+
'ssl': False,
322+
'ssl_verify': False,
323+
'aws_region': 'us-west-2',
324+
'sparql': {
325+
'path': 'a_path'
326+
},
327+
'gremlin': {
328+
'traversal_source': 'g',
329+
'username': '',
330+
'password': '',
331+
'message_serializer': 'graphbinaryv1',
332+
'connection_protocol': 'http'
333+
},
334+
'neo4j': {
335+
'username': 'neo4j',
336+
'password': 'password',
337+
'auth': True,
338+
'database': None
339+
}
340+
}
341+
342+
config = get_config_from_dict(input_config, neptune_hosts=NEPTUNE_CONFIG_HOST_IDENTIFIERS)
343+
self.assertEqual(config.to_dict(), expected_config)
344+
124345
def test_generate_configuration_with_defaults_neptune_reg(self):
125346
config = Configuration(self.neptune_host_reg, self.port)
126347
c = generate_config(config.host, config.port, auth_mode=config.auth_mode, ssl=config.ssl,

0 commit comments

Comments
 (0)