Fix serialization of union types at root level (fixes #168)

avro/src/main/java/com/fasterxml/jackson/dataformat/avro/ser/ArrayWriteContext.java

@@ -11,7 +11,7 @@ public class ArrayWriteContext
     protected final GenericArray<Object> _array;
 
     public ArrayWriteContext(AvroWriteContext parent, AvroGenerator generator,
-            GenericArray<Object> array)
+            GenericArray<Object> array, Object currValue)
     {
         super(TYPE_ARRAY, parent, generator, array.getSchema());
         _array = array;
@@ -21,23 +21,15 @@ public ArrayWriteContext(AvroWriteContext parent, AvroGenerator generator,
     public Object rawValue() { return _array; }
 
     @Override
-    public final AvroWriteContext createChildArrayContext() throws JsonMappingException {
+    public final AvroWriteContext createChildArrayContext(Object currValue) throws JsonMappingException {
         GenericArray<Object> arr = _createArray(_schema.getElementType());
         _array.add(arr);
-        return new ArrayWriteContext(this, _generator, arr);
+        return new ArrayWriteContext(this, _generator, arr, currValue);
     }
 
     @Override
-    public final AvroWriteContext createChildObjectContext() throws JsonMappingException
-    {
-        AvroWriteContext child = _createObjectContext(_schema.getElementType());
-        _array.add(child.rawValue());
-        return child;
-    }
-
-    @Override
-    public AvroWriteContext createChildObjectContext(Object object) throws JsonMappingException {
-        AvroWriteContext child = _createObjectContext(_schema.getElementType(), object);
+    public AvroWriteContext createChildObjectContext(Object currValue) throws JsonMappingException {
+        AvroWriteContext child = _createObjectContext(_schema.getElementType(), currValue);
         _array.add(child.rawValue());
         return child;
     }

avro/src/main/java/com/fasterxml/jackson/dataformat/avro/ser/AvroWriteContext.java

@@ -54,7 +54,18 @@ protected AvroWriteContext(int type, AvroWriteContext parent,
         _generator = generator;
         _schema = schema;
     }
-
+
+    protected AvroWriteContext(int type, AvroWriteContext parent,
+            AvroGenerator generator, Schema schema, Object currValue)
+    {
+        super();
+        _type = type;
+        _parent = parent;
+        _generator = generator;
+        _schema = schema;
+        _currentValue = currValue;
+    }
+
     // // // Factory methods
 
     public static AvroWriteContext createRootContext(AvroGenerator generator, Schema schema,
@@ -70,15 +81,22 @@ public static AvroWriteContext nullContext() {
         return NullContext.instance;
     }
 
-    public abstract AvroWriteContext createChildArrayContext() throws JsonMappingException;
-    public abstract AvroWriteContext createChildObjectContext() throws JsonMappingException;
+    public final AvroWriteContext createChildArrayContext() throws JsonMappingException {
+        return createChildArrayContext(null);
+    }
+
+    public abstract AvroWriteContext createChildArrayContext(Object currValue) throws JsonMappingException;
+
+    public AvroWriteContext createChildObjectContext() throws JsonMappingException {
+        return createChildObjectContext(null);
+    }
+
+    public abstract AvroWriteContext createChildObjectContext(Object currValue) throws JsonMappingException;
 
     public void complete() throws IOException {
         throw new IllegalStateException("Can not be called on "+getClass().getName());
     }
 
-    public AvroWriteContext createChildObjectContext(Object object) throws JsonMappingException { return createChildObjectContext(); }
-
     /*
     /**********************************************************
     /* Accessors
@@ -154,6 +172,28 @@ public final String toString()
 
     // // // Shared helper methods
 
+    protected GenericRecord _createRecord(Schema schema, Object currValue) throws JsonMappingException
+    {
+        Type type = schema.getType();
+        if (type == Schema.Type.UNION) {
+            try {
+                schema = resolveUnionSchema(schema, currValue);
+            } catch (UnresolvedUnionException e) {
+                // couldn't find an exact match
+                schema = _recordOrMapFromUnion(schema);
+            }
+        }
+        if (type == Schema.Type.MAP) {
+            throw new IllegalStateException("_createRecord should never be called for elements of type MAP");
+        }
+        try {
+            return new GenericData.Record(schema);
+        } catch (RuntimeException e) {
+            // alas, generator not passed to us
+            throw new JsonMappingException(null, "Failed to create Record type from "+type, e);
+        }
+    }
+
     protected GenericRecord _createRecord(Schema schema) throws JsonMappingException
     {
         // Quick check: if type is Union, need to find actual record type...
@@ -191,22 +231,23 @@ protected AvroWriteContext _createObjectContext(Schema schema) throws JsonMappin
         return _createObjectContext(schema, null); // Object doesn't matter as long as schema isn't a union
     }
 
-    protected AvroWriteContext _createObjectContext(Schema schema, Object object) throws JsonMappingException
+    protected AvroWriteContext _createObjectContext(Schema schema, Object currValue)
+            throws JsonMappingException
     {
         Type type = schema.getType();
         if (type == Schema.Type.UNION) {
             try {
-                schema = resolveUnionSchema(schema, object);
+                schema = resolveUnionSchema(schema, currValue);
             } catch (UnresolvedUnionException e) {
                 // couldn't find an exact match
                 schema = _recordOrMapFromUnion(schema);
             }
             type = schema.getType();
         }
         if (type == Schema.Type.MAP) {
-            return new MapWriteContext(this, _generator, schema);
+            return new MapWriteContext(this, _generator, schema, currValue);
         }
-        return new ObjectWriteContext(this, _generator, _createRecord(schema));
+        return new ObjectWriteContext(this, _generator, _createRecord(schema), currValue);
     }
 
     protected Schema _recordOrMapFromUnion(Schema unionSchema)
@@ -503,17 +544,17 @@ private NullContext() {
         public Object rawValue() { return null; }
 
         @Override
-        public final AvroWriteContext createChildArrayContext() {
+        public final AvroWriteContext createChildArrayContext(Object currValue) {
             _reportError();
             return null;
         }
-        
+
         @Override
-        public final AvroWriteContext createChildObjectContext() {
+        public final AvroWriteContext createChildObjectContext(Object currValue) {
             _reportError();
             return null;
         }
-    
+
         @Override
         public void writeValue(Object value) {
             _reportError();

avro/src/main/java/com/fasterxml/jackson/dataformat/avro/ser/KeyValueContext.java

@@ -14,9 +14,9 @@ abstract class KeyValueContext extends AvroWriteContext
     protected boolean _expectValue = false;
 
     protected KeyValueContext(AvroWriteContext parent, AvroGenerator generator,
-            Schema schema)
+            Schema schema, Object currValue)
     {
-        super(TYPE_OBJECT, parent, generator, schema);
+        super(TYPE_OBJECT, parent, generator, schema, currValue);
     }
 
     @Override

avro/src/main/java/com/fasterxml/jackson/dataformat/avro/ser/MapWriteContext.java

@@ -17,9 +17,10 @@ public final class MapWriteContext
 {
     protected final Map<String,Object> _data;
 
-    public MapWriteContext(AvroWriteContext parent, AvroGenerator generator, Schema schema)
+    public MapWriteContext(AvroWriteContext parent, AvroGenerator generator,
+            Schema schema, Object currValue)
     {
-        super(parent, generator, schema);
+        super(parent, generator, schema, currValue);
         _data = new HashMap<String,Object>();
     }
 
@@ -33,21 +34,12 @@ public final boolean writeFieldName(String name)
         _expectValue = true;
         return true;
     }
-
-    @Override
-    public final AvroWriteContext createChildArrayContext() {
-        _verifyValueWrite();
-        AvroWriteContext child = new ArrayWriteContext(this, _generator,
-                _createArray(_schema.getValueType()));
-        _data.put(_currentName, child.rawValue());
-        return child;
-    }
 
     @Override
-    public final AvroWriteContext createChildObjectContext() throws JsonMappingException
-    {
+    public final AvroWriteContext createChildArrayContext(Object currValue) {
         _verifyValueWrite();
-        AvroWriteContext child = _createObjectContext(_schema.getValueType());
+        AvroWriteContext child = new ArrayWriteContext(this, _generator,
+                _createArray(_schema.getValueType()), currValue);
         _data.put(_currentName, child.rawValue());
         return child;
     }

avro/src/main/java/com/fasterxml/jackson/dataformat/avro/ser/NopWriteContext.java

@@ -8,23 +8,24 @@
  */
 public class NopWriteContext extends AvroWriteContext
 {
-    public NopWriteContext(int type, AvroWriteContext parent, AvroGenerator generator) {
-        super(type, parent, generator, null);
+    public NopWriteContext(int type, AvroWriteContext parent, AvroGenerator generator,
+            Object currValue) {
+        super(type, parent, generator, null, currValue);
     }
 
     @Override
     public Object rawValue() { return null; }
 
     @Override
-    public final AvroWriteContext createChildArrayContext() throws JsonMappingException {
-        return new NopWriteContext(TYPE_ARRAY, this, _generator);
+    public final AvroWriteContext createChildArrayContext(Object currValue) throws JsonMappingException {
+        return new NopWriteContext(TYPE_ARRAY, this, _generator, currValue);
     }
-    
+
     @Override
-    public final AvroWriteContext createChildObjectContext() throws JsonMappingException {
-        return new NopWriteContext(TYPE_OBJECT, this, _generator);
+    public final AvroWriteContext createChildObjectContext(Object currValue) throws JsonMappingException {
+        return new NopWriteContext(TYPE_OBJECT, this, _generator, currValue);
     }
-    
+
     @Override
     public void writeValue(Object value) { }
 

avro/src/main/java/com/fasterxml/jackson/dataformat/avro/ser/ObjectWriteContext.java

@@ -23,49 +23,50 @@ public final class ObjectWriteContext
     protected Schema.Field _nextField;
 
     public ObjectWriteContext(AvroWriteContext parent, AvroGenerator generator,
-            GenericRecord record)
+            GenericRecord record, Object currValue)
     {
-        super(parent, generator, record.getSchema());
+        super(parent, generator, record.getSchema(), currValue);
         _record = record;
     }
 
     @Override
     public Object rawValue() { return _record; }
 
+
     @Override
-    public final AvroWriteContext createChildArrayContext()
+    public final AvroWriteContext createChildArrayContext(Object currValue)
     {
         _verifyValueWrite();
         Schema.Field field = _findField();
         if (field == null) { // unknown, to ignore
-            return new NopWriteContext(TYPE_ARRAY, this, _generator);
+            return new NopWriteContext(TYPE_ARRAY, this, _generator, currValue);
         }
-        AvroWriteContext child = new ArrayWriteContext(this, _generator, _createArray(field.schema()));
+        AvroWriteContext child = new ArrayWriteContext(this, _generator,
+                _createArray(field.schema()), currValue);
         _record.put(_currentName, child.rawValue());
         return child;
     }
 
     @Override
-    public final AvroWriteContext createChildObjectContext() throws JsonMappingException
-    {
+    public AvroWriteContext createChildObjectContext() throws JsonMappingException {
         _verifyValueWrite();
         Schema.Field field = _findField();
         if (field == null) { // unknown, to ignore
-            return new NopWriteContext(TYPE_OBJECT, this, _generator);
+            return new NopWriteContext(TYPE_OBJECT, this, _generator, null);
         }
         AvroWriteContext child = _createObjectContext(field.schema());
         _record.put(_currentName, child.rawValue());
         return child;
     }
 
     @Override
-    public AvroWriteContext createChildObjectContext(Object object) throws JsonMappingException {
+    public AvroWriteContext createChildObjectContext(Object currValue) throws JsonMappingException {
         _verifyValueWrite();
         Schema.Field field = _findField();
         if (field == null) { // unknown, to ignore
-            return new NopWriteContext(TYPE_OBJECT, this, _generator);
+            return new NopWriteContext(TYPE_OBJECT, this, _generator, currValue);
         }
-        AvroWriteContext child = _createObjectContext(field.schema(), object);
+        AvroWriteContext child = _createObjectContext(field.schema(), currValue);
         _record.put(_currentName, child.rawValue());
         return child;
     }

avro/src/main/java/com/fasterxml/jackson/dataformat/avro/ser/RootContext.java

@@ -38,10 +38,9 @@ public RootContext(AvroGenerator generator, Schema schema, BinaryEncoder encoder
 
     @Override
     public Object rawValue() { return _rootValue; }
-    
+
     @Override
-    public final AvroWriteContext createChildArrayContext() throws JsonMappingException
-    {
+    public final AvroWriteContext createChildArrayContext(Object currValue) throws JsonMappingException {
         // verify that root type is array (or compatible)
         switch (_schema.getType()) {
         case ARRAY:
@@ -53,24 +52,23 @@ public final AvroWriteContext createChildArrayContext() throws JsonMappingExcept
         }
         GenericArray<Object> arr = _createArray(_schema);
         _rootValue = arr;
-        return new ArrayWriteContext(this, _generator, arr);
+        return new ArrayWriteContext(this, _generator, arr, currValue);
     }
-    
+
     @Override
-    public final AvroWriteContext createChildObjectContext() throws JsonMappingException
-    {
+    public final AvroWriteContext createChildObjectContext(Object currValue) throws JsonMappingException {
         // verify that root type is record (or compatible)
         switch (_schema.getType()) {
         case RECORD:
         case UNION: // maybe
             {
-                GenericRecord rec = _createRecord(_schema);
+                GenericRecord rec = _createRecord(_schema, currValue);
                 _rootValue = rec;
-                return new ObjectWriteContext(this, _generator, rec);
+                return new ObjectWriteContext(this, _generator, rec, currValue);
             }
         case MAP: // used to not be supported
             {
-                MapWriteContext ctxt = new MapWriteContext(this, _generator, _schema);
+                MapWriteContext ctxt = new MapWriteContext(this, _generator, _schema, currValue);
                 _rootValue = ctxt.rawValue();
                 return ctxt;
             }