From d0706d31242bd6e542d7d4c112ae3b065cbbda00 Mon Sep 17 00:00:00 2001 From: Michael Chin Date: Fri, 8 Dec 2023 19:20:22 -0800 Subject: [PATCH 1/2] Allow custom Gremlin serializers for Neptune --- .../configuration/generate_config.py | 3 +- .../configuration/get_config.py | 1 - test/unit/configuration/test_configuration.py | 43 ++++++++++++++++++- .../test_configuration_from_main.py | 11 ++++- 4 files changed, 53 insertions(+), 5 deletions(-) diff --git a/src/graph_notebook/configuration/generate_config.py b/src/graph_notebook/configuration/generate_config.py index 19a1a2ba..4619af25 100644 --- a/src/graph_notebook/configuration/generate_config.py +++ b/src/graph_notebook/configuration/generate_config.py @@ -141,7 +141,8 @@ def __init__(self, host: str, port: int, self.auth_mode = auth_mode self.load_from_s3_arn = load_from_s3_arn self.aws_region = aws_region - self.gremlin = GremlinSection() + self.gremlin = GremlinSection(message_serializer=gremlin_section.message_serializer) \ + if gremlin_section is not None else GremlinSection() self.neo4j = Neo4JSection() else: self.is_neptune_config = False diff --git a/src/graph_notebook/configuration/get_config.py b/src/graph_notebook/configuration/get_config.py index 719f1df4..334b3a3b 100644 --- a/src/graph_notebook/configuration/get_config.py +++ b/src/graph_notebook/configuration/get_config.py @@ -13,7 +13,6 @@ neptune_params = ['neptune_service', 'auth_mode', 'load_from_s3_arn', 'aws_region'] -neptune_params = ['auth_mode', 'load_from_s3_arn', 'aws_region'] def get_config_from_dict(data: dict, neptune_hosts: list = NEPTUNE_CONFIG_HOST_IDENTIFIERS) -> Configuration: diff --git a/test/unit/configuration/test_configuration.py b/test/unit/configuration/test_configuration.py index 59a9fef7..d295b015 100644 --- a/test/unit/configuration/test_configuration.py +++ b/test/unit/configuration/test_configuration.py @@ -7,7 +7,8 @@ import unittest from graph_notebook.configuration.get_config import get_config -from graph_notebook.configuration.generate_config import Configuration, DEFAULT_AUTH_MODE, AuthModeEnum, generate_config +from graph_notebook.configuration.generate_config import Configuration, DEFAULT_AUTH_MODE, AuthModeEnum, \ + generate_config, GremlinSection from graph_notebook.neptune.client import NEPTUNE_DB_SERVICE_NAME, NEPTUNE_ANALYTICS_SERVICE_NAME @@ -219,6 +220,46 @@ def test_configuration_neptune_proxy_host_with_whitespace_using_setter(self): self.assertEqual(config.proxy_host, self.neptune_host_reg) self.assertEqual(config._proxy_host, self.neptune_host_reg) + def test_configuration_gremlinsection_generic_default(self): + config = Configuration('localhost', self.port) + self.assertEqual(config.gremlin.traversal_source, 'g') + self.assertEqual(config.gremlin.username, '') + self.assertEqual(config.gremlin.password, '') + self.assertEqual(config.gremlin.message_serializer, 'graphsonv3') + + def test_configuration_gremlinsection_generic_override(self): + config = Configuration('localhost', + self.port, + gremlin_section=GremlinSection(traversal_source='t', + username='foo', + password='bar', + message_serializer='graphbinary'), + ) + self.assertEqual(config.gremlin.traversal_source, 't') + self.assertEqual(config.gremlin.username, 'foo') + self.assertEqual(config.gremlin.password, 'bar') + self.assertEqual(config.gremlin.message_serializer, 'graphbinaryv1') + + def test_configuration_gremlinsection_neptune_default(self): + config = Configuration(self.neptune_host_reg, self.port) + self.assertEqual(config.gremlin.traversal_source, 'g') + self.assertEqual(config.gremlin.username, '') + self.assertEqual(config.gremlin.password, '') + self.assertEqual(config.gremlin.message_serializer, 'graphsonv3') + + def test_configuration_gremlinsection_neptune_override(self): + config = Configuration(self.neptune_host_reg, + self.port, + gremlin_section=GremlinSection(traversal_source='t', + username='foo', + password='bar', + message_serializer='graphbinary'), + ) + self.assertEqual(config.gremlin.traversal_source, 'g') + self.assertEqual(config.gremlin.username, '') + self.assertEqual(config.gremlin.password, '') + self.assertEqual(config.gremlin.message_serializer, 'graphbinaryv1') + def test_configuration_neptune_service_default(self): config = Configuration(self.neptune_host_reg, self.port) self.assertEqual(config.neptune_service, self.neptune_service_db) diff --git a/test/unit/configuration/test_configuration_from_main.py b/test/unit/configuration/test_configuration_from_main.py index 275948f9..66293589 100644 --- a/test/unit/configuration/test_configuration_from_main.py +++ b/test/unit/configuration/test_configuration_from_main.py @@ -6,7 +6,7 @@ import os import unittest -from graph_notebook.configuration.generate_config import AuthModeEnum, Configuration +from graph_notebook.configuration.generate_config import AuthModeEnum, Configuration, GremlinSection from graph_notebook.configuration.get_config import get_config @@ -59,6 +59,12 @@ def test_generate_configuration_main_override_defaults_neptune_no_verify(self): ssl=True, ssl_verify=False) self.generate_config_from_main_and_test(expected_config, host_type='neptune') + def test_generate_configuration_main_override_defaults_neptune_with_serializer(self): + expected_config = Configuration(self.neptune_host_reg, self.port, neptune_service='neptune-graph', + auth_mode=AuthModeEnum.IAM, load_from_s3_arn='loader_arn', ssl=False, + gremlin_section=GremlinSection(message_serializer='graphbinary')) + self.generate_config_from_main_and_test(expected_config, host_type='neptune') + def test_generate_configuration_main_override_defaults_neptune_cn(self): expected_config = Configuration(self.neptune_host_cn, self.port, neptune_service='neptune-graph', auth_mode=AuthModeEnum.IAM, load_from_s3_arn='loader_arn', ssl=False) @@ -115,7 +121,8 @@ def generate_config_from_main_and_test(self, source_config: Configuration, host_ f'--load_from_s3_arn "{source_config.load_from_s3_arn}" ' f'--proxy_host "{source_config.proxy_host}" ' f'--proxy_port "{source_config.proxy_port}" ' - f'--config_destination="{self.test_file_path}" ') + f'--gremlin_serializer "{source_config.gremlin.message_serializer}" ' + f'--config_destination="{self.test_file_path}"') else: result = os.system(f'{self.python_cmd} -m graph_notebook.configuration.generate_config ' f'--host "{source_config.host}" --port "{source_config.port}" ' From 6cbb36a276ba7302bb76c70a5df60b5c81e30b12 Mon Sep 17 00:00:00 2001 From: Michael Chin Date: Fri, 8 Dec 2023 19:25:44 -0800 Subject: [PATCH 2/2] update changelog --- ChangeLog.md | 1 + 1 file changed, 1 insertion(+) diff --git a/ChangeLog.md b/ChangeLog.md index 3b76eaae..d9ff4ddf 100644 --- a/ChangeLog.md +++ b/ChangeLog.md @@ -4,6 +4,7 @@ Starting with v1.31.6, this file will contain a record of major features and upd ## Upcoming - Fixed `neptune_ml_utils` imports in `03-Neptune-ML` samples ([Link to PR](https://github.com/aws/graph-notebook/pull/546)) +- Enable Gremlin `message_serializer` config field for Neptune endpoints ([Link to PR](https://github.com/aws/graph-notebook/pull/547)) ## Release 4.0.1 (Nov 29, 2023) - Fixed @neptune_db_only magics decorator ([Link to PR](https://github.com/aws/graph-notebook/pull/543))