Implement Spring Batch BigQuery

spring-batch-bigquery/pom.xml

@@ -0,0 +1,114 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+  ~ Copyright 2002-2021 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.
+  -->
+
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
+
+    <modelVersion>4.0.0</modelVersion>
+
+    <groupId>org.springframework.batch</groupId>
+    <artifactId>spring-batch-bigquery</artifactId>
+    <version>0.1.0-SNAPSHOT</version>
+
+    <properties>
+        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
+
+        <!-- Dependent on Spring Batch core -->
+        <java.version>8</java.version>
+    </properties>
+
+    <dependencies>
+        <dependency>
+            <groupId>org.springframework.batch</groupId>
+            <artifactId>spring-batch-core</artifactId>
+            <version>4.3.2</version>
+        </dependency>
+
+        <dependency>
+            <groupId>com.google.cloud</groupId>
+            <artifactId>google-cloud-bigquery</artifactId>
+            <version>1.128.0</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.avro</groupId>
+            <artifactId>avro</artifactId>
+            <version>1.10.2</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.hadoop</groupId>
+            <artifactId>hadoop-common</artifactId>
+            <version>3.3.0</version>
+        </dependency>
+
+        <dependency>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-lang3</artifactId>
+            <version>3.12.0</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-collections4</artifactId>
+            <version>4.4</version>
+        </dependency>
+
+
+        <!-- Test Dependencies -->
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter-api</artifactId>
+            <version>5.7.1</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>com.fasterxml.jackson.dataformat</groupId>
+            <artifactId>jackson-dataformat-csv</artifactId>
+            <version>2.12.3</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-core</artifactId>
+            <version>3.9.0</version>
+            <scope>test</scope>
+        </dependency>
+
+    </dependencies>
+
+    <build>
+        <plugins>
+            <!-- Compiles source code -->
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-compiler-plugin</artifactId>
+                <version>3.8.1</version>
+                <configuration>
+                    <source>${java.version}</source>
+                    <target>${java.version}</target>
+                </configuration>
+            </plugin>
+
+            <!-- Runs tests -->
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-surefire-plugin</artifactId>
+                <version>2.22.2</version>
+            </plugin>
+        </plugins>
+    </build>
+
+</project>

spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/item/google/bigquery/BigQueryItemWriter.java

@@ -0,0 +1,439 @@
+/*
+ * Copyright 2002-2021 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.extensions.item.google.bigquery;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.lang.reflect.InvocationTargetException;
+import java.net.URI;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Paths;
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import com.google.cloud.bigquery.BigQuery;
+import com.google.cloud.bigquery.Dataset;
+import com.google.cloud.bigquery.DatasetInfo;
+import com.google.cloud.bigquery.FormatOptions;
+import com.google.cloud.bigquery.Job;
+import com.google.cloud.bigquery.Table;
+import com.google.cloud.bigquery.TableDataWriteChannel;
+import com.google.cloud.bigquery.TableDefinition;
+import com.google.cloud.bigquery.WriteChannelConfiguration;
+import com.google.common.collect.Iterables;
+import org.apache.avro.file.CodecFactory;
+import org.apache.avro.file.DataFileWriter;
+import org.apache.avro.specific.SpecificDatumWriter;
+import org.apache.avro.specific.SpecificRecordBase;
+import org.apache.commons.collections4.CollectionUtils;
+import org.apache.commons.io.IOUtils;
+import org.apache.commons.lang3.ArrayUtils;
+import org.apache.commons.lang3.BooleanUtils;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.apache.hadoop.fs.Path;
+
+import org.springframework.batch.item.ItemWriter;
+import org.springframework.beans.factory.InitializingBean;
+import org.springframework.core.convert.converter.Converter;
+import org.springframework.core.io.FileSystemResource;
+import org.springframework.util.Assert;
+import org.springframework.util.ObjectUtils;
+
+/**
+ * {@link ItemWriter} for Google BigQuery.
+ * This writer uses java client from Google, so we cannot control this flow fully.
+ * Take into account that this writer produces {@link com.google.cloud.bigquery.JobConfiguration.Type#LOAD} {@link Job}.
+ *
+ * <p>Supported formats:
+ * <ul>
+ *     <li>JSON</li>
+ *     <li>CSV</li>
+ *     <li>Avro</li>
+ * </ul>
+ *
+ * <p>For example if you generate {@link TableDataWriteChannel} and you {@link TableDataWriteChannel#close()} it,
+ * there is no guarantee that single {@link com.google.cloud.bigquery.Job} will be created.
+ *
+ * <p>It does not support save state feature. It is thread-safe.
+ *
+ * @author Vova Perebykivskyi
+ * @since 0.1.0
+ * @see <a href="https://cloud.google.com/bigquery">BigQuery</a>
+ * @see <a href="https://github.com/googleapis/java-bigquery">BigQuery Java Client on GitHub</a>
+ */
+public class BigQueryItemWriter<T> implements ItemWriter<T>, InitializingBean {
+
+    private final Log logger = LogFactory.getLog(getClass());
+    private final AtomicLong bigQueryWriteCounter = new AtomicLong();
+
+    /**
+     * Used for simple conversion.
+     */
+    private Converter<T, byte[]> rowMapper;
+
+
+    private BigQuery bigQuery;
+
+    /**
+     * You can specify here some specific dataset configuration, like location.
+     * This dataset will be created.
+     */
+    private DatasetInfo datasetInfo;
+
+    private WriteChannelConfiguration writeChannelConfig;
+
+    /**
+     * Your custom logic with {@link Job}.
+     * {@link Job} will be assigned after {@link TableDataWriteChannel#close()}.
+     */
+    private Consumer<Job> jobConsumer;
+
+    private Table getTable() {
+        return this.bigQuery.getTable(this.writeChannelConfig.getDestinationTable());
+    }
+
+    public void setRowMapper(Converter<T, byte[]> rowMapper) {
+        this.rowMapper = rowMapper;
+    }
+
+    public void setDatasetInfo(DatasetInfo datasetInfo) {
+        this.datasetInfo = datasetInfo;
+    }
+
+    public void setJobConsumer(Consumer<Job> consumer) {
+        this.jobConsumer = consumer;
+    }
+
+    public void setWriteChannelConfig(WriteChannelConfiguration writeChannelConfig) {
+        this.writeChannelConfig = writeChannelConfig;
+    }
+
+    public void setBigQuery(BigQuery bigQuery) {
+        this.bigQuery = bigQuery;
+    }
+
+    @Override
+    public void write(List<? extends T> items) throws Exception {
+        if (CollectionUtils.isNotEmpty(items)) {
+            if (this.logger.isDebugEnabled()) {
+                this.logger.debug("Mapping data");
+            }
+
+            ByteBuffer byteBuffer = mapDataToBigQueryFormat(items);
+            doWriteDataToBigQuery(byteBuffer);
+        }
+    }
+
+    private ByteBuffer mapDataToBigQueryFormat(List<? extends T> items)
+            throws IOException, InvocationTargetException, NoSuchMethodException, InstantiationException, IllegalAccessException {
+
+        ByteBuffer byteBuffer;
+        try (ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {
+
+            List<byte[]> data = convertObjectsToByteArrays(items);
+
+            for (byte[] byteArray : data) {
+                outputStream.write(byteArray);
+            }
+
+            /*
+             * It is extremely important to create larger ByteBuffer,
+             * if you call TableDataWriteChannel too many times, it leads to BigQuery exceptions.
+             */
+            byteBuffer = ByteBuffer.wrap(outputStream.toByteArray());
+        }
+        return byteBuffer;
+    }
+
+    private void doWriteDataToBigQuery(ByteBuffer byteBuffer) throws IOException {
+        if (this.logger.isDebugEnabled()) {
+            this.logger.debug("Writing data to BigQuery");
+        }
+
+        TableDataWriteChannel writeChannel = null;
+
+        try (TableDataWriteChannel writer = getWriteChannel()) {
+            /* TableDataWriteChannel is not thread safe */
+            writer.write(byteBuffer);
+            writeChannel = writer;
+        }
+        finally {
+            if (this.logger.isDebugEnabled()) {
+                this.logger.debug("Write operation submitted: " + bigQueryWriteCounter.incrementAndGet());
+            }
+
+            if (Objects.nonNull(writeChannel) && Objects.nonNull(this.jobConsumer)) {
+                this.jobConsumer.accept(writeChannel.getJob());
+            }
+        }
+    }
+
+    private List<byte[]> convertObjectsToByteArrays(List<? extends T> items)
+            throws IOException, NoSuchMethodException, IllegalAccessException, InvocationTargetException, InstantiationException {
+
+        Stream<byte[]> byteArrayStream = Stream.empty();
+
+        if (isJson()) {
+            byteArrayStream = getJsonByteArrayStream(items);
+        }
+        else if (isCsv()) {
+            byteArrayStream = getCsvByteArrayStream(items);
+        }
+        else if (isAvro()) {
+            byteArrayStream = getAvroByteArrayStream(items);
+        }
+        else if (isParquet() || isOrc()) {
+            throw new UnsupportedOperationException("Not supported right now");
+            /*byteArrayStream = getHadoopPathByteArrayStream(items);*/
+        }
+
+        return byteArrayStream.collect(Collectors.toList());
+    }
+
+    /**
+     * Row could be read as typical {@link String}.
+     */
+    private Stream<byte[]> getJsonByteArrayStream(List<? extends T> items) {
+        return items
+                .stream()
+                .map(this.rowMapper::convert)
+                .filter(ArrayUtils::isNotEmpty)
+                .map(String::new)
+                .map(this::convertToNdJson)
+                .filter(value -> !ObjectUtils.isEmpty(value))
+                .map(row -> row.getBytes(StandardCharsets.UTF_8));
+    }
+
+    /**
+     * BigQuery uses ndjson https://github.com/ndjson/ndjson-spec.
+     * It is expected that to pass here JSON line generated by
+     * {@link com.fasterxml.jackson.databind.ObjectMapper} or any other JSON parser.
+     */
+    private String convertToNdJson(String json) {
+        return json.concat(org.apache.commons.lang3.StringUtils.LF);
+    }
+
+    /**
+     * Row could be read as typical {@link String}.
+     */
+    private Stream<byte[]> getCsvByteArrayStream(List<? extends T> items) {
+        return items
+                .stream()
+                .map(this.rowMapper::convert)
+                .filter(ArrayUtils::isNotEmpty)
+                .map(String::new)
+                .filter(value -> !ObjectUtils.isEmpty(value))
+                .map(row -> row.getBytes(StandardCharsets.UTF_8));
+    }
+
+    /**
+     * Generates Avro file and writes it to {@link OutputStream}.
+     * @see <a href="https://github.com/GoogleCloudPlatform/bigquery-ingest-avro-dataflow-sample">Avro example</a>
+     */
+    private Stream<byte[]> getAvroByteArrayStream(List<? extends T> items)
+            throws IOException, NoSuchMethodException, IllegalAccessException, InvocationTargetException, InstantiationException {
+
+        T firstElement = Iterables.getFirst(items, null);
+        Assert.notNull(firstElement, "Collection is empty");
+
+        Class classType = firstElement.getClass();
+
+        SpecificRecordBase objectInstance = (SpecificRecordBase) classType.getDeclaredConstructor().newInstance();
+        SpecificDatumWriter<? super T> avroWriter = new SpecificDatumWriter<>(classType);
+
+        try (ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
+             DataFileWriter<? super T> avroFileWriter = new DataFileWriter<>(avroWriter)) {
+
+            /*
+             * Input data - 500 rows.
+             * Statistic gathered from BigQuery (com.google.cloud.bigquery.JobStatistics.LoadStatistics#getInputBytes).
+             *
+             * 41,229      input bytes Avro (no codec)
+             * 14,691      input bytes Avro (Deflate lvl 8)
+             * 41,122      input bytes CSV
+             */
+            CodecFactory compressionCodec = CodecFactory.deflateCodec(8);
+
+            /* Order of lines (code) should not be changed */
+            avroFileWriter.setCodec(compressionCodec);
+            avroFileWriter.create(objectInstance.getSchema(), outputStream);
+
+            for (T item : items) {
+                avroFileWriter.append(item);
+            }
+
+            /* At this point of time only schema present in output stream */
+            avroFileWriter.flush();
+
+            return Stream.of(outputStream.toByteArray());
+        }
+    }
+
+    /**
+     * @return {@link TableDataWriteChannel} that should be closed manually.
+     * @see <a href="https://github.com/googleapis/google-cloud-java/blob/969bbeef18f004fd51fd46c5def1ae5c644cae3c/google-cloud-examples/src/main/java/com/google/cloud/examples/bigquery/snippets/BigQuerySnippets.java">Examples</a>
+     */
+    private TableDataWriteChannel getWriteChannel() {
+        return this.bigQuery.writer(this.writeChannelConfig);
+    }
+
+    @Override
+    public void afterPropertiesSet() {
+        Assert.isTrue(BooleanUtils.isFalse(isBigtable()), "Google BigTable is not supported");
+        Assert.isTrue(BooleanUtils.isFalse(isGoogleSheets()), "Google Sheets is not supported");
+        Assert.isTrue(BooleanUtils.isFalse(isDatastore()), "Google Datastore is not supported");
+        Assert.isTrue(BooleanUtils.isFalse(isParquet()), "Parquet is not supported");
+        Assert.isTrue(BooleanUtils.isFalse(isOrc()), "Orc is not supported");
+
+        Assert.notNull(this.bigQuery, "BigQuery service must be provided");
+        Assert.notNull(this.writeChannelConfig, "Write channel configuration must be provided");
+
+        if (BooleanUtils.isFalse(isAvro())) {
+            Table table = getTable();
+            if (BooleanUtils.isFalse(BooleanUtils.toBoolean(this.writeChannelConfig.getAutodetect()))) {
+                Assert.notNull(this.writeChannelConfig.getSchema(), "Schema must be provided");
+                if (tableHasDefinedSchema(table)) {
+                    Assert.isTrue(
+                            table.getDefinition().getSchema().equals(this.writeChannelConfig.getSchema()),
+                            "Schema should be the same"
+                    );
+                }
+            }
+            else {
+                if ((isCsv() || isJson()) && this.logger.isWarnEnabled() && tableHasDefinedSchema(table)) {
+                    this.logger.warn("Mixing autodetect mode with already defined schema may lead to errors on BigQuery side");
+                }
+            }
+        }
+        else {
+            Assert.isNull(this.writeChannelConfig.getSchema(), "Avro does not require schema");
+            Assert.isNull(this.writeChannelConfig.getAutodetect(), "Avro does not require autodetection");
+        }
+
+        Assert.notNull(this.writeChannelConfig.getFormat(), "Data format must be provided");
+
+        if (isCsv() || isJson()) {
+            Assert.notNull(this.rowMapper, "Row mapper must be provided");
+        }
+
+        if (Objects.nonNull(this.datasetInfo)) {
+            Assert.isTrue(
+                    Objects.equals(this.datasetInfo.getDatasetId().getDataset(), this.writeChannelConfig.getDestinationTable().getDataset()),
+                    "Dataset should be configured properly"
+            );
+        }
+
+        createDataset();
+    }
+
+    private boolean isCsv() {
+        return FormatOptions.csv().getType().equals(this.writeChannelConfig.getFormat());
+    }
+
+    private boolean isJson() {
+        return FormatOptions.json().getType().equals(this.writeChannelConfig.getFormat());
+    }
+
+    private boolean isAvro() {
+        return FormatOptions.avro().getType().equals(this.writeChannelConfig.getFormat());
+    }
+
+    private boolean isParquet() {
+        return FormatOptions.parquet().getType().equals(this.writeChannelConfig.getFormat());
+    }
+
+    private boolean isOrc() {
+        return FormatOptions.orc().getType().equals(this.writeChannelConfig.getFormat());
+    }
+
+    private boolean isBigtable() {
+        return FormatOptions.bigtable().getType().equals(this.writeChannelConfig.getFormat());
+    }
+
+    private boolean isGoogleSheets() {
+        return FormatOptions.googleSheets().getType().equals(this.writeChannelConfig.getFormat());
+    }
+
+    private boolean isDatastore() {
+        return FormatOptions.datastoreBackup().getType().equals(this.writeChannelConfig.getFormat());
+    }
+
+    private boolean tableHasDefinedSchema(Table table) {
+        return Optional
+                .ofNullable(table)
+                .map(Table::getDefinition)
+                .map(TableDefinition.class::cast)
+                .map(TableDefinition::getSchema)
+                .isPresent();
+    }
+
+    private void createDataset() {
+        if (Objects.nonNull(this.datasetInfo)) {
+            Dataset foundDataset = this.bigQuery.getDataset(this.writeChannelConfig.getDestinationTable().getDataset());
+
+            if (Objects.isNull(foundDataset)) {
+                this.bigQuery.create(this.datasetInfo);
+            }
+        }
+    }
+
+    /**
+     * It is expected that for {@link BigQueryItemWriter#isParquet()} and {@link BigQueryItemWriter#isOrc()} you use
+     * {@link Path} because there is no convenient way to write records to {@link OutputStream}.
+     * Not supported right now.
+     */
+    private Stream<byte[]> getHadoopPathByteArrayStream(List<? extends T> items) throws IOException {
+        Stream<byte[]> result = Stream.empty();
+
+        T firstElement = Iterables.getFirst(items, null);
+        Assert.notNull(firstElement, "Collection is empty");
+        Class<?> classType = firstElement.getClass();
+
+        if (classType.isAssignableFrom(Path.class)) {
+            List<URI> uris = items.stream()
+                    .map(Path.class::cast)
+                    .map(Path::toUri)
+                    .collect(Collectors.toList());
+
+            return getFileBasedByteArrayStream(uris);
+        }
+
+        return result;
+    }
+
+    private Stream<byte[]> getFileBasedByteArrayStream(List<URI> items) throws IOException {
+        try (ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {
+            for (URI uri : items) {
+                try (InputStream inputStream = new FileSystemResource(Paths.get(uri)).getInputStream()) {
+                    IOUtils.copy(inputStream, outputStream);
+                }
+            }
+            return Stream.of(outputStream.toByteArray());
+        }
+    }
+
+}

spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/item/google/bigquery/builder/BigQueryItemWriterBuilder.java

@@ -0,0 +1,83 @@
+/*
+ * Copyright 2002-2021 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.extensions.item.google.bigquery.builder;
+
+import java.util.function.Consumer;
+
+import com.google.cloud.bigquery.BigQuery;
+import com.google.cloud.bigquery.DatasetInfo;
+import com.google.cloud.bigquery.Job;
+import com.google.cloud.bigquery.WriteChannelConfiguration;
+
+import org.springframework.batch.extensions.item.google.bigquery.BigQueryItemWriter;
+import org.springframework.core.convert.converter.Converter;
+
+/**
+ * A builder for {@link BigQueryItemWriter}.
+ *
+ * @author Vova Perebykivskyi
+ * @since 0.1.0
+ * @see BigQueryItemWriter
+ * @see <a href="https://github.com/spring-projects/spring-batch-extensions/tree/master/spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/item/google/bigquery/builder/BigQueryItemWriterBuilderTests.java">Examples</a>
+ */
+public class BigQueryItemWriterBuilder<T> {
+
+    private Converter<T, byte[]> rowMapper;
+    private Consumer<Job> jobConsumer;
+
+    private DatasetInfo datasetInfo;
+    private WriteChannelConfiguration writeChannelConfig;
+    private BigQuery bigQuery;
+
+    public BigQueryItemWriterBuilder<T> rowMapper(Converter<T, byte[]> rowMapper) {
+        this.rowMapper = rowMapper;
+        return this;
+    }
+
+    public BigQueryItemWriterBuilder<T> datasetInfo(DatasetInfo datasetInfo) {
+        this.datasetInfo = datasetInfo;
+        return this;
+    }
+
+    public BigQueryItemWriterBuilder<T> jobConsumer(Consumer<Job> consumer) {
+        this.jobConsumer = consumer;
+        return this;
+    }
+
+    public BigQueryItemWriterBuilder<T> writeChannelConfig(WriteChannelConfiguration configuration) {
+        this.writeChannelConfig = configuration;
+        return this;
+    }
+
+    public BigQueryItemWriterBuilder<T> bigQuery(BigQuery bigQuery) {
+        this.bigQuery = bigQuery;
+        return this;
+    }
+
+    public BigQueryItemWriter<T> build() {
+        BigQueryItemWriter<T> writer = new BigQueryItemWriter<>();
+
+        writer.setRowMapper(this.rowMapper);
+        writer.setDatasetInfo(this.datasetInfo);
+        writer.setJobConsumer(this.jobConsumer);
+        writer.setWriteChannelConfig(this.writeChannelConfig);
+        writer.setBigQuery(this.bigQuery);
+
+        return writer;
+    }
+
+}

spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/item/google/bigquery/package-info.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2002-2021 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.
+ */
+
+/**
+ * Google BigQuery related functionality.
+ *
+ * @see <a href="https://cloud.google.com/bigquery/">Google BigQuery</a>
+ * @author Vova Perebykivskyi
+ */
+@NonNullApi
+package org.springframework.batch.extensions.item.google.bigquery;
+
+import org.springframework.lang.NonNullApi;

spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/item/google/bigquery/builder/BigQueryItemWriterBuilderTests.java

@@ -0,0 +1,161 @@
+/*
+ * Copyright 2002-2021 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.extensions.item.google.bigquery.builder;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.dataformat.csv.CsvMapper;
+import com.google.cloud.bigquery.BigQuery;
+import com.google.cloud.bigquery.DatasetInfo;
+import com.google.cloud.bigquery.Field;
+import com.google.cloud.bigquery.FormatOptions;
+import com.google.cloud.bigquery.Schema;
+import com.google.cloud.bigquery.StandardSQLTypeName;
+import com.google.cloud.bigquery.TableId;
+import com.google.cloud.bigquery.WriteChannelConfiguration;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+
+import org.springframework.batch.extensions.item.google.bigquery.BigQueryItemWriter;
+
+class BigQueryItemWriterBuilderTests {
+
+    private static final String DATASET_NAME = "my_dataset";
+
+    private final Log logger = LogFactory.getLog(getClass());
+
+    /**
+     * Example how CSV writer is expected to be built without {@link org.springframework.context.annotation.Bean} annotation.
+     */
+    @Test
+    void testCsvWriter() {
+        BigQuery mockedBigQuery = prepareMockedBigQuery();
+        CsvMapper csvMapper = new CsvMapper();
+
+        WriteChannelConfiguration writeConfiguration = WriteChannelConfiguration
+                .newBuilder(TableId.of(DATASET_NAME, "csv_table"))
+                .setAutodetect(true)
+                .setFormatOptions(FormatOptions.csv())
+                .build();
+
+        BigQueryItemWriter<PersonDto> writer = new BigQueryItemWriterBuilder<PersonDto>()
+                .bigQuery(mockedBigQuery)
+                .rowMapper(dto -> convertDtoToCsvByteArray(csvMapper, dto))
+                .writeChannelConfig(writeConfiguration)
+                .datasetInfo(DatasetInfo.newBuilder(DATASET_NAME).setLocation("europe-west-2").build())
+                .build();
+
+        writer.afterPropertiesSet();
+
+        Assertions.assertNotNull(writer);
+    }
+
+    /**
+     * Example how JSON writer is expected to be built without {@link org.springframework.context.annotation.Bean} annotation.
+     */
+    @Test
+    void testJsonWriter() {
+        BigQuery mockedBigQuery = prepareMockedBigQuery();
+        ObjectMapper objectMapper = new ObjectMapper();
+
+        WriteChannelConfiguration writeConfiguration = WriteChannelConfiguration
+                .newBuilder(TableId.of(DATASET_NAME, "json_table"))
+                .setFormatOptions(FormatOptions.json())
+                .setSchema(Schema.of(
+                        Field.newBuilder("name", StandardSQLTypeName.STRING).setMode(Field.Mode.REQUIRED).build()
+                ))
+                .build();
+
+        BigQueryItemWriter<PersonDto> writer = new BigQueryItemWriterBuilder<PersonDto>()
+                .bigQuery(mockedBigQuery)
+                .rowMapper(dto -> convertDtoToJsonByteArray(objectMapper, dto))
+                .writeChannelConfig(writeConfiguration)
+                .jobConsumer(job -> this.logger.debug("Job with id: " + job.getJobId() + " is created"))
+                .build();
+
+        writer.afterPropertiesSet();
+
+        Assertions.assertNotNull(writer);
+    }
+
+    /**
+     * Example how Apache Avro writer is expected to be built without {@link org.springframework.context.annotation.Bean} annotation.
+     */
+    @Test
+    void testAvroWriter() {
+        BigQuery mockedBigQuery = prepareMockedBigQuery();
+
+        WriteChannelConfiguration writeConfiguration = WriteChannelConfiguration
+                .newBuilder(TableId.of(DATASET_NAME, "avro_table"))
+                .setFormatOptions(FormatOptions.avro())
+                .build();
+
+        BigQueryItemWriter<PersonDto> writer = new BigQueryItemWriterBuilder<PersonDto>()
+                .bigQuery(mockedBigQuery)
+                .writeChannelConfig(writeConfiguration)
+                .build();
+
+        writer.afterPropertiesSet();
+
+        Assertions.assertNotNull(writer);
+    }
+
+    private byte[] convertDtoToJsonByteArray(ObjectMapper objectMapper, PersonDto dto)  {
+        try {
+            return objectMapper.writeValueAsBytes(dto);
+        } catch (JsonProcessingException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private byte[] convertDtoToCsvByteArray(CsvMapper csvMapper, PersonDto dto) {
+        try {
+            return csvMapper.writerWithSchemaFor(PersonDto.class).writeValueAsBytes(dto);
+        }
+        catch (JsonProcessingException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private BigQuery prepareMockedBigQuery() {
+        BigQuery mockedBigQuery = Mockito.mock(BigQuery.class);
+
+        Mockito
+                .when(mockedBigQuery.getTable(Mockito.any()))
+                .thenReturn(null);
+
+        Mockito
+                .when(mockedBigQuery.getDataset(Mockito.anyString()))
+                .thenReturn(null);
+
+        return mockedBigQuery;
+    }
+
+
+    static class PersonDto {
+
+        private String name;
+
+        public String getName() {
+            return name;
+        }
+    }
+
+}