[BEAM-13965] Use TypeDeserializer if type information is available to…

… support polymorphic types. (#16913)
apache · Feb 25, 2022 · daabeba · daabeba
1 parent c75807c
commit daabeba
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 10 deletions.
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptionsFactory.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/options/PipelineOptionsFactory.java
@@ -36,11 +36,13 @@
 import com.fasterxml.jackson.databind.annotation.JsonSerialize;
 import com.fasterxml.jackson.databind.deser.DefaultDeserializationContext;
 import com.fasterxml.jackson.databind.deser.impl.MethodProperty;
+import com.fasterxml.jackson.databind.deser.impl.TypeWrappedDeserializer;
 import com.fasterxml.jackson.databind.introspect.AnnotatedMember;
 import com.fasterxml.jackson.databind.introspect.AnnotatedMethod;
 import com.fasterxml.jackson.databind.introspect.AnnotationCollector;
 import com.fasterxml.jackson.databind.introspect.BeanPropertyDefinition;
 import com.fasterxml.jackson.databind.introspect.TypeResolutionContext;
+import com.fasterxml.jackson.databind.jsontype.TypeDeserializer;
 import com.fasterxml.jackson.databind.node.TreeTraversingParser;
 import com.fasterxml.jackson.databind.ser.DefaultSerializerProvider;
 import com.fasterxml.jackson.databind.type.TypeBindings;
@@ -1730,21 +1732,23 @@ private static JsonDeserializer<Object> computeDeserializerForMethod(Method meth
       BeanProperty prop = createBeanProperty(method);
       AnnotatedMember annotatedMethod = prop.getMember();
 
+      DefaultDeserializationContext context = DESERIALIZATION_CONTEXT.get();
       Object maybeDeserializerClass =
-          DESERIALIZATION_CONTEXT
-              .get()
-              .getAnnotationIntrospector()
-              .findDeserializer(annotatedMethod);
+          context.getAnnotationIntrospector().findDeserializer(annotatedMethod);
 
       JsonDeserializer<Object> jsonDeserializer =
-          DESERIALIZATION_CONTEXT
-              .get()
-              .deserializerInstance(annotatedMethod, maybeDeserializerClass);
+          context.deserializerInstance(annotatedMethod, maybeDeserializerClass);
 
       if (jsonDeserializer == null) {
-        jsonDeserializer =
-            DESERIALIZATION_CONTEXT.get().findContextualValueDeserializer(prop.getType(), prop);
+        jsonDeserializer = context.findContextualValueDeserializer(prop.getType(), prop);
       }
+
+      TypeDeserializer typeDeserializer =
+          context.getFactory().findTypeDeserializer(context.getConfig(), prop.getType());
+      if (typeDeserializer != null) {
+        jsonDeserializer = new TypeWrappedDeserializer(typeDeserializer, jsonDeserializer);
+      }
+
       return jsonDeserializer;
     } catch (JsonMappingException e) {
       throw new RuntimeException(e);

diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/options/PipelineOptionsFactoryTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/options/PipelineOptionsFactoryTest.java
@@ -39,6 +39,8 @@
 
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonSubTypes;
+import com.fasterxml.jackson.annotation.JsonTypeInfo;
 import com.fasterxml.jackson.core.JsonGenerator;
 import com.fasterxml.jackson.core.JsonParser;
 import com.fasterxml.jackson.core.JsonProcessingException;
@@ -60,7 +62,6 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 import java.util.Set;
 import org.apache.beam.model.jobmanagement.v1.JobApi.PipelineOptionDescriptor;
 import org.apache.beam.model.jobmanagement.v1.JobApi.PipelineOptionType;
@@ -1070,6 +1071,53 @@ public void testComplexTypes() {
     assertEquals("value2", options.getObjectValue().get().value2);
   }
 
+  @JsonTypeInfo(use = JsonTypeInfo.Id.NAME)
+  @JsonSubTypes({
+    @JsonSubTypes.Type(value = PolymorphicTypeOne.class, name = "one"),
+    @JsonSubTypes.Type(value = PolymorphicTypeTwo.class, name = "two")
+  })
+  public abstract static class PolymorphicType {
+    String key;
+
+    @JsonProperty("key")
+    public String getKey() {
+      return key;
+    }
+
+    public void setKey(String key) {
+      this.key = key;
+    }
+  }
+
+  public static class PolymorphicTypeOne extends PolymorphicType {}
+
+  public static class PolymorphicTypeTwo extends PolymorphicType {}
+
+  public interface PolymorphicTypes extends PipelineOptions {
+    PolymorphicType getObject();
+
+    void setObject(PolymorphicType value);
+
+    ValueProvider<PolymorphicType> getObjectValue();
+
+    void setObjectValue(ValueProvider<PolymorphicType> value);
+  }
+
+  @Test
+  public void testPolymorphicType() {
+    String[] args =
+        new String[] {
+          "--object={\"key\":\"value\",\"@type\":\"one\"}",
+          "--objectValue={\"key\":\"value\",\"@type\":\"two\"}"
+        };
+    PolymorphicTypes options = PipelineOptionsFactory.fromArgs(args).as(PolymorphicTypes.class);
+    assertEquals("value", options.getObject().key);
+    assertEquals(PolymorphicTypeOne.class, options.getObject().getClass());
+
+    assertEquals("value", options.getObjectValue().get().key);
+    assertEquals(PolymorphicTypeTwo.class, options.getObjectValue().get().getClass());
+  }
+
   @Test
   public void testMissingArgument() {
     String[] args = new String[] {};