[Avro] Add support for @Union and polymorphic types

avro/src/main/java/com/fasterxml/jackson/dataformat/avro/AvroAnnotationIntrospector.java

@@ -1,20 +1,27 @@
 package com.fasterxml.jackson.dataformat.avro;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
 import org.apache.avro.reflect.*;
 
 import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonTypeInfo;
 import com.fasterxml.jackson.core.Version;
 import com.fasterxml.jackson.databind.AnnotationIntrospector;
+import com.fasterxml.jackson.databind.JavaType;
 import com.fasterxml.jackson.databind.PropertyName;
 import com.fasterxml.jackson.databind.cfg.MapperConfig;
 import com.fasterxml.jackson.databind.introspect.Annotated;
 import com.fasterxml.jackson.databind.introspect.AnnotatedClass;
 import com.fasterxml.jackson.databind.introspect.AnnotatedConstructor;
 import com.fasterxml.jackson.databind.introspect.AnnotatedMember;
+import com.fasterxml.jackson.databind.jsontype.NamedType;
+import com.fasterxml.jackson.databind.jsontype.TypeResolverBuilder;
 import com.fasterxml.jackson.databind.ser.std.ToStringSerializer;
+import com.fasterxml.jackson.dataformat.avro.schema.AvroSchemaHelper;
+
 /**
  * Adds support for the following annotations from the Apache Avro implementation:
  * <ul>
@@ -26,6 +33,7 @@
  * <li>{@link Nullable @Nullable} - Alias for <code>JsonProperty(required = false)</code></li>
  * <li>{@link Stringable @Stringable} - Alias for <code>JsonCreator</code> on the constructor and <code>JsonValue</code> on
  * the {@link #toString()} method. </li>
+ * <li>{@link Union @Union} - Alias for <code>JsonSubTypes</code></li>
  * </ul>
  *
  * @since 2.9
@@ -70,7 +78,7 @@ public List<PropertyName> findPropertyAliases(Annotated m) {
     }
 
     protected PropertyName _findName(Annotated a)
-    {
+	{
         AvroName ann = _findAnnotation(a, AvroName.class);
         return (ann == null) ? null : PropertyName.construct(ann.value());
     }
@@ -107,4 +115,41 @@ public Object findSerializer(Annotated a) {
         }
         return null;
     }
+
+    @Override
+    public List<NamedType> findSubtypes(Annotated a) {
+        Union union = _findAnnotation(a, Union.class);
+        if (union == null) {
+            return null;
+        }
+        ArrayList<NamedType> names = new ArrayList<>(union.value().length);
+        for (Class<?> subtype : union.value()) {
+            names.add(new NamedType(subtype, AvroSchemaHelper.getTypeId(subtype)));
+        }
+        return names;
+    }
+
+    @Override
+    public TypeResolverBuilder<?> findTypeResolver(MapperConfig<?> config, AnnotatedClass ac, JavaType baseType) {
+        return _findTypeResolver(config, ac, baseType);
+    }
+
+    @Override
+    public TypeResolverBuilder<?> findPropertyTypeResolver(MapperConfig<?> config, AnnotatedMember am, JavaType baseType) {
+        return _findTypeResolver(config, am, baseType);
+    }
+
+    @Override
+    public TypeResolverBuilder<?> findPropertyContentTypeResolver(MapperConfig<?> config, AnnotatedMember am, JavaType containerType) {
+        return _findTypeResolver(config, am, containerType);
+    }
+
+    protected TypeResolverBuilder<?> _findTypeResolver(MapperConfig<?> config, Annotated ann, JavaType baseType) {
+        TypeResolverBuilder<?> resolver = new AvroTypeResolverBuilder();
+        JsonTypeInfo typeInfo = ann.getAnnotation(JsonTypeInfo.class);
+        if (typeInfo != null && typeInfo.defaultImpl() != JsonTypeInfo.class) {
+            resolver = resolver.defaultImpl(typeInfo.defaultImpl());
+        }
+        return resolver;
+    }
 }

avro/src/main/java/com/fasterxml/jackson/dataformat/avro/AvroGenerator.java

@@ -1,18 +1,26 @@
 package com.fasterxml.jackson.dataformat.avro;
 
-import com.fasterxml.jackson.core.*;
-import com.fasterxml.jackson.core.base.GeneratorBase;
-import com.fasterxml.jackson.core.io.IOContext;
-import com.fasterxml.jackson.dataformat.avro.ser.AvroWriteContext;
-
-import org.apache.avro.io.BinaryEncoder;
-
 import java.io.IOException;
 import java.io.OutputStream;
 import java.math.BigDecimal;
 import java.math.BigInteger;
 import java.nio.ByteBuffer;
 
+import org.apache.avro.io.BinaryEncoder;
+
+import com.fasterxml.jackson.core.Base64Variant;
+import com.fasterxml.jackson.core.FormatFeature;
+import com.fasterxml.jackson.core.FormatSchema;
+import com.fasterxml.jackson.core.JsonGenerationException;
+import com.fasterxml.jackson.core.JsonGenerator;
+import com.fasterxml.jackson.core.ObjectCodec;
+import com.fasterxml.jackson.core.PrettyPrinter;
+import com.fasterxml.jackson.core.SerializableString;
+import com.fasterxml.jackson.core.Version;
+import com.fasterxml.jackson.core.base.GeneratorBase;
+import com.fasterxml.jackson.core.io.IOContext;
+import com.fasterxml.jackson.dataformat.avro.ser.AvroWriteContext;
+
 public class AvroGenerator extends GeneratorBase
 {
     /**
@@ -381,6 +389,17 @@ public final void writeStartObject() throws IOException {
         _complete = false;
     }
 
+    @Override
+    public void writeStartObject(Object forValue) throws IOException {
+        _avroContext = _avroContext.createChildObjectContext(forValue);
+        _complete = false;
+        if(this._writeContext != null && forValue != null) {
+            this._writeContext.setCurrentValue(forValue);
+        }
+
+        this.setCurrentValue(forValue);
+    }
+
     @Override
     public final void writeEndObject() throws IOException
     {

avro/src/main/java/com/fasterxml/jackson/dataformat/avro/AvroMapper.java

@@ -1,16 +1,16 @@
 package com.fasterxml.jackson.dataformat.avro;
 
-import java.io.*;
+import java.io.File;
+import java.io.IOException;
+import java.io.InputStream;
 
 import org.apache.avro.Schema;
 
 import com.fasterxml.jackson.core.Version;
-
 import com.fasterxml.jackson.databind.JavaType;
 import com.fasterxml.jackson.databind.JsonMappingException;
 import com.fasterxml.jackson.databind.Module;
 import com.fasterxml.jackson.databind.ObjectMapper;
-
 import com.fasterxml.jackson.dataformat.avro.schema.AvroSchemaGenerator;
 
 /**

avro/src/main/java/com/fasterxml/jackson/dataformat/avro/AvroModule.java

@@ -37,6 +37,8 @@ public AvroModule()
         addSerializer(File.class, new ToStringSerializer(File.class));
         // 08-Mar-2016, tatu: to fix [dataformat-avro#35], need to prune 'schema' property:
         setSerializerModifier(new AvroSerializerModifier());
+        // Override untyped deserializer to one that checks for type information in the schema before going to default handling
+        addDeserializer(Object.class, new AvroUntypedDeserializer());
     }
 
     @Override

avro/src/main/java/com/fasterxml/jackson/dataformat/avro/AvroParser.java

@@ -1,7 +1,6 @@
 package com.fasterxml.jackson.dataformat.avro;
 
 import java.io.IOException;
-import java.io.InputStream;
 import java.io.Writer;
 import java.math.BigDecimal;
 
@@ -229,6 +228,16 @@ public void setSchema(FormatSchema schema)
 
     protected abstract void _initSchema(AvroSchema schema) throws JsonProcessingException;
 
+    @Override
+    public boolean canReadTypeId() {
+        return true;
+    }
+
+    @Override
+    public Object getTypeId() throws IOException {
+        return _avroContext != null ? _avroContext.getTypeId() : null;
+    }
+
     /*
     /**********************************************************
     /* Location info

avro/src/main/java/com/fasterxml/jackson/dataformat/avro/AvroTypeDeserializer.java

@@ -0,0 +1,74 @@
+package com.fasterxml.jackson.dataformat.avro;
+
+import java.io.IOException;
+
+import com.fasterxml.jackson.annotation.JsonTypeInfo;
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.databind.BeanProperty;
+import com.fasterxml.jackson.databind.DeserializationContext;
+import com.fasterxml.jackson.databind.JavaType;
+import com.fasterxml.jackson.databind.JsonDeserializer;
+import com.fasterxml.jackson.databind.jsontype.TypeDeserializer;
+import com.fasterxml.jackson.databind.jsontype.TypeIdResolver;
+import com.fasterxml.jackson.databind.jsontype.impl.TypeDeserializerBase;
+import com.fasterxml.jackson.dataformat.avro.schema.AvroSchemaHelper;
+
+public class AvroTypeDeserializer extends TypeDeserializerBase {
+
+    protected AvroTypeDeserializer(JavaType baseType, TypeIdResolver idRes, String typePropertyName, boolean typeIdVisible,
+                                   JavaType defaultImpl) {
+        super(baseType, idRes, typePropertyName, typeIdVisible, defaultImpl);
+    }
+
+    protected AvroTypeDeserializer(TypeDeserializerBase src, BeanProperty property) {
+        super(src, property);
+    }
+
+    @Override
+    public TypeDeserializer forProperty(BeanProperty prop) {
+        return new AvroTypeDeserializer(this, prop);
+    }
+
+    @Override
+    public JsonTypeInfo.As getTypeInclusion() {
+        // Don't do any restructuring of the incoming JSON tokens
+        return JsonTypeInfo.As.EXISTING_PROPERTY;
+    }
+
+    @Override
+    public Object deserializeTypedFromObject(JsonParser p, DeserializationContext ctxt) throws IOException {
+        return deserializeTypedFromAny(p, ctxt);
+    }
+
+    @Override
+    public Object deserializeTypedFromArray(JsonParser p, DeserializationContext ctxt) throws IOException {
+        return deserializeTypedFromAny(p, ctxt);
+    }
+
+    @Override
+    public Object deserializeTypedFromScalar(JsonParser p, DeserializationContext ctxt) throws IOException {
+        return deserializeTypedFromAny(p, ctxt);
+    }
+
+    @Override
+    public Object deserializeTypedFromAny(JsonParser p, DeserializationContext ctxt) throws IOException {
+        if (p.getTypeId() == null && getDefaultImpl() == null) {
+            JsonDeserializer<Object> deser = _findDeserializer(ctxt, AvroSchemaHelper.getTypeId(_baseType));
+            if (deser == null) {
+                ctxt.reportInputMismatch(_baseType, "No (native) type id found when one was expected for polymorphic type handling");
+                return null;
+            }
+            return deser.deserialize(p, ctxt);
+        }
+        return _deserializeWithNativeTypeId(p, ctxt, p.getTypeId());
+    }
+
+    @Override
+    protected JavaType _handleUnknownTypeId(DeserializationContext ctxt, String typeId)
+        throws IOException {
+        if (ctxt.hasValueDeserializerFor(_baseType, null)) {
+            return _baseType;
+        }
+        return super._handleUnknownTypeId(ctxt, typeId);
+    }
+}

avro/src/main/java/com/fasterxml/jackson/dataformat/avro/AvroTypeIdResolver.java

@@ -0,0 +1,63 @@
+package com.fasterxml.jackson.dataformat.avro;
+
+import java.io.IOException;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.Map;
+
+import com.fasterxml.jackson.databind.DatabindContext;
+import com.fasterxml.jackson.databind.JavaType;
+import com.fasterxml.jackson.databind.exc.InvalidTypeIdException;
+import com.fasterxml.jackson.databind.jsontype.NamedType;
+import com.fasterxml.jackson.databind.jsontype.impl.ClassNameIdResolver;
+import com.fasterxml.jackson.databind.type.TypeFactory;
+
+/**
+ * {@link com.fasterxml.jackson.databind.jsontype.TypeIdResolver} for Avro type IDs embedded in schemas. Avro generally uses class names,
+ * but we want to also support named subtypes so that developers can easily remap the embedded type IDs to a different runtime class.
+ */
+public class AvroTypeIdResolver extends ClassNameIdResolver {
+
+    private final Map<String, Class<?>> _idTypes = new HashMap<>();
+
+    private final Map<Class<?>, String> _typeIds = new HashMap<>();
+
+    public AvroTypeIdResolver(JavaType baseType, TypeFactory typeFactory, Collection<NamedType> subTypes) {
+        this(baseType, typeFactory);
+        if (subTypes != null) {
+            for (NamedType namedType : subTypes) {
+                registerSubtype(namedType.getType(), namedType.getName());
+            }
+        }
+    }
+
+    public AvroTypeIdResolver(JavaType baseType, TypeFactory typeFactory) {
+        super(baseType, typeFactory);
+    }
+
+    @Override
+    public void registerSubtype(Class<?> type, String name) {
+        _idTypes.put(name, type);
+        _typeIds.put(type, name);
+    }
+
+    @Override
+    protected JavaType _typeFromId(String id, DatabindContext ctxt) throws IOException {
+        // base types don't have subclasses
+        if (_baseType.isPrimitive()) {
+            return _baseType;
+        }
+        // check if there's a specific type we should be using for this ID
+        Class<?> subType = _idTypes.get(id);
+        if (subType != null) {
+            id = _idFrom(null, subType, _typeFactory);
+        }
+        try {
+            return super._typeFromId(id, ctxt);
+        } catch (InvalidTypeIdException | IllegalArgumentException e) {
+            // AvroTypeDeserializer expects null if we can't map the type ID to a class; It will throw an appropriate error if we can't
+            // find a usable type.
+            return null;
+        }
+    }
+}

avro/src/main/java/com/fasterxml/jackson/dataformat/avro/AvroTypeResolverBuilder.java

@@ -0,0 +1,50 @@
+package com.fasterxml.jackson.dataformat.avro;
+
+import java.util.Collection;
+
+import com.fasterxml.jackson.databind.DeserializationConfig;
+import com.fasterxml.jackson.databind.JavaType;
+import com.fasterxml.jackson.databind.SerializationConfig;
+import com.fasterxml.jackson.databind.cfg.MapperConfig;
+import com.fasterxml.jackson.databind.jsontype.NamedType;
+import com.fasterxml.jackson.databind.jsontype.TypeDeserializer;
+import com.fasterxml.jackson.databind.jsontype.TypeIdResolver;
+import com.fasterxml.jackson.databind.jsontype.TypeSerializer;
+import com.fasterxml.jackson.databind.jsontype.impl.StdTypeResolverBuilder;
+
+
+public class AvroTypeResolverBuilder extends StdTypeResolverBuilder {
+
+    public AvroTypeResolverBuilder() {
+        super();
+        typeIdVisibility(false).typeProperty("@class");
+    }
+
+    @Override
+    public TypeSerializer buildTypeSerializer(SerializationConfig config, JavaType baseType, Collection<NamedType> subtypes) {
+        // All type information is encoded in the schema, never in the data.
+        return null;
+    }
+
+    @Override
+    public TypeDeserializer buildTypeDeserializer(DeserializationConfig config, JavaType baseType, Collection<NamedType> subtypes) {
+        JavaType defaultImpl = null;
+        if (getDefaultImpl() != null) {
+            defaultImpl = config.constructType(getDefaultImpl());
+        }
+
+        return new AvroTypeDeserializer(baseType,
+                                        idResolver(config, baseType, subtypes, true, true),
+                                        getTypeProperty(),
+                                        isTypeIdVisible(),
+                                        defaultImpl
+        );
+
+    }
+
+    @Override
+    protected TypeIdResolver idResolver(MapperConfig<?> config, JavaType baseType, Collection<NamedType> subtypes, boolean forSer,
+                                        boolean forDeser) {
+        return new AvroTypeIdResolver(baseType, config.getTypeFactory(), subtypes);
+    }
+}