Allow custom Gremlin serializers for Neptune (#547)

* Allow custom Gremlin serializers for Neptune

* update changelog

Author: michaelnchin

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))

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

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:
 

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)

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}" '