/**
 * Copyright 2015 Confluent Inc.
 *
 * Licensed 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 io.confluent.kafkarest;

import com.fasterxml.jackson.databind.JsonNode;

import org.apache.avro.Schema;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.ProducerRecord;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;

import io.confluent.kafka.serializers.KafkaAvroSerializer;
import io.confluent.kafkarest.converters.AvroConverter;
import io.confluent.kafkarest.converters.ConversionException;
import io.confluent.kafkarest.entities.ProduceRecord;
import io.confluent.kafkarest.entities.SchemaHolder;
import io.confluent.rest.exceptions.RestException;

public class AvroRestProducer implements RestProducer<JsonNode, JsonNode> {

  protected final KafkaProducer<Object, Object> producer;
  protected final KafkaAvroSerializer keySerializer;
  protected final KafkaAvroSerializer valueSerializer;

  public AvroRestProducer(KafkaProducer<Object, Object> producer,
                          KafkaAvroSerializer keySerializer,
                          KafkaAvroSerializer valueSerializer) {
    this.producer = producer;
    this.keySerializer = keySerializer;
    this.valueSerializer = valueSerializer;
  }

  public void produce(ProduceTask task, String topic, Integer partition,
                      Collection<? extends ProduceRecord<JsonNode, JsonNode>> records) {
    SchemaHolder schemaHolder = task.getSchemaHolder();
    Schema keySchema = null, valueSchema = null;
    Integer keySchemaId = schemaHolder.getKeySchemaId();
    Integer valueSchemaId = schemaHolder.getValueSchemaId();
    try {
      // If both ID and schema are null, that may be ok. Validation of the ProduceTask by the
      // caller should have checked this already.
      if (keySchemaId != null) {
        keySchema = keySerializer.getByID(keySchemaId);
      } else if (schemaHolder.getKeySchema() != null) {
        keySchema = new Schema.Parser().parse(schemaHolder.getKeySchema());
        keySchemaId = keySerializer.register(topic + "-key", keySchema);
      }

      if (valueSchemaId != null) {
        valueSchema = valueSerializer.getByID(valueSchemaId);
      } else if (schemaHolder.getValueSchema() != null) {
        valueSchema = new Schema.Parser().parse(schemaHolder.getValueSchema());
        valueSchemaId = valueSerializer.register(topic + "-value", valueSchema);
      }
    } catch (IOException e) {
      // FIXME We should return more specific error codes (unavailable vs registration failed in
      // a way that isn't retriable?).
      throw new RestException("Schema registration or lookup failed", 408, 40801, e);
    }

    // Store the schema IDs in the task. These will be used to include the IDs in the response
    task.setSchemaIds(keySchemaId, valueSchemaId);

    // Convert everything to Avro before doing any sends so if any conversion fails we can kill
    // the entire request so we don't get partially sent requests
    ArrayList<ProducerRecord<Object, Object>> kafkaRecords
        = new ArrayList<ProducerRecord<Object, Object>>();
    try {
      for (ProduceRecord<JsonNode, JsonNode> record : records) {
        // Beware of null schemas and NullNodes here: we need to avoid attempting the conversion
        // if there isn't a schema. Validation will have already checked that all the keys/values
        // were NullNodes.
        Object key = (keySchema != null ? AvroConverter.toAvro(record.getKey(), keySchema) : null);
        Object value = (valueSchema != null
                        ? AvroConverter.toAvro(record.getValue(), valueSchema) : null);
        kafkaRecords.add(new ProducerRecord(topic, partition, key, value));
      }
    } catch (ConversionException e) {
      throw Errors.jsonAvroConversionException();
    }
    for (ProducerRecord<Object, Object> rec : kafkaRecords) {
      producer.send(rec, task);
    }
  }

  public void close() {
    producer.close();
  }
}

@@ -1237,16 +1237,101 @@
 ception;
+%0Aimport io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException;
 %0A%0Apublic
@@ -3058,18 +3058,26 @@
  catch (
-IO
+RestClient
 Exceptio
@@ -3306,24 +3306,140 @@
 1, e);%0A    %7D
+ catch (IOException e) %7B%0A      throw new RestException(%22Schema registration or lookup failed%22, 408, 40801, e);%0A    %7D
 %0A%0A    // Sto

Handle RestClientException checked exceptions that are now thrown by the schema-registry-client.