lang/java/protobuf/src/main/java/org/apache/avro/protobuf/ProtobufData.java - avro - Git at Google

 /**
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
 package org.apache.avro.protobuf;

 import java.util.List;
 import java.util.Arrays;
 import java.util.ArrayList;
 import java.util.Map;
 import java.util.IdentityHashMap;
 import java.util.concurrent.ConcurrentHashMap;

 import java.io.IOException;
 import java.io.File;

 import org.apache.avro.Schema;
 import org.apache.avro.Schema.Field;
 import org.apache.avro.generic.GenericData;
 import org.apache.avro.specific.SpecificData;
 import org.apache.avro.io.DatumReader;
 import org.apache.avro.io.DatumWriter;

 import com.google.protobuf.ByteString;
 import com.google.protobuf.Message;
 import com.google.protobuf.Message.Builder;
 import com.google.protobuf.MessageOrBuilder;
 import com.google.protobuf.Descriptors.Descriptor;
 import com.google.protobuf.Descriptors.FieldDescriptor;
 import com.google.protobuf.Descriptors.EnumDescriptor;
 import com.google.protobuf.Descriptors.EnumValueDescriptor;
 import com.google.protobuf.Descriptors.FileDescriptor;
 import com.google.protobuf.DescriptorProtos.FileOptions;

 import org.apache.avro.util.ClassUtils;
 import org.codehaus.jackson.JsonFactory;
 import org.codehaus.jackson.JsonNode;
 import org.codehaus.jackson.map.ObjectMapper;
 import org.codehaus.jackson.node.JsonNodeFactory;

 /** Utilities for serializing Protobuf data in Avro format. */
 public class ProtobufData extends GenericData {
   private static final String PROTOBUF_TYPE = "protobuf";

   private static final ProtobufData INSTANCE = new ProtobufData();

   protected ProtobufData() {}

   /** Return the singleton instance. */
   public static ProtobufData get() { return INSTANCE; }

   @Override
   public DatumReader createDatumReader(Schema schema) {
     return new ProtobufDatumReader(schema, schema, this);
   }

   @Override
   public DatumWriter createDatumWriter(Schema schema) {
     return new ProtobufDatumWriter(schema, this);
   }

   @Override
   public void setField(Object r, String n, int pos, Object o) {
     setField(r, n, pos, o, getRecordState(r, getSchema(r.getClass())));
   }

   @Override
   public Object getField(Object r, String name, int pos) {
     return getField(r, name, pos, getRecordState(r, getSchema(r.getClass())));
   }

   @Override
   protected void setField(Object r, String n, int pos, Object o, Object state) {
     Builder b = (Builder)r;
     FieldDescriptor f = ((FieldDescriptor[])state)[pos];
     switch (f.getType()) {
     case MESSAGE:
       if (o == null) {
         b.clearField(f);
         break;
       }
     default:
       b.setField(f, o);
     }
   }

   @Override
   protected Object getField(Object record, String name, int pos, Object state) {
     Message m = (Message)record;
     FieldDescriptor f = ((FieldDescriptor[])state)[pos];
     switch (f.getType()) {
     case MESSAGE:
       if (!f.isRepeated() && !m.hasField(f))
         return null;
     default:
       return m.getField(f);
     }
   }

   private final Map<Descriptor,FieldDescriptor[]> fieldCache =
     new ConcurrentHashMap<Descriptor,FieldDescriptor[]>();

   @Override
   protected Object getRecordState(Object r, Schema s) {
     Descriptor d = ((MessageOrBuilder)r).getDescriptorForType();
     FieldDescriptor[] fields = fieldCache.get(d);
     if (fields == null) {                         // cache miss
       fields = new FieldDescriptor[s.getFields().size()];
       for (Field f : s.getFields())
         fields[f.pos()] = d.findFieldByName(f.name());
       fieldCache.put(d, fields);                  // update cache
     }
     return fields;
   }

   @Override
   protected boolean isRecord(Object datum) {
     return datum instanceof Message;
   }

   @Override
   public Object newRecord(Object old, Schema schema) {
     try {
       Class c = ClassUtils.forName(SpecificData.getClassName(schema));
       if (c == null)
         return newRecord(old, schema);            // punt to generic
       if (c.isInstance(old))
         return old;                               // reuse instance
       return c.getMethod("newBuilder").invoke(null);

     } catch (Exception e) {
       throw new RuntimeException(e);
     }
   }

   @Override
   protected boolean isArray(Object datum) {
     return datum instanceof List;
   }

   @Override
   protected boolean isBytes(Object datum) {
     return datum instanceof ByteString;
   }

   @Override
   protected Schema getRecordSchema(Object record) {
     return getSchema(((Message)record).getDescriptorForType());
   }

   private final Map<Class,Schema> schemaCache
     = new ConcurrentHashMap<Class,Schema>();

   /** Return a record schema given a protobuf message class. */
   public Schema getSchema(Class c) {
     Schema schema = schemaCache.get(c);

     if (schema == null) {                         // cache miss
       try {
         Object descriptor = c.getMethod("getDescriptor").invoke(null);
         if (c.isEnum())
           schema = getSchema((EnumDescriptor)descriptor);
         else
           schema = getSchema((Descriptor)descriptor);
       } catch (Exception e) {
         throw new RuntimeException(e);
       }
       schemaCache.put(c, schema);                 // update cache
     }
     return schema;
   }

   private static final ThreadLocal<Map<Descriptor,Schema>> SEEN
     = new ThreadLocal<Map<Descriptor,Schema>>() {
     protected Map<Descriptor,Schema> initialValue() {
       return new IdentityHashMap<Descriptor,Schema>();
     }
   };

   private Schema getSchema(Descriptor descriptor) {
     Map<Descriptor,Schema> seen = SEEN.get();
     if (seen.containsKey(descriptor))             // stop recursion
       return seen.get(descriptor);
     boolean first = seen.isEmpty();
     try {
       Schema result =
         Schema.createRecord(descriptor.getName(), null,
                             getNamespace(descriptor.getFile(),
                                          descriptor.getContainingType()),
                             false);

       seen.put(descriptor, result);

       List<Field> fields = new ArrayList<Field>();
       for (FieldDescriptor f : descriptor.getFields())
         fields.add(new Field(f.getName(), getSchema(f), null, getDefault(f)));
       result.setFields(fields);
       return result;

     } finally {
       if (first)
         seen.clear();
     }
   }

   private String getNamespace(FileDescriptor fd, Descriptor containing) {
     FileOptions o = fd.getOptions();
     String p = o.hasJavaPackage()
       ? o.getJavaPackage()
       : fd.getPackage();
     String outer;
     if (o.hasJavaOuterClassname()) {
       outer = o.getJavaOuterClassname();
     } else {
       outer = new File(fd.getName()).getName();
       outer = outer.substring(0, outer.lastIndexOf('.'));
       outer = toCamelCase(outer);
     }
     String inner = "";
     while (containing != null) {
       inner = containing.getName() + "$" + inner;
       containing = containing.getContainingType();
     }
     return p + "." + outer + "$" + inner;
   }

   private static String toCamelCase(String s){
     String[] parts = s.split("_");
     String camelCaseString = "";
     for (String part : parts) {
       camelCaseString = camelCaseString + cap(part);
     }
     return camelCaseString;
   }

   private static String cap(String s) {
     return s.substring(0, 1).toUpperCase() + s.substring(1).toLowerCase();
   }

   private static final Schema NULL = Schema.create(Schema.Type.NULL);

   private Schema getSchema(FieldDescriptor f) {
     Schema s = getNonRepeatedSchema(f);
     if (f.isRepeated())
       s = Schema.createArray(s);
     return s;
   }

   private Schema getNonRepeatedSchema(FieldDescriptor f) {
     Schema result;
     switch (f.getType()) {
     case BOOL:
       return Schema.create(Schema.Type.BOOLEAN);
     case FLOAT:
       return Schema.create(Schema.Type.FLOAT);
     case DOUBLE:
       return Schema.create(Schema.Type.DOUBLE);
     case STRING:
       Schema s = Schema.create(Schema.Type.STRING);
       GenericData.setStringType(s, GenericData.StringType.String);
       return s;
     case BYTES:
       return Schema.create(Schema.Type.BYTES);
     case INT32: case UINT32: case SINT32: case FIXED32: case SFIXED32:
       return Schema.create(Schema.Type.INT);
     case INT64: case UINT64: case SINT64: case FIXED64: case SFIXED64:
       return Schema.create(Schema.Type.LONG);
     case ENUM:
       return getSchema(f.getEnumType());
     case MESSAGE:
       result = getSchema(f.getMessageType());
       if (f.isOptional())
         // wrap optional record fields in a union with null
         result = Schema.createUnion(Arrays.asList(new Schema[] {NULL, result}));
       return result;
     case GROUP:                                   // groups are deprecated
     default:
       throw new RuntimeException("Unexpected type: "+f.getType());
     }
   }

   private Schema getSchema(EnumDescriptor d) {
     List<String> symbols = new ArrayList<String>();
     for (EnumValueDescriptor e : d.getValues()) {
       symbols.add(e.getName());
     }
     return Schema.createEnum(d.getName(), null,
                              getNamespace(d.getFile(), d.getContainingType()),
                              symbols);
   }

   private static final JsonFactory FACTORY = new JsonFactory();
   private static final ObjectMapper MAPPER = new ObjectMapper(FACTORY);
   private static final JsonNodeFactory NODES = JsonNodeFactory.instance;

   private JsonNode getDefault(FieldDescriptor f) {
     if (f.isRequired() || f.isRepeated())         // no default
       return null;

     if (f.hasDefaultValue()) {                    // parse spec'd default value
       Object value = f.getDefaultValue();
       switch (f.getType()) {
       case ENUM:
         value = ((EnumValueDescriptor)value).getName();
         break;
       }
       String json = toString(value);
       try {
         return MAPPER.readTree(FACTORY.createJsonParser(json));
       } catch (IOException e) {
         throw new RuntimeException(e);
       }
     }

     switch (f.getType()) {                        // generate default for type
     case BOOL:
       return NODES.booleanNode(false);
     case FLOAT: case DOUBLE:
     case INT32: case UINT32: case SINT32: case FIXED32: case SFIXED32:
     case INT64: case UINT64: case SINT64: case FIXED64: case SFIXED64:
       return NODES.numberNode(0);
     case STRING: case BYTES:
       return NODES.textNode("");
     case ENUM:
       return NODES.textNode(f.getEnumType().getValues().get(0).getName());
     case MESSAGE:
       return NODES.nullNode();
     case GROUP:                                   // groups are deprecated
     default:
       throw new RuntimeException("Unexpected type: "+f.getType());
     }

   }

 }
	/**
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/
	package org.apache.avro.protobuf;

	import java.util.List;
	import java.util.Arrays;
	import java.util.ArrayList;
	import java.util.Map;
	import java.util.IdentityHashMap;
	import java.util.concurrent.ConcurrentHashMap;

	import java.io.IOException;
	import java.io.File;

	import org.apache.avro.Schema;
	import org.apache.avro.Schema.Field;
	import org.apache.avro.generic.GenericData;
	import org.apache.avro.specific.SpecificData;
	import org.apache.avro.io.DatumReader;
	import org.apache.avro.io.DatumWriter;

	import com.google.protobuf.ByteString;
	import com.google.protobuf.Message;
	import com.google.protobuf.Message.Builder;
	import com.google.protobuf.MessageOrBuilder;
	import com.google.protobuf.Descriptors.Descriptor;
	import com.google.protobuf.Descriptors.FieldDescriptor;
	import com.google.protobuf.Descriptors.EnumDescriptor;
	import com.google.protobuf.Descriptors.EnumValueDescriptor;
	import com.google.protobuf.Descriptors.FileDescriptor;
	import com.google.protobuf.DescriptorProtos.FileOptions;

	import org.apache.avro.util.ClassUtils;
	import org.codehaus.jackson.JsonFactory;
	import org.codehaus.jackson.JsonNode;
	import org.codehaus.jackson.map.ObjectMapper;
	import org.codehaus.jackson.node.JsonNodeFactory;

	/** Utilities for serializing Protobuf data in Avro format. */
	public class ProtobufData extends GenericData {
	private static final String PROTOBUF_TYPE = "protobuf";

	private static final ProtobufData INSTANCE = new ProtobufData();

	protected ProtobufData() {}

	/** Return the singleton instance. */
	public static ProtobufData get() { return INSTANCE; }

	@Override
	public DatumReader createDatumReader(Schema schema) {
	return new ProtobufDatumReader(schema, schema, this);
	}

	@Override
	public DatumWriter createDatumWriter(Schema schema) {
	return new ProtobufDatumWriter(schema, this);
	}

	@Override
	public void setField(Object r, String n, int pos, Object o) {
	setField(r, n, pos, o, getRecordState(r, getSchema(r.getClass())));
	}

	@Override
	public Object getField(Object r, String name, int pos) {
	return getField(r, name, pos, getRecordState(r, getSchema(r.getClass())));
	}

	@Override
	protected void setField(Object r, String n, int pos, Object o, Object state) {
	Builder b = (Builder)r;
	FieldDescriptor f = ((FieldDescriptor[])state)[pos];
	switch (f.getType()) {
	case MESSAGE:
	if (o == null) {
	b.clearField(f);
	break;
	}
	default:
	b.setField(f, o);
	}
	}

	@Override
	protected Object getField(Object record, String name, int pos, Object state) {
	Message m = (Message)record;
	FieldDescriptor f = ((FieldDescriptor[])state)[pos];
	switch (f.getType()) {
	case MESSAGE:
	if (!f.isRepeated() && !m.hasField(f))
	return null;
	default:
	return m.getField(f);
	}
	}

	private final Map<Descriptor,FieldDescriptor[]> fieldCache =
	new ConcurrentHashMap<Descriptor,FieldDescriptor[]>();

	@Override
	protected Object getRecordState(Object r, Schema s) {
	Descriptor d = ((MessageOrBuilder)r).getDescriptorForType();
	FieldDescriptor[] fields = fieldCache.get(d);
	if (fields == null) { // cache miss
	fields = new FieldDescriptor[s.getFields().size()];
	for (Field f : s.getFields())
	fields[f.pos()] = d.findFieldByName(f.name());
	fieldCache.put(d, fields); // update cache
	}
	return fields;
	}

	@Override
	protected boolean isRecord(Object datum) {
	return datum instanceof Message;
	}

	@Override
	public Object newRecord(Object old, Schema schema) {
	try {
	Class c = ClassUtils.forName(SpecificData.getClassName(schema));
	if (c == null)
	return newRecord(old, schema); // punt to generic
	if (c.isInstance(old))
	return old; // reuse instance
	return c.getMethod("newBuilder").invoke(null);

	} catch (Exception e) {
	throw new RuntimeException(e);
	}
	}

	@Override
	protected boolean isArray(Object datum) {
	return datum instanceof List;
	}

	@Override
	protected boolean isBytes(Object datum) {
	return datum instanceof ByteString;
	}

	@Override
	protected Schema getRecordSchema(Object record) {
	return getSchema(((Message)record).getDescriptorForType());
	}

	private final Map<Class,Schema> schemaCache
	= new ConcurrentHashMap<Class,Schema>();

	/** Return a record schema given a protobuf message class. */
	public Schema getSchema(Class c) {
	Schema schema = schemaCache.get(c);

	if (schema == null) { // cache miss
	try {
	Object descriptor = c.getMethod("getDescriptor").invoke(null);
	if (c.isEnum())
	schema = getSchema((EnumDescriptor)descriptor);
	else
	schema = getSchema((Descriptor)descriptor);
	} catch (Exception e) {
	throw new RuntimeException(e);
	}
	schemaCache.put(c, schema); // update cache
	}
	return schema;
	}

	private static final ThreadLocal<Map<Descriptor,Schema>> SEEN
	= new ThreadLocal<Map<Descriptor,Schema>>() {
	protected Map<Descriptor,Schema> initialValue() {
	return new IdentityHashMap<Descriptor,Schema>();
	}
	};

	private Schema getSchema(Descriptor descriptor) {
	Map<Descriptor,Schema> seen = SEEN.get();
	if (seen.containsKey(descriptor)) // stop recursion
	return seen.get(descriptor);
	boolean first = seen.isEmpty();
	try {
	Schema result =
	Schema.createRecord(descriptor.getName(), null,
	getNamespace(descriptor.getFile(),
	descriptor.getContainingType()),
	false);

	seen.put(descriptor, result);

	List<Field> fields = new ArrayList<Field>();
	for (FieldDescriptor f : descriptor.getFields())
	fields.add(new Field(f.getName(), getSchema(f), null, getDefault(f)));
	result.setFields(fields);
	return result;

	} finally {
	if (first)
	seen.clear();
	}
	}

	private String getNamespace(FileDescriptor fd, Descriptor containing) {
	FileOptions o = fd.getOptions();
	String p = o.hasJavaPackage()
	? o.getJavaPackage()
	: fd.getPackage();
	String outer;
	if (o.hasJavaOuterClassname()) {
	outer = o.getJavaOuterClassname();
	} else {
	outer = new File(fd.getName()).getName();
	outer = outer.substring(0, outer.lastIndexOf('.'));
	outer = toCamelCase(outer);
	}
	String inner = "";
	while (containing != null) {
	inner = containing.getName() + "$" + inner;
	containing = containing.getContainingType();
	}
	return p + "." + outer + "$" + inner;
	}

	private static String toCamelCase(String s){
	String[] parts = s.split("_");
	String camelCaseString = "";
	for (String part : parts) {
	camelCaseString = camelCaseString + cap(part);
	}
	return camelCaseString;
	}

	private static String cap(String s) {
	return s.substring(0, 1).toUpperCase() + s.substring(1).toLowerCase();
	}

	private static final Schema NULL = Schema.create(Schema.Type.NULL);

	private Schema getSchema(FieldDescriptor f) {
	Schema s = getNonRepeatedSchema(f);
	if (f.isRepeated())
	s = Schema.createArray(s);
	return s;
	}

	private Schema getNonRepeatedSchema(FieldDescriptor f) {
	Schema result;
	switch (f.getType()) {
	case BOOL:
	return Schema.create(Schema.Type.BOOLEAN);
	case FLOAT:
	return Schema.create(Schema.Type.FLOAT);
	case DOUBLE:
	return Schema.create(Schema.Type.DOUBLE);
	case STRING:
	Schema s = Schema.create(Schema.Type.STRING);
	GenericData.setStringType(s, GenericData.StringType.String);
	return s;
	case BYTES:
	return Schema.create(Schema.Type.BYTES);
	case INT32: case UINT32: case SINT32: case FIXED32: case SFIXED32:
	return Schema.create(Schema.Type.INT);
	case INT64: case UINT64: case SINT64: case FIXED64: case SFIXED64:
	return Schema.create(Schema.Type.LONG);
	case ENUM:
	return getSchema(f.getEnumType());
	case MESSAGE:
	result = getSchema(f.getMessageType());
	if (f.isOptional())
	// wrap optional record fields in a union with null
	result = Schema.createUnion(Arrays.asList(new Schema[] {NULL, result}));
	return result;
	case GROUP: // groups are deprecated
	default:
	throw new RuntimeException("Unexpected type: "+f.getType());
	}
	}

	private Schema getSchema(EnumDescriptor d) {
	List<String> symbols = new ArrayList<String>();
	for (EnumValueDescriptor e : d.getValues()) {
	symbols.add(e.getName());
	}
	return Schema.createEnum(d.getName(), null,
	getNamespace(d.getFile(), d.getContainingType()),
	symbols);
	}

	private static final JsonFactory FACTORY = new JsonFactory();
	private static final ObjectMapper MAPPER = new ObjectMapper(FACTORY);
	private static final JsonNodeFactory NODES = JsonNodeFactory.instance;

	private JsonNode getDefault(FieldDescriptor f) {
	if (f.isRequired() \|\| f.isRepeated()) // no default
	return null;

	if (f.hasDefaultValue()) { // parse spec'd default value
	Object value = f.getDefaultValue();
	switch (f.getType()) {
	case ENUM:
	value = ((EnumValueDescriptor)value).getName();
	break;
	}
	String json = toString(value);
	try {
	return MAPPER.readTree(FACTORY.createJsonParser(json));
	} catch (IOException e) {
	throw new RuntimeException(e);
	}
	}

	switch (f.getType()) { // generate default for type
	case BOOL:
	return NODES.booleanNode(false);
	case FLOAT: case DOUBLE:
	case INT32: case UINT32: case SINT32: case FIXED32: case SFIXED32:
	case INT64: case UINT64: case SINT64: case FIXED64: case SFIXED64:
	return NODES.numberNode(0);
	case STRING: case BYTES:
	return NODES.textNode("");
	case ENUM:
	return NODES.textNode(f.getEnumType().getValues().get(0).getName());
	case MESSAGE:
	return NODES.nullNode();
	case GROUP: // groups are deprecated
	default:
	throw new RuntimeException("Unexpected type: "+f.getType());
	}

	}

	}