Implement AvroItemReader and AvroItemWriter

Resolves BATCH-2833

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/avro/AvroItemReader.java

@@ -0,0 +1,177 @@
+/*
+ * Copyright 2019 the original author or authors.
+ *
+ *  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
+ *
+ *          https://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.springframework.batch.item.avro;
+
+import java.io.IOException;
+import java.io.InputStream;
+
+import org.apache.avro.Schema;
+import org.apache.avro.file.DataFileStream;
+import org.apache.avro.generic.GenericDatumReader;
+import org.apache.avro.generic.GenericRecord;
+import org.apache.avro.io.BinaryDecoder;
+import org.apache.avro.io.DatumReader;
+import org.apache.avro.io.DecoderFactory;
+import org.apache.avro.reflect.ReflectDatumReader;
+import org.apache.avro.specific.SpecificDatumReader;
+import org.apache.avro.specific.SpecificRecordBase;
+
+import org.springframework.batch.item.ItemReader;
+import org.springframework.batch.item.ItemStreamException;
+import org.springframework.batch.item.support.AbstractItemCountingItemStreamItemReader;
+import org.springframework.core.io.Resource;
+import org.springframework.util.Assert;
+
+/**
+ * An {@link ItemReader} that deserializes data from a {@link Resource} containing serialized Avro objects.
+ *
+ * @author David Turanski
+ * @since 4.2
+ */
+public class AvroItemReader<T> extends AbstractItemCountingItemStreamItemReader<T> {
+
+	private boolean embeddedSchema = true;
+
+	private InputStreamReader<T> inputStreamReader;
+
+	private DataFileStream<T> dataFileReader;
+
+	private InputStream inputStream;
+
+	private DatumReader<T> datumReader;
+
+	/**
+	 *
+	 * @param resource the {@link Resource} containing objects serialized with Avro.
+	 * @param clazz the data type to be deserialized.
+	 */
+	public AvroItemReader(Resource resource, Class<T> clazz) {
+		Assert.notNull(resource, "'resource' is required.");
+		Assert.notNull(clazz, "'class' is required.");
+
+		try {
+			this.inputStream = resource.getInputStream();
+			this.datumReader = datumReaderForClass(clazz);
+		}
+		catch (IOException e) {
+			throw new IllegalArgumentException(e.getMessage(), e);
+		}
+	}
+
+	/**
+	 *
+	 * @param data the {@link Resource} containing the data to be read.
+	 * @param schema the {@link Resource} containing the Avro schema.
+	 */
+	public AvroItemReader(Resource data, Resource schema) {
+		Assert.notNull(data, "'data' is required.");
+		Assert.state(data.exists(), "'data' " + data.getFilename() +" does not exist.");
+		Assert.notNull(schema, "'schema' is required");
+		Assert.state(schema.exists(), "'schema' " + schema.getFilename() +" does not exist.");
+		try {
+			this.inputStream = data.getInputStream();
+			Schema avroSchema = new Schema.Parser().parse(schema.getInputStream());
+			this.datumReader = new GenericDatumReader<>(avroSchema);
+		}
+		catch (IOException e) {
+			throw new IllegalArgumentException(e.getMessage(), e);
+		}
+	}
+
+	/**
+	 * Disable or enable reading an embedded Avro schema. True by default.
+	 * @param embeddedSchema set to false to if the input does not embed an Avro schema.
+	 */
+	public void setEmbeddedSchema(boolean embeddedSchema) {
+		this.embeddedSchema = embeddedSchema;
+	}
+
+
+	@Override
+	protected T doRead() throws Exception {
+	    if (inputStreamReader != null) {
+            return inputStreamReader.read();
+        }
+	    return dataFileReader.hasNext()? dataFileReader.next(): null;
+	}
+
+	@Override
+	protected void doOpen() throws Exception {
+		initializeReader();
+	}
+
+	@Override
+	protected void doClose() throws Exception {
+		if (this.inputStreamReader != null) {
+			inputStreamReader.close();
+			return;
+		}
+		dataFileReader.close();
+	}
+
+	private void  initializeReader() throws IOException {
+		if (this.embeddedSchema) {
+			this.dataFileReader = new DataFileStream<>(this.inputStream, this.datumReader);
+		} else {
+			this.inputStreamReader = createInputStreamReader(this.inputStream, this.datumReader);
+		}
+
+	}
+
+	private InputStreamReader<T> createInputStreamReader(InputStream inputStream, DatumReader<T> datumReader) {
+		return new InputStreamReader<>(inputStream, datumReader);
+	}
+
+	private static <T> DatumReader<T> datumReaderForClass(Class<T> clazz) {
+		if (SpecificRecordBase.class.isAssignableFrom(clazz)){
+			return new SpecificDatumReader<>(clazz);
+		}
+		if (GenericRecord.class.isAssignableFrom(clazz)) {
+			return new GenericDatumReader<>();
+		}
+		return new ReflectDatumReader<>(clazz);
+	}
+
+
+	private static class InputStreamReader<T> {
+        private final DatumReader<T> datumReader;
+
+        private final BinaryDecoder binaryDecoder;
+
+        private final InputStream inputStream;
+
+        private InputStreamReader(InputStream inputStream, DatumReader<T> datumReader) {
+            this.inputStream = inputStream;
+            this.datumReader = datumReader;
+            this.binaryDecoder = DecoderFactory.get().binaryDecoder(inputStream, null);
+        }
+
+        private T read() throws Exception {
+            if (!binaryDecoder.isEnd()) {
+                return datumReader.read(null, binaryDecoder);
+            }
+            return null;
+        }
+
+        private void close() {
+			try {
+				this.inputStream.close();
+			} catch (IOException e) {
+				throw new ItemStreamException(e.getMessage(), e);
+			}
+		}
+    }
+}

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/avro/AvroItemWriter.java

@@ -0,0 +1,200 @@
+/*
+ * Copyright 2019 the original author or authors.
+ *
+ * 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
+ *
+ *      https://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.springframework.batch.item.avro;
+
+import java.io.IOException;
+import java.io.OutputStream;
+import java.util.List;
+
+import org.apache.avro.Schema;
+import org.apache.avro.file.DataFileWriter;
+import org.apache.avro.generic.GenericDatumWriter;
+import org.apache.avro.generic.GenericRecord;
+import org.apache.avro.io.BinaryEncoder;
+import org.apache.avro.io.DatumWriter;
+import org.apache.avro.io.EncoderFactory;
+import org.apache.avro.reflect.ReflectDatumWriter;
+import org.apache.avro.specific.SpecificDatumWriter;
+import org.apache.avro.specific.SpecificRecordBase;
+
+import org.springframework.batch.item.ExecutionContext;
+import org.springframework.batch.item.ItemStreamException;
+import org.springframework.batch.item.ItemWriter;
+import org.springframework.batch.item.support.AbstractItemStreamItemWriter;
+import org.springframework.core.io.Resource;
+import org.springframework.core.io.WritableResource;
+import org.springframework.util.Assert;
+
+/**
+ * An {@link ItemWriter} that serializes data to an {@link WritableResource} using Avro.
+ *
+ * This does not support restart on failure.
+ *
+ * @since 4.2
+ * @author David Turanski
+ */
+public class AvroItemWriter<T> extends AbstractItemStreamItemWriter<T> {
+
+	private DataFileWriter<T> dataFileWriter;
+
+	private OutputStreamWriter<T> outputStreamWriter;
+
+	private WritableResource resource;
+
+	private Schema schema;
+
+	private Resource schemaResource;
+
+	private Class<T> clazz;
+
+	private boolean embedSchema = true;
+
+
+	/**
+	 *
+	 * @param resource a {@link WritableResource} to which the objects will be serialized.
+	 * @param schema a {@link Resource} containing the Avro schema.
+	 * @param clazz the data type to be serialized.
+	 */
+	public AvroItemWriter(WritableResource resource, Resource schema, Class<T> clazz) {
+		this.schemaResource = schema;
+		this.resource = resource;
+		this.clazz = clazz;
+	}
+
+	/**
+	 * This constructor will create an ItemWriter that does not embedded Avro schema.
+	 *
+	 * @param resource a {@link WritableResource} to which the objects will be serialized.
+	 * @param clazz the data type to be serialized.
+	 */
+	public AvroItemWriter(WritableResource resource, Class<T> clazz) {
+		this(resource, null, clazz);
+		embedSchema = false;
+	}
+
+	@Override
+	public void write(List<? extends T> items) throws Exception {
+		items.forEach(item -> {
+			try {
+				if (this.dataFileWriter != null) {
+					this.dataFileWriter.append(item);
+				}
+				else {
+					this.outputStreamWriter.write(item);
+				}
+			}
+			catch (Exception e) {
+				throw new ItemStreamException(e.getMessage(), e);
+			}
+		});
+	}
+
+	/**
+	 * @see org.springframework.batch.item.ItemStream#open(ExecutionContext)
+	 */
+	@Override
+	public void open(ExecutionContext executionContext) {
+		try {
+			initializeWriter();
+		} catch (IOException e) {
+			throw new ItemStreamException(e.getMessage(), e);
+		}
+	}
+
+	@Override
+	public void close() {
+		try {
+			if (this.dataFileWriter != null) {
+				dataFileWriter.close();
+			}
+			else {
+				this.outputStreamWriter.close();
+			}
+		}
+		catch (IOException e) {
+			throw new ItemStreamException(e.getMessage(), e);
+		}
+	}
+
+	private void  initializeWriter() throws IOException {
+		Assert.notNull(resource, "'resource' is required.");
+		Assert.notNull(clazz, "'class' is required.");
+
+		if (this.embedSchema) {
+			Assert.notNull(this.schemaResource, "'schema' is required.");
+			Assert.state(this.schemaResource.exists(),
+					"'schema' " + this.schemaResource.getFilename() + " does not exist.");
+			try {
+				this.schema = new Schema.Parser().parse(this.schemaResource.getInputStream());
+			} catch (IOException e) {
+				throw new IllegalArgumentException(e.getMessage(), e);
+			}
+			this.dataFileWriter = new DataFileWriter<>(datumWriterForClass(this.clazz));
+			this.dataFileWriter.create(this.schema, this.resource.getOutputStream());
+		} else {
+			this.outputStreamWriter = createOutputStreamWriter(resource.getOutputStream(),
+					datumWriterForClass(this.clazz));
+		}
+
+	}
+
+	private static <T> DatumWriter<T> datumWriterForClass(Class<T> clazz) {
+		if (GenericRecord.class.isAssignableFrom(clazz)) {
+			return new GenericDatumWriter<>();
+
+		}
+		if (SpecificRecordBase.class.isAssignableFrom(clazz)){
+			return new SpecificDatumWriter<>(clazz);
+		}
+
+		return new ReflectDatumWriter<>(clazz);
+	}
+
+	private AvroItemWriter.OutputStreamWriter<T> createOutputStreamWriter(OutputStream outputStream,
+			DatumWriter<T> datumWriter) {
+		return new AvroItemWriter.OutputStreamWriter<>(outputStream, datumWriter);
+	}
+
+	private static class OutputStreamWriter<T> {
+		private final DatumWriter<T> datumWriter;
+
+		private final BinaryEncoder binaryEncoder;
+
+		private final OutputStream outputStream;
+
+		private OutputStreamWriter(OutputStream outputStream, DatumWriter<T> datumWriter) {
+			this.outputStream = outputStream;
+			this.datumWriter = datumWriter;
+			this.binaryEncoder = EncoderFactory.get().binaryEncoder(outputStream, null);
+		}
+
+		private void write(T datum) throws Exception {
+			datumWriter.write(datum, binaryEncoder);
+			binaryEncoder.flush();
+		}
+
+		private void close() {
+			try {
+				this.outputStream.close();
+			}
+			catch (IOException e) {
+				throw new ItemStreamException(e.getMessage(), e);
+			}
+		}
+	}
+}

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/avro/builder/AvroItemReaderBuilder.java

@@ -0,0 +1,203 @@
+/*
+ * Copyright 2019 the original author or authors.
+ *
+ * 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
+ *
+ *      https://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.springframework.batch.item.avro.builder;
+
+import org.apache.avro.Schema;
+
+import org.springframework.batch.item.avro.AvroItemReader;
+import org.springframework.core.io.ByteArrayResource;
+import org.springframework.core.io.Resource;
+import org.springframework.util.Assert;
+import org.springframework.util.StringUtils;
+
+/**
+ * A builder implementation for the {@link AvroItemReader}
+ *
+ * @author David Turanski
+ * @since 4.2
+ */
+public class AvroItemReaderBuilder<T> {
+
+	private boolean saveState = true;
+
+	private String name = AvroItemReader.class.getSimpleName();
+
+	private int maxItemCount = Integer.MAX_VALUE;
+
+	private int currentItemCount;
+
+	private Resource schema;
+
+	private Resource resource;
+
+	private Class<T> type;
+
+	private boolean embeddedSchema =true;
+
+
+	/**
+	 * Configure a {@link Resource} containing Avro serialized objects.
+	 * @param resource an existing Resource.
+	 * @return The current instance of the builder.
+	 */
+	public AvroItemReaderBuilder<T> resource(Resource resource) {
+		Assert.notNull(resource, "A 'resource' is required.");
+		Assert.state(resource.exists(), "Resource " + resource.getFilename() + " does not exist.");
+		this.resource = resource;
+		return this;
+	}
+
+
+	/**
+	 * Configure an Avro {@link Schema} from a {@link Resource}.
+	 * @param schema an existing schema Resource.
+	 * @return The current instance of the builder.
+	 */
+	public AvroItemReaderBuilder<T> schema(Resource schema) {
+		Assert.notNull(schema, "A 'schema' Resource is required.");
+		Assert.state(schema.exists(), "Resource " + schema.getFilename() + " does not exist.");
+		this.schema = schema;
+		return this;
+	}
+
+	/**
+	 * Configure an Avro {@link Schema} from a String.
+	 * @param schemaString the schema String.
+	 * @return The current instance of the builder.
+	 */
+	public AvroItemReaderBuilder<T> schema(String schemaString) {
+		Assert.hasText(schemaString, "A 'schema' is required.");
+		this.schema = new ByteArrayResource(schemaString.getBytes());
+		return this;
+	}
+
+	/**
+	 * Configure a type to be deserialized.
+	 * @param type the class to be deserialized.
+	 * @return The current instance of the builder.
+	 */
+	public AvroItemReaderBuilder<T> type(Class<T> type) {
+		Assert.notNull(type, "A 'type' is required.");
+		this.type = type;
+		return this;
+	}
+
+	/**
+	 * Disable or enable reading an embedded Avro schema. True by default.
+	 * @param embeddedSchema set to false to if the input does not contain an Avro schema.
+	 */
+	public AvroItemReaderBuilder<T> embeddedSchema(boolean embeddedSchema) {
+		this.embeddedSchema = embeddedSchema;
+		return this;
+	}
+
+	/**
+	 * Configure if the state of the {@link org.springframework.batch.item.ItemStreamSupport}
+	 * should be persisted within the {@link org.springframework.batch.item.ExecutionContext}
+	 * for restart purposes.
+	 *
+	 * @param saveState defaults to true
+	 * @return The current instance of the builder.
+	 */
+	public AvroItemReaderBuilder<T> saveState(boolean saveState) {
+		this.saveState = saveState;
+		return this;
+	}
+
+	/**
+	 * The name used to calculate the key within the
+	 * {@link org.springframework.batch.item.ExecutionContext}. Required if
+	 * {@link #saveState(boolean)} is set to true.
+	 *
+	 * @param name name of the reader instance
+	 * @return The current instance of the builder.
+	 * @see org.springframework.batch.item.ItemStreamSupport#setName(String)
+	 */
+	public AvroItemReaderBuilder<T> name(String name) {
+		this.name = name;
+		return this;
+	}
+
+	/**
+	 * Configure the max number of items to be read.
+	 *
+	 * @param maxItemCount the max items to be read
+	 * @return The current instance of the builder.
+	 * @see org.springframework.batch.item.support.AbstractItemCountingItemStreamItemReader#setMaxItemCount(int)
+	 */
+	public AvroItemReaderBuilder<T> maxItemCount(int maxItemCount) {
+		this.maxItemCount = maxItemCount;
+		return this;
+	}
+
+	/**
+	 * Index for the current item. Used on restarts to indicate where to start from.
+	 *
+	 * @param currentItemCount current index
+	 * @return this instance for method chaining
+	 * @see org.springframework.batch.item.support.AbstractItemCountingItemStreamItemReader#setCurrentItemCount(int)
+	 */
+	public AvroItemReaderBuilder<T> currentItemCount(int currentItemCount) {
+		this.currentItemCount = currentItemCount;
+		return this;
+	}
+
+
+	/**
+	 * Build an instance of {@link AvroItemReader}.
+	 * @return the instance;
+	 */
+	public AvroItemReader<T> build() {
+		AvroItemReader<T> avroItemReader;
+
+		Assert.notNull(this.resource, "A 'resource' is required.");
+
+		if (this.type != null) {
+			avroItemReader = buildForType();
+		}
+		else {
+			avroItemReader = buildForSchema();
+		}
+
+		avroItemReader.setSaveState(this.saveState);
+
+		if(this.saveState) {
+			Assert.state(StringUtils.hasText(this.name),
+					"A name is required when saveState is set to true.");
+		}
+
+		avroItemReader.setName(this.name);
+		avroItemReader.setCurrentItemCount(this.currentItemCount);
+		avroItemReader.setMaxItemCount(this.maxItemCount);
+		avroItemReader.setEmbeddedSchema(this.embeddedSchema);
+
+		return avroItemReader;
+	}
+
+
+	private AvroItemReader<T> buildForType() {
+		Assert.isNull(this.schema, "You cannot specify a schema and 'type'.");
+		return new AvroItemReader<>(this.resource, this.type);
+	}
+
+	private AvroItemReader<T> buildForSchema() {
+		Assert.notNull(this.schema, "'schema' is required.");
+		return new AvroItemReader<>(resource, schema);
+	}
+
+
+}

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/avro/builder/AvroItemWriterBuilder.java

@@ -0,0 +1,118 @@
+/*
+ * Copyright 2019 the original author or authors.
+ *
+ * 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
+ *
+ *      https://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.springframework.batch.item.avro.builder;
+
+import org.springframework.batch.item.avro.AvroItemWriter;
+import org.springframework.core.io.ByteArrayResource;
+import org.springframework.core.io.Resource;
+import org.springframework.core.io.WritableResource;
+import org.springframework.util.Assert;
+
+/**
+ * @author David Turanski
+ * @since 4.2
+ */
+public class AvroItemWriterBuilder<T> {
+	private Class<T> type;
+
+	private WritableResource resource;
+
+	private Resource schema;
+
+	private String name  = AvroItemWriter.class.getSimpleName();
+
+	/**
+	 *
+	 * @param resource the {@link WritableResource} used to write the serialized data.
+	 * @return The current instance of the builder.
+	 */
+	public AvroItemWriterBuilder<T> resource(WritableResource resource) {
+		Assert.notNull(resource, "A 'resource' is required.");
+		this.resource = resource;
+		return this;
+	}
+
+	/**
+	 *
+	 * @param schema the Resource containing the schema JSON used to serialize the output.
+	 * @return The current instance of the builder.
+	 */
+	public AvroItemWriterBuilder<T> schema(Resource schema) {
+		Assert.notNull(schema, "A 'schema' is required.");
+		Assert.state(schema.exists(), "Resource " + schema.getFilename() + "does not exist.");
+		this.schema = schema;
+		return this;
+	}
+
+
+	/**
+	 *
+	 * @param schemaString the String containing the schema JSON used to serialize the output.
+	 * @return The current instance of the builder.
+	 */
+	public AvroItemWriterBuilder<T> schema(String schemaString) {
+		Assert.hasText(schemaString, "A 'schemaString' is required.");
+		this.schema = new ByteArrayResource(schemaString.getBytes());
+		return this;
+	}
+
+
+	/**
+	 *
+	 * @param type the Class of objects to be serialized.
+	 * @return The current instance of the builder.
+	 */
+	public AvroItemWriterBuilder<T> type(Class<T> type) {
+		Assert.notNull(type, "A 'type' is required.");
+		this.type = type;
+		return this;
+	}
+
+	/**
+	 * The name used to calculate the key within the
+	 * {@link org.springframework.batch.item.ExecutionContext}.
+	 *
+	 * @param name name of the reader instance
+	 * @return The current instance of the builder.
+	 * @see org.springframework.batch.item.ItemStreamSupport#setName(String)
+	 */
+	public AvroItemWriterBuilder<T> name(String name) {
+		Assert.hasText(name, "A 'name' is required.");
+		this.name = name;
+		return this;
+	}
+
+	/**
+	 * Build an instance of {@link AvroItemWriter}.
+	 *
+	 * @return the instance;
+	 */
+	public AvroItemWriter<T> build() {
+
+		Assert.notNull(this.resource, "A 'resource' is required.");
+
+		Assert.notNull(this.type, "A 'type' is required.");
+
+
+		AvroItemWriter<T> avroItemWriter = this.schema != null ?
+				new AvroItemWriter(this.resource, this.schema, this.type):
+				new AvroItemWriter<>(this.resource, this.type);
+		avroItemWriter.setName(name);
+		return avroItemWriter;
+	}
+
+}