[tinkerpop] 01/01: TINKERPOP-2150 Support configuring the registry builder in config

[jo...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

jorgebg pushed a commit to branch TINKERPOP-2150
in repository https://gitbox.apache.org/repos/asf/tinkerpop.git

commit 2c91de023f12642866fa0e3d6f6e25c540bab21e
Author: Jorge Bay Gondra <jo...@gmail.com>
AuthorDate: Tue Feb 5 13:07:27 2019 +0100

    TINKERPOP-2150 Support configuring the registry builder in config
---
 .../driver/ser/GraphBinaryMessageSerializerV1.java | 20 ++++++++++++++--
 .../binary/GraphBinaryMessageSerializerV1Test.java | 27 ++++++++++++++++++++++
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/gremlin-driver/src/main/java/org/apache/tinkerpop/gremlin/driver/ser/GraphBinaryMessageSerializerV1.java b/gremlin-driver/src/main/java/org/apache/tinkerpop/gremlin/driver/ser/GraphBinaryMessageSerializerV1.java
index cd5405b..c0defd4 100644
--- a/gremlin-driver/src/main/java/org/apache/tinkerpop/gremlin/driver/ser/GraphBinaryMessageSerializerV1.java
+++ b/gremlin-driver/src/main/java/org/apache/tinkerpop/gremlin/driver/ser/GraphBinaryMessageSerializerV1.java
@@ -34,6 +34,7 @@ import org.apache.tinkerpop.gremlin.structure.Graph;
 import org.apache.tinkerpop.gremlin.structure.io.IoRegistry;
 import org.javatuples.Pair;
 
+import java.lang.reflect.Constructor;
 import java.lang.reflect.Method;
 import java.util.List;
 import java.util.Map;
@@ -42,6 +43,7 @@ import static java.nio.charset.StandardCharsets.UTF_8;
 
 public class GraphBinaryMessageSerializerV1 extends AbstractMessageSerializer {
     public static final String TOKEN_CUSTOM = "custom";
+    public static final String TOKEN_BUILDER = "builder";
     private static final String MIME_TYPE = SerTokens.MIME_GRAPHBINARY_V1D0;
     private static final byte[] HEADER = MIME_TYPE.getBytes(UTF_8);
 
@@ -71,9 +73,23 @@ public class GraphBinaryMessageSerializerV1 extends AbstractMessageSerializer {
 
     @Override
     public void configure(final Map<String, Object> config, final Map<String, Graph> graphs) {
-        final List<String> classNameList = getListStringFromConfig(TOKEN_IO_REGISTRIES, config);
-        final TypeSerializerRegistry.Builder builder = TypeSerializerRegistry.build();
+        final String builderClassName = (String) config.get(TOKEN_BUILDER);
+        final TypeSerializerRegistry.Builder builder;
+
+        if (builderClassName != null) {
+            try {
+                Class<?> clazz = Class.forName(builderClassName);
+                Constructor<?> ctor = clazz.getConstructor();
+                builder = (TypeSerializerRegistry.Builder) ctor.newInstance();
+            } catch (Exception ex) {
+                throw new IllegalStateException(ex);
+            }
 
+        } else {
+            builder = TypeSerializerRegistry.build();
+        }
+
+        final List<String> classNameList = getListStringFromConfig(TOKEN_IO_REGISTRIES, config);
         classNameList.forEach(className -> {
             try {
                 final Class<?> clazz = Class.forName(className);
diff --git a/gremlin-driver/src/test/java/org/apache/tinkerpop/gremlin/driver/ser/binary/GraphBinaryMessageSerializerV1Test.java b/gremlin-driver/src/test/java/org/apache/tinkerpop/gremlin/driver/ser/binary/GraphBinaryMessageSerializerV1Test.java
index d3683de..0e8e7ed 100644
--- a/gremlin-driver/src/test/java/org/apache/tinkerpop/gremlin/driver/ser/binary/GraphBinaryMessageSerializerV1Test.java
+++ b/gremlin-driver/src/test/java/org/apache/tinkerpop/gremlin/driver/ser/binary/GraphBinaryMessageSerializerV1Test.java
@@ -28,7 +28,10 @@ import org.apache.tinkerpop.gremlin.driver.ser.SerializationException;
 import org.junit.Test;
 import org.mockito.internal.matchers.apachecommons.ReflectionEquals;
 
+import java.util.HashMap;
+import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThat;
@@ -131,6 +134,20 @@ public class GraphBinaryMessageSerializerV1Test {
         assertResponseEquals(response, deserialized);
     }
 
+    @Test
+    public void shouldSupportConfigurationOfRegistryBuilder() {
+        final GraphBinaryMessageSerializerV1 serializer = new GraphBinaryMessageSerializerV1();
+        final Map<String, Object> config = new HashMap<>();
+        int counter = SampleTypeSerializerRegistryBuilder.createCounter.get();
+
+        config.put(GraphBinaryMessageSerializerV1.TOKEN_BUILDER, "org.apache.tinkerpop.gremlin.driver.ser.binary.GraphBinaryMessageSerializerV1Test$SampleTypeSerializerRegistryBuilder");
+        serializer.configure(config, null);
+
+        counter = SampleTypeSerializerRegistryBuilder.createCounter.get() - counter;
+        // There should be a call to `create()`
+        assertEquals(1, counter);
+    }
+
     private static void assertResponseEquals(ResponseMessage expected, ResponseMessage actual) {
         assertEquals(expected.getRequestId(), actual.getRequestId());
         // Status
@@ -141,4 +158,14 @@ public class GraphBinaryMessageSerializerV1Test {
         assertEquals(expected.getResult().getData(), actual.getResult().getData());
         assertEquals(expected.getResult().getMeta(), actual.getResult().getMeta());
     }
+
+    public static class SampleTypeSerializerRegistryBuilder extends TypeSerializerRegistry.Builder {
+        public static AtomicInteger createCounter = new AtomicInteger();
+
+        @Override
+        public TypeSerializerRegistry create() {
+            createCounter.incrementAndGet();
+            return super.create();
+        }
+    }
 }