Feature/use reader schema for specific records

connectors/activemq/build.sbt

@@ -2,7 +2,7 @@ organization := "com.github.dcshock"
 
 name := "forklift-activemq"
 
-version := "2.0"
+version := "2.1"
 
 javacOptions ++= Seq("-source", "1.8")
 
@@ -22,7 +22,7 @@ resolvers ++= Seq(
 )
 
 libraryDependencies ++= Seq(
-  "com.github.dcshock" % "forklift" % "2.0",
+  "com.github.dcshock" % "forklift" % "2.2",
   "org.apache.activemq" % "activemq-client" % "5.14.0",
   "org.apache.activemq" % "activemq-broker" % "5.14.0",
   "com.fasterxml.jackson.core" % "jackson-databind" % "2.7.3",

connectors/activemq/src/main/java/forklift/connectors/ActiveMQConnector.java

@@ -86,8 +86,7 @@ public synchronized Session getSession()
         }
     }
 
-    @Override
-    public ForkliftConsumerI getQueue(String name)
+    private ForkliftConsumerI getQueue(String name)
       throws ConnectorException {
         final Session s = getSession();
         try {
@@ -97,8 +96,7 @@ public ForkliftConsumerI getQueue(String name)
         }
     }
 
-    @Override
-    public ForkliftConsumerI getTopic(String name)
+    private ForkliftConsumerI getTopic(String name)
       throws ConnectorException {
         final Session s = getSession();
         try {
@@ -111,7 +109,7 @@ public ForkliftConsumerI getTopic(String name)
     @Override
     public ForkliftConsumerI getConsumerForSource(SourceI source) throws ConnectorException {
         return source
-            .apply(QueueSource.class, queue -> getQueue(queue.getName()))
+            .apply(QueueSource.class, queue ->   getQueue(queue.getName()))
             .apply(TopicSource.class, topic -> getTopic(topic.getName()))
             .apply(GroupedTopicSource.class, topic -> getGroupedTopic(topic))
             .apply(RoleInputSource.class, roleSource -> {

connectors/activemq/src/test/java/forklift/activemq/test/ActiveMQForkliftConnector.java

@@ -34,16 +34,6 @@ public Connection getConnection() throws ConnectorException {
         return TestServiceManager.getConnector().getConnection();
     }
 
-    @Override
-    public ForkliftConsumerI getQueue(String name) throws ConnectorException {
-        return TestServiceManager.getConnector().getQueue(name);
-    }
-
-    @Override
-    public ForkliftConsumerI getTopic(String name) throws ConnectorException {
-        return TestServiceManager.getConnector().getTopic(name);
-    }
-
     @Override
     public ForkliftConsumerI getConsumerForSource(SourceI source) throws ConnectorException {
         return TestServiceManager.getConnector().getConsumerForSource(source);

connectors/kafka/build.sbt

@@ -2,7 +2,7 @@ organization := "com.github.dcshock"
 
 name := "forklift-kafka"
 
-version := "2.0"
+version := "2.1"
 
 //required for some test dependencies
 scalaVersion := "2.11.7"
@@ -26,7 +26,7 @@ resolvers ++= Seq(
 )
 
 libraryDependencies ++= Seq(
-  "com.github.dcshock" % "forklift" % "2.0" ,
+  "com.github.dcshock" % "forklift" % "2.2" ,
   "com.fasterxml.jackson.core" % "jackson-databind" % "2.7.3",
   "com.fasterxml.jackson.datatype" % "jackson-datatype-jsr310" % "2.7.3",
   "org.apache.kafka" % "kafka-clients" % "0.10.1.1-cp1" exclude("org.slf4j","slf4j-log4j12"),

connectors/kafka/src/main/java/forklift/connectors/KafkaConnector.java

@@ -4,31 +4,40 @@
 import forklift.consumer.KafkaTopicConsumer;
 import forklift.consumer.wrapper.RoleInputConsumerWrapper;
 import forklift.controller.KafkaController;
+import forklift.decorators.Message;
+import forklift.message.ForkliftAvroMessageUtils;
 import forklift.message.MessageStream;
 import forklift.producers.ForkliftProducerI;
 import forklift.producers.KafkaForkliftProducer;
+import forklift.serializers.ForkliftKafkaAvroDeserializer;
 import forklift.source.ActionSource;
 import forklift.source.LogicalSource;
 import forklift.source.SourceI;
 import forklift.source.sources.GroupedTopicSource;
-import forklift.source.sources.RoleInputSource;
 import forklift.source.sources.QueueSource;
+import forklift.source.sources.RoleInputSource;
 import forklift.source.sources.TopicSource;
-
 import io.confluent.kafka.serializers.KafkaAvroDeserializer;
 import io.confluent.kafka.serializers.KafkaAvroDeserializerConfig;
 import io.confluent.kafka.serializers.KafkaAvroSerializer;
 import io.confluent.kafka.serializers.KafkaAvroSerializerConfig;
+import org.apache.avro.Schema;
+import org.apache.avro.specific.SpecificRecord;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.clients.producer.KafkaProducer;
 import org.apache.kafka.clients.producer.ProducerConfig;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.io.IOException;
+import java.lang.reflect.Field;
+import java.lang.reflect.InvocationTargetException;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Properties;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -101,20 +110,42 @@ private KafkaProducer createKafkaProducer() {
         return new KafkaProducer(producerProperties);
     }
 
-    private KafkaController createController(String topicName) {
+    private KafkaController createController(GroupedTopicSource source) {
+
+        ForkliftKafkaAvroDeserializer deserializer = new ForkliftKafkaAvroDeserializer();
+        Schema readerSchema = null;
+        try {
+            readerSchema = readerSchemaFromSource(source);
+        } catch (Exception e) {
+            log.error("Unable to generate reader schema, falling back to writer schema.  Defaults may be lost");
+        }
+
         Properties props = new Properties();
         props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, kafkaHosts);
         props.put(ConsumerConfig.GROUP_ID_CONFIG, groupId);
         props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false);
         props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, io.confluent.kafka.serializers.KafkaAvroDeserializer.class);
-        props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, io.confluent.kafka.serializers.KafkaAvroDeserializer.class);
+        props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ForkliftKafkaAvroDeserializer.class);
+        if (readerSchema != null) {
+            props.put(ForkliftKafkaAvroDeserializer.READER_SCHEMA_CONFIG, readerSchema);
+        }
         props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
         props.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 200);
         props.put(KafkaAvroDeserializerConfig.SCHEMA_REGISTRY_URL_CONFIG, schemaRegistries);
         props.put(KafkaAvroDeserializerConfig.SPECIFIC_AVRO_READER_CONFIG, false);
 
         final KafkaConsumer<?, ?> kafkaConsumer = new KafkaConsumer(props);
-        return new KafkaController(kafkaConsumer, new MessageStream(), topicName);
+        return new KafkaController(kafkaConsumer, new MessageStream(), source.getName());
+    }
+
+    private Schema readerSchemaFromSource(SourceI source) throws NoSuchMethodException, InvocationTargetException, IllegalAccessException, IOException {
+        for (Field field : source.getContextClass().getDeclaredFields()) {
+            if (field.isAnnotationPresent(Message.class) && SpecificRecord.class.isAssignableFrom(field.getType())) {
+                Schema schema = (Schema) field.getType().getMethod("getClassSchema").invoke(null);
+                return ForkliftAvroMessageUtils.addForkliftPropertiesToSchema(schema);
+            }
+        }
+        return null;
     }
 
     @Override
@@ -138,21 +169,30 @@ public synchronized void stop() throws ConnectorException {
     @Override
     public ForkliftConsumerI getConsumerForSource(SourceI source) throws ConnectorException {
         return source
-            .apply(QueueSource.class, queue -> getQueue(queue.getName()))
-            .apply(TopicSource.class, topic -> getTopic(topic.getName()))
-            .apply(GroupedTopicSource.class, topic -> getGroupedTopic(topic))
-            .apply(RoleInputSource.class, roleSource -> {
-                final ForkliftConsumerI rawConsumer = getConsumerForSource(roleSource.getActionSource(this));
-                return new RoleInputConsumerWrapper(rawConsumer);
-             })
-            .elseUnsupportedError();
+                   .apply(QueueSource.class, queue -> {
+                       GroupedTopicSource groupedTopicSource = new GroupedTopicSource(queue.getName(), groupId);
+                       groupedTopicSource.setContextClass(source.getContextClass());
+                       ForkliftConsumerI consumerI = getGroupedTopic(groupedTopicSource);
+                       return consumerI;
+                   })
+                   .apply(TopicSource.class, topic -> {
+                       GroupedTopicSource groupedTopicSource = new GroupedTopicSource(topic.getName(), groupId);
+                       groupedTopicSource.setContextClass(source.getContextClass());
+                       ForkliftConsumerI consumerI = getGroupedTopic(groupedTopicSource);
+                       return consumerI;
+                   })
+                   .apply(GroupedTopicSource.class, topic -> getGroupedTopic(topic))
+                   .apply(RoleInputSource.class, roleSource -> {
+                       final ForkliftConsumerI rawConsumer = getConsumerForSource(roleSource.getActionSource(this));
+                       return new RoleInputConsumerWrapper(rawConsumer);
+                   })
+                   .elseUnsupportedError();
     }
 
     public synchronized ForkliftConsumerI getGroupedTopic(GroupedTopicSource source) throws ConnectorException {
         if (!source.groupSpecified()) {
             source.overrideGroup(groupId);
         }
-
         if (!source.getGroup().equals(groupId)) { //TODO actually support GroupedTopics
             throw new ConnectorException("Unexpected group '" + source.getGroup() + "'; only the connector group '" + groupId + "' is allowed");
         }
@@ -161,24 +201,13 @@ public synchronized ForkliftConsumerI getGroupedTopic(GroupedTopicSource source)
         if (controller != null && controller.isRunning()) {
             log.warn("Consumer for topic already exists under this controller's groupname.  Messages will be divided amongst consumers.");
         } else {
-            controller = createController(source.getName());
+            controller = createController(source);
             this.controllers.put(source.getName(), controller);
             controller.start();
         }
         return new KafkaTopicConsumer(source.getName(), controller);
     }
-
-    @Override
-    public ForkliftConsumerI getQueue(String name) throws ConnectorException {
-        return getGroupedTopic(new GroupedTopicSource(name, groupId));
-    }
-
-    @Override
-    public ForkliftConsumerI getTopic(String name) throws ConnectorException {
-        return getGroupedTopic(new GroupedTopicSource(name, groupId));
-    }
-
-
+
     @Override
     public ForkliftProducerI getQueueProducer(String name) {
         return getTopicProducer(name);
@@ -195,8 +224,8 @@ public synchronized ForkliftProducerI getTopicProducer(String name) {
     @Override
     public ActionSource mapSource(LogicalSource source) {
         return source
-            .apply(RoleInputSource.class, roleSource -> mapRoleInputSource(roleSource))
-            .get();
+                   .apply(RoleInputSource.class, roleSource -> mapRoleInputSource(roleSource))
+                   .get();
     }
 
     protected GroupedTopicSource mapRoleInputSource(RoleInputSource roleSource) {

connectors/kafka/src/main/java/forklift/message/ForkliftAvroMessageUtils.java

@@ -0,0 +1,37 @@
+package forklift.message;
+
+import com.fasterxml.jackson.databind.DeserializationFeature;
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ArrayNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
+import forklift.connectors.KafkaSerializer;
+import org.apache.avro.Schema;
+
+import java.io.IOException;
+
+public class ForkliftAvroMessageUtils {
+
+    private static final ObjectMapper mapper = new ObjectMapper().registerModule(new JavaTimeModule())
+                                                                 .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
+
+    /**
+     * Adds the forkliftProperties field to the passed in Schema.
+     *
+     * @param schema the schema which forklift properties should be added
+     * @return Schema with forlift properties added
+     * @throws IOException
+     */
+    public static Schema addForkliftPropertiesToSchema(Schema schema) throws IOException {
+        String originalJson = schema.toString(false);
+        JsonNode propertiesField = mapper.readTree(KafkaSerializer.SCHEMA_FIELD_VALUE_PROPERTIES);
+        ObjectNode schemaNode = (ObjectNode) mapper.readTree(originalJson);
+        ArrayNode fieldsNode = (ArrayNode) schemaNode.get("fields");
+        fieldsNode.add(propertiesField);
+        schemaNode.set("fields", fieldsNode);
+        Schema.Parser parser = new Schema.Parser();
+        return parser.parse(mapper.writeValueAsString(schemaNode));
+    }
+
+}