Spring/spring-batch-extensions
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[bq] 0.2 upgrade dependencies & refactor writers

Browse Source
This commit is contained in:
Volodymyr
2022-12-17 11:43:32 +02:00
committed by GitHub
parent ce115bbfe8
commit 79465739ce
12 changed files with 566 additions and 328 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
CONTRIBUTING.md
View File

@@ -114,7 +114,7 @@ Please carefully follow the same [code style as Spring Framework](https://github
```java
/*
* Copyright 2002-2021 the original author or authors.
* Copyright 2002-2022 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.

12
spring-batch-bigquery/README.adoc
View File

@@ -2,27 +2,27 @@
Spring Batch extension which contains an `ItemWriter` implementation for https://cloud.google.com/bigquery[BigQuery] based on https://github.com/googleapis/java-bigquery[Java BigQuery]. It supports writing https://en.wikipedia.org/wiki/Comma-separated_values[CSV], https://en.wikipedia.org/wiki/JSON[JSON] using https://cloud.google.com/bigquery/docs/batch-loading-data[load jobs].
## Configuration of `BigQueryItemWriter`
## Configuration of `BigQueryCsvItemWriter`
Next to the https://docs.spring.io/spring-batch/reference/html/configureJob.html[configuration of Spring Batch] one needs to configure the `BigQueryItemWriter`.
Next to the https://docs.spring.io/spring-batch/reference/html/configureJob.html[configuration of Spring Batch] one needs to configure the `BigQueryCsvItemWriter`.
```java
```javaBigQueryCsv
@Bean
BigQueryItemWriter<MyDto> bigQueryCsvWriter() {
BigQueryCsvItemWriter<MyDto> bigQueryCsvWriter() {
WriteChannelConfiguration writeConfiguration = WriteChannelConfiguration
.newBuilder(TableId.of("csv_dataset", "csv_table"))
.setAutodetect(true)
.setFormatOptions(FormatOptions.csv())
.build();
BigQueryItemWriter<MyDto> writer = new BigQueryItemWriterBuilder<MyDto>()
BigQueryCsvItemWriter<MyDto> writer = new BigQueryCsvItemWriterBuilder<MyDto>()
.bigQuery(mockedBigQuery)
.writeChannelConfig(writeConfiguration)
.build();
}
```
Additional examples could be found in https://github.com/spring-projects/spring-batch-extensions/blob/main/spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/builder/BigQueryItemWriterBuilderTests.java[here].
Additional examples could be found in https://github.com/spring-projects/spring-batch-extensions/blob/main/spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/writer/builder/[here].
## Configuration properties
[cols="1,1,4"]

20
spring-batch-bigquery/pom.xml
View File

@@ -1,6 +1,6 @@
<?xml version="1.0" encoding="UTF-8"?>
<!--
~ Copyright 2002-2021 the original author or authors.
~ Copyright 2002-2022 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.
@@ -27,7 +27,7 @@
<developers>
<developer>
<id>Dgray16</id>
<name>Vova Perebykivskyi</name>
<name>Volodymyr Perebykivskyi</name>
<email>vova235@gmail.com</email>
</developer>
</developers>
@@ -50,25 +50,25 @@
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
<!-- Dependent on Spring Batch core -->
<java.version>8</java.version>
<java.version>17</java.version>
</properties>
<dependencies>
<dependency>
<groupId>org.springframework.batch</groupId>
<artifactId>spring-batch-core</artifactId>
<version>4.3.3</version>
<version>5.0.0</version>
</dependency>
<dependency>
<groupId>com.google.cloud</groupId>
<artifactId>google-cloud-bigquery</artifactId>
<version>1.133.0</version>
<version>2.19.1</version>
</dependency>
<dependency>
<groupId>com.fasterxml.jackson.dataformat</groupId>
<artifactId>jackson-dataformat-csv</artifactId>
<version>2.12.3</version>
<version>2.14.1</version>
</dependency>
<dependency>
@@ -87,13 +87,13 @@
<dependency>
<groupId>org.junit.jupiter</groupId>
<artifactId>junit-jupiter-api</artifactId>
<version>5.7.2</version>
<version>5.9.1</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.mockito</groupId>
<artifactId>mockito-core</artifactId>
<version>3.11.1</version>
<version>4.9.0</version>
<scope>test</scope>
</dependency>
@@ -105,7 +105,7 @@
<plugin>
<groupId>org.apache.maven.plugins</groupId>
<artifactId>maven-compiler-plugin</artifactId>
<version>3.8.1</version>
<version>3.10.1</version>
<configuration>
<source>${java.version}</source>
<target>${java.version}</target>
@@ -123,7 +123,7 @@
<plugin>
<groupId>org.apache.maven.plugins</groupId>
<artifactId>maven-javadoc-plugin</artifactId>
<version>3.2.0</version>
<version>3.4.1</version>
<executions>
<execution>
<id>attach-javadocs</id>

26
spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/package-info.java
View File

@@ -1,26 +0,0 @@
/*
* 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.bigquery;
import org.springframework.lang.NonNullApi;

245
spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/BigQueryItemWriter.java → spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/writer/BigQueryBaseItemWriter.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2002-2021 the original author or authors.
* Copyright 2002-2022 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.
@@ -14,24 +14,8 @@
* limitations under the License.
*/
package org.springframework.batch.extensions.bigquery;
package org.springframework.batch.extensions.bigquery.writer;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
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.stream.Collectors;
import java.util.stream.Stream;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.ObjectWriter;
import com.fasterxml.jackson.dataformat.csv.CsvMapper;
import com.google.cloud.bigquery.BigQuery;
import com.google.cloud.bigquery.Dataset;
import com.google.cloud.bigquery.DatasetInfo;
@@ -42,75 +26,45 @@ import com.google.cloud.bigquery.TableDataWriteChannel;
import com.google.cloud.bigquery.TableDefinition;
import com.google.cloud.bigquery.TableId;
import com.google.cloud.bigquery.WriteChannelConfiguration;
import org.apache.commons.collections4.CollectionUtils;
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.springframework.batch.item.Chunk;
import org.springframework.batch.item.ItemWriter;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.core.convert.converter.Converter;
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>
* </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.
* Take into account that BigQuery has rate limits and it is very easy to exceed those in concurrent environment.
* @see <a href="https://cloud.google.com/bigquery/quotas">BigQuery Quotas &amp; Limits</a>
*
* @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 {
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
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.Supplier;
private final Log logger = LogFactory.getLog(getClass());
public abstract class BigQueryBaseItemWriter<T> implements ItemWriter<T> {
protected final Log logger = LogFactory.getLog(getClass());
private final AtomicLong bigQueryWriteCounter = new AtomicLong();
/**
* Used for simple conversion.
*/
private Converter<T, byte[]> rowMapper;
private ObjectWriter objectWriter;
private Class itemClass;
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());
}
protected WriteChannelConfiguration writeChannelConfig;
private BigQuery bigQuery;
public void setRowMapper(Converter<T, byte[]> rowMapper) {
this.rowMapper = rowMapper;
protected Table getTable() {
return this.bigQuery.getTable(this.writeChannelConfig.getDestinationTable());
}
public void setDatasetInfo(DatasetInfo datasetInfo) {
@@ -130,9 +84,10 @@ public class BigQueryItemWriter<T> implements ItemWriter<T>, InitializingBean {
}
@Override
public void write(List<? extends T> items) throws Exception {
if (CollectionUtils.isNotEmpty(items)) {
initializeProperties(items);
public void write(Chunk<? extends T> chunk) throws Exception {
if (BooleanUtils.isFalse(chunk.isEmpty())) {
List<? extends T> items = chunk.getItems();
doInitializeProperties(items);
if (this.logger.isDebugEnabled()) {
this.logger.debug(String.format("Mapping %d elements", items.size()));
@@ -143,27 +98,6 @@ public class BigQueryItemWriter<T> implements ItemWriter<T>, InitializingBean {
}
}
/** Actual type of incoming data can be obtained only in runtime */
private synchronized void initializeProperties(List<? extends T> items) {
if (Objects.isNull(this.itemClass)) {
if (isCsv() || isJson()) {
T firstItem = items.stream().findFirst().orElseThrow(RuntimeException::new);
this.itemClass = firstItem.getClass();
if (Objects.isNull(this.rowMapper)) {
if (isCsv()) {
this.objectWriter = new CsvMapper().writerWithTypedSchemaFor(this.itemClass);
}
else if (isJson()) {
this.objectWriter = new ObjectMapper().writerFor(this.itemClass);
}
}
logger.debug("Writer setup is completed");
}
}
}
private ByteBuffer mapDataToBigQueryFormat(List<? extends T> items) throws IOException {
ByteBuffer byteBuffer;
try (ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {
@@ -211,69 +145,6 @@ public class BigQueryItemWriter<T> implements ItemWriter<T>, InitializingBean {
}
}
private List<byte[]> convertObjectsToByteArrays(List<? extends T> items) {
Stream<byte[]> byteArrayStream = Stream.empty();
if (isJson()) {
byteArrayStream = getJsonByteArrayStream(items);
}
else if (isCsv()) {
byteArrayStream = getCsvByteArrayStream(items);
}
else if (isParquet() || isOrc() || isAvro()) {
throw new UnsupportedOperationException("Not supported right now");
}
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::mapItemToCsvOrJson)
.filter(ArrayUtils::isNotEmpty)
.map(String::new)
.map(this::convertToNdJson)
.filter(value -> !ObjectUtils.isEmpty(value))
.map(row -> row.getBytes(StandardCharsets.UTF_8));
}
private byte[] mapItemToCsvOrJson(T t) {
byte[] result = null;
try {
result = Objects.isNull(rowMapper) ? objectWriter.writeValueAsBytes(t) : rowMapper.convert(t);
}
catch (JsonProcessingException e) {
logger.error("Error during processing of the line: ", e);
}
return result;
}
/**
* 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::mapItemToCsvOrJson)
.filter(ArrayUtils::isNotEmpty)
.map(String::new)
.filter(value -> !ObjectUtils.isEmpty(value))
.map(row -> row.getBytes(StandardCharsets.UTF_8));
}
/**
* @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>
@@ -282,8 +153,7 @@ public class BigQueryItemWriter<T> implements ItemWriter<T>, InitializingBean {
return this.bigQuery.writer(this.writeChannelConfig);
}
@Override
public void afterPropertiesSet() {
protected void baseAfterPropertiesSet(Supplier<Void> formatSpecificChecks) {
Assert.notNull(this.bigQuery, "BigQuery service must be provided");
Assert.notNull(this.writeChannelConfig, "Write channel configuration must be provided");
@@ -294,33 +164,11 @@ public class BigQueryItemWriter<T> implements ItemWriter<T>, InitializingBean {
Assert.isTrue(BooleanUtils.isFalse(isOrc()), "Orc is not supported");
Assert.isTrue(BooleanUtils.isFalse(isAvro()), "Avro is not supported");
if (BooleanUtils.isFalse(isAvro())) {
Table table = getTable();
if (BooleanUtils.toBoolean(this.writeChannelConfig.getAutodetect())) {
if ((isCsv() || isJson()) && tableHasDefinedSchema(table) && this.logger.isWarnEnabled()) {
this.logger.warn("Mixing autodetect mode with already defined schema may lead to errors on BigQuery side");
}
}
else {
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 {
Assert.isNull(this.writeChannelConfig.getSchema(), "Avro does not require schema");
Assert.isNull(this.writeChannelConfig.getAutodetect(), "Avro does not require autodetection");
}
formatSpecificChecks.get();
Assert.notNull(this.writeChannelConfig.getFormat(), "Data format must be provided");
String dataset = this.writeChannelConfig.getDestinationTable().getDataset();
if (Objects.isNull(this.datasetInfo)) {
this.datasetInfo = DatasetInfo.newBuilder(dataset).build();
}
@@ -333,12 +181,19 @@ public class BigQueryItemWriter<T> implements ItemWriter<T>, InitializingBean {
createDataset();
}
private boolean isCsv() {
return FormatOptions.csv().getType().equals(this.writeChannelConfig.getFormat());
}
private void createDataset() {
TableId tableId = this.writeChannelConfig.getDestinationTable();
String datasetToCheck = tableId.getDataset();
private boolean isJson() {
return FormatOptions.json().getType().equals(this.writeChannelConfig.getFormat());
if (Objects.nonNull(datasetToCheck)) {
Dataset foundDataset = this.bigQuery.getDataset(datasetToCheck);
if (Objects.isNull(foundDataset)) {
if (Objects.nonNull(this.datasetInfo)) {
this.bigQuery.create(this.datasetInfo);
}
}
}
}
private boolean isAvro() {
@@ -365,7 +220,15 @@ public class BigQueryItemWriter<T> implements ItemWriter<T>, InitializingBean {
return FormatOptions.datastoreBackup().getType().equals(this.writeChannelConfig.getFormat());
}
private boolean tableHasDefinedSchema(Table table) {
protected boolean isCsv() {
return FormatOptions.csv().getType().equals(this.writeChannelConfig.getFormat());
}
protected boolean isJson() {
return FormatOptions.json().getType().equals(this.writeChannelConfig.getFormat());
}
protected boolean tableHasDefinedSchema(Table table) {
return Optional
.ofNullable(table)
.map(Table::getDefinition)
@@ -374,19 +237,7 @@ public class BigQueryItemWriter<T> implements ItemWriter<T>, InitializingBean {
.isPresent();
}
private void createDataset() {
TableId tableId = this.writeChannelConfig.getDestinationTable();
String datasetToCheck = tableId.getDataset();
if (Objects.nonNull(datasetToCheck)) {
Dataset foundDataset = this.bigQuery.getDataset(datasetToCheck);
if (Objects.isNull(foundDataset)) {
if (Objects.nonNull(this.datasetInfo)) {
this.bigQuery.create(this.datasetInfo);
}
}
}
}
protected abstract void doInitializeProperties(List<? extends T> items);
protected abstract List<byte[]> convertObjectsToByteArrays(List<? extends T> items);
}

110
spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/writer/BigQueryCsvItemWriter.java Normal file
View File

@@ -0,0 +1,110 @@
/*
* Copyright 2002-2022 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.bigquery.writer;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectWriter;
import com.fasterxml.jackson.dataformat.csv.CsvMapper;
import com.google.cloud.bigquery.Table;
import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.lang3.BooleanUtils;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.core.convert.converter.Converter;
import org.springframework.util.Assert;
import org.springframework.util.ObjectUtils;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;
public class BigQueryCsvItemWriter<T> extends BigQueryBaseItemWriter<T> implements InitializingBean {
protected Converter<T, byte[]> rowMapper;
protected ObjectWriter objectWriter;
protected Class itemClass;
/**
* Actual type of incoming data can be obtained only in runtime
*/
@Override
protected synchronized void doInitializeProperties(List<? extends T> items) {
if (Objects.isNull(this.itemClass)) {
T firstItem = items.stream().findFirst().orElseThrow(RuntimeException::new);
this.itemClass = firstItem.getClass();
if (Objects.isNull(this.rowMapper)) {
this.objectWriter = new CsvMapper().writerWithTypedSchemaFor(this.itemClass);
}
logger.debug("Writer setup is completed");
}
}
public void setRowMapper(Converter<T, byte[]> rowMapper) {
this.rowMapper = rowMapper;
}
@Override
protected List<byte[]> convertObjectsToByteArrays(List<? extends T> items) {
return items
.stream()
.map(this::mapItemToCsv)
.filter(ArrayUtils::isNotEmpty)
.map(String::new)
.filter(value -> !ObjectUtils.isEmpty(value))
.map(row -> row.getBytes(StandardCharsets.UTF_8))
.collect(Collectors.toList());
}
@Override
public void afterPropertiesSet() {
super.baseAfterPropertiesSet(() -> {
Table table = getTable();
if (BooleanUtils.toBoolean(super.writeChannelConfig.getAutodetect())) {
if ((tableHasDefinedSchema(table) && super.logger.isWarnEnabled())) {
super.logger.warn("Mixing autodetect mode with already defined schema may lead to errors on BigQuery side");
}
} else {
Assert.notNull(super.writeChannelConfig.getSchema(), "Schema must be provided");
if (tableHasDefinedSchema(table)) {
Assert.isTrue(
table.getDefinition().getSchema().equals(super.writeChannelConfig.getSchema()),
"Schema should be the same"
);
}
}
return null;
});
}
protected byte[] mapItemToCsv(T t) {
byte[] result = null;
try {
result = Objects.isNull(rowMapper) ? objectWriter.writeValueAsBytes(t) : rowMapper.convert(t);
}
catch (JsonProcessingException e) {
logger.error("Error during processing of the line: ", e);
}
return result;
}
}

116
spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/writer/BigQueryJsonItemWriter.java Normal file
View File

@@ -0,0 +1,116 @@
/*
* Copyright 2002-2022 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.bigquery.writer;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.ObjectWriter;
import com.google.cloud.bigquery.Table;
import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.lang3.BooleanUtils;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.core.convert.converter.Converter;
import org.springframework.util.Assert;
import org.springframework.util.ObjectUtils;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;
public class BigQueryJsonItemWriter<T> extends BigQueryBaseItemWriter<T> implements InitializingBean {
protected Converter<T, byte[]> rowMapper;
protected ObjectWriter objectWriter;
protected Class itemClass;
@Override
protected void doInitializeProperties(List<? extends T> items) {
if (Objects.isNull(this.itemClass)) {
T firstItem = items.stream().findFirst().orElseThrow(RuntimeException::new);
this.itemClass = firstItem.getClass();
if (Objects.isNull(this.rowMapper)) {
this.objectWriter = new ObjectMapper().writerFor(this.itemClass);
}
super.logger.debug("Writer setup is completed");
}
}
public void setRowMapper(Converter<T, byte[]> rowMapper) {
this.rowMapper = rowMapper;
}
@Override
protected List<byte[]> convertObjectsToByteArrays(List<? extends T> items) {
return items
.stream()
.map(this::mapItemToJson)
.filter(ArrayUtils::isNotEmpty)
.map(String::new)
.map(this::convertToNdJson)
.filter(value -> !ObjectUtils.isEmpty(value))
.map(row -> row.getBytes(StandardCharsets.UTF_8))
.collect(Collectors.toList());
}
@Override
public void afterPropertiesSet() {
super.baseAfterPropertiesSet(() -> {
Table table = getTable();
if (BooleanUtils.toBoolean(super.writeChannelConfig.getAutodetect())) {
if ((tableHasDefinedSchema(table) && super.logger.isWarnEnabled())) {
super.logger.warn("Mixing autodetect mode with already defined schema may lead to errors on BigQuery side");
}
} else {
Assert.notNull(super.writeChannelConfig.getSchema(), "Schema must be provided");
if (tableHasDefinedSchema(table)) {
Assert.isTrue(
table.getDefinition().getSchema().equals(super.writeChannelConfig.getSchema()),
"Schema should be the same"
);
}
}
return null;
});
}
protected byte[] mapItemToJson(T t) {
byte[] result = null;
try {
result = Objects.isNull(rowMapper) ? objectWriter.writeValueAsBytes(t) : rowMapper.convert(t);
}
catch (JsonProcessingException e) {
logger.error("Error during processing of the line: ", e);
}
return result;
}
/**
* 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);
}
}

39
spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/builder/BigQueryItemWriterBuilder.java → spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/writer/builder/BigQueryCsvItemWriterBuilder.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2002-2021 the original author or authors.
* Copyright 2002-2022 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.
@@ -14,62 +14,61 @@
* limitations under the License.
*/
package org.springframework.batch.extensions.bigquery.builder;
import java.util.function.Consumer;
package org.springframework.batch.extensions.bigquery.writer.builder;
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.bigquery.BigQueryItemWriter;
import org.springframework.batch.extensions.bigquery.writer.BigQueryCsvItemWriter;
import org.springframework.core.convert.converter.Converter;
import java.util.function.Consumer;
/**
* A builder for {@link BigQueryItemWriter}.
* A builder for {@link BigQueryCsvItemWriter}.
*
* @author Vova Perebykivskyi
* @since 0.1.0
* @see BigQueryItemWriter
* @see <a href="https://github.com/spring-projects/spring-batch-extensions/tree/main/spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/builder/BigQueryItemWriterBuilderTests.java">Examples</a>
* @author Volodymyr Perebykivskyi
* @since 0.2.0
* @see BigQueryCsvItemWriter
* @see <a href="https://github.com/spring-projects/spring-batch-extensions/tree/main/spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/builder/BigQueryCsvItemWriterBuilderTests.java">Examples</a>
*/
public class BigQueryItemWriterBuilder<T> {
public class BigQueryCsvItemWriterBuilder<T> {
private Converter<T, byte[]> rowMapper;
private Consumer<Job> jobConsumer;
private Consumer<Job> jobConsumer;
private DatasetInfo datasetInfo;
private WriteChannelConfiguration writeChannelConfig;
private BigQuery bigQuery;
public BigQueryItemWriterBuilder<T> rowMapper(Converter<T, byte[]> rowMapper) {
public BigQueryCsvItemWriterBuilder<T> rowMapper(Converter<T, byte[]> rowMapper) {
this.rowMapper = rowMapper;
return this;
}
public BigQueryItemWriterBuilder<T> datasetInfo(DatasetInfo datasetInfo) {
public BigQueryCsvItemWriterBuilder<T> datasetInfo(DatasetInfo datasetInfo) {
this.datasetInfo = datasetInfo;
return this;
}
public BigQueryItemWriterBuilder<T> jobConsumer(Consumer<Job> consumer) {
public BigQueryCsvItemWriterBuilder<T> jobConsumer(Consumer<Job> consumer) {
this.jobConsumer = consumer;
return this;
}
public BigQueryItemWriterBuilder<T> writeChannelConfig(WriteChannelConfiguration configuration) {
public BigQueryCsvItemWriterBuilder<T> writeChannelConfig(WriteChannelConfiguration configuration) {
this.writeChannelConfig = configuration;
return this;
}
public BigQueryItemWriterBuilder<T> bigQuery(BigQuery bigQuery) {
public BigQueryCsvItemWriterBuilder<T> bigQuery(BigQuery bigQuery) {
this.bigQuery = bigQuery;
return this;
}
public BigQueryItemWriter<T> build() {
BigQueryItemWriter<T> writer = new BigQueryItemWriter<>();
public BigQueryCsvItemWriter<T> build() {
BigQueryCsvItemWriter<T> writer = new BigQueryCsvItemWriter<>();
writer.setRowMapper(this.rowMapper);
writer.setWriteChannelConfig(this.writeChannelConfig);

82
spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/writer/builder/BigQueryJsonItemWriterBuilder.java Normal file
View File

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

43
spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/writer/package-info.java Normal file
View File

@@ -0,0 +1,43 @@
/*
* Copyright 2002-2022 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.
*
* These writers use 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 com.google.cloud.bigquery.Job}.
*
* <p>Supported formats:
* <ul>
* <li>JSON</li>
* <li>CSV</li>
* </ul>
*
* <p>For example if you generate {@link com.google.cloud.bigquery.TableDataWriteChannel} and you {@link com.google.cloud.bigquery.TableDataWriteChannel#close()} it,
* there is no guarantee that single {@link com.google.cloud.bigquery.Job} will be created.
*
* Take into account that BigQuery has rate limits, and it is very easy to exceed those in concurrent environment.
* @see <a href="https://cloud.google.com/bigquery/quotas">BigQuery Quotas &amp; Limits</a>
*
* @author Volodymyr Perebykivskyi
* @since 0.1.0
* @see <a href="https://cloud.google.com/bigquery/">Google BigQuery</a>
* @see <a href="https://github.com/googleapis/java-bigquery">BigQuery Java Client on GitHub</a>
*/
@NonNullApi
package org.springframework.batch.extensions.bigquery.writer;
import org.springframework.lang.NonNullApi;

77
spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/builder/BigQueryItemWriterBuilderTests.java → spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/writer/builder/BigQueryCsvItemWriterBuilderTests.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2002-2021 the original author or authors.
* Copyright 2002-2022 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.
@@ -14,17 +14,13 @@
* limitations under the License.
*/
package org.springframework.batch.extensions.bigquery.builder;
package org.springframework.batch.extensions.bigquery.writer.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;
@@ -32,10 +28,10 @@ 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.bigquery.writer.BigQueryCsvItemWriter;
import org.springframework.batch.extensions.bigquery.writer.builder.BigQueryCsvItemWriterBuilder;
import org.springframework.batch.extensions.bigquery.BigQueryItemWriter;
class BigQueryItemWriterBuilderTests {
class BigQueryCsvItemWriterBuilderTests {
private static final String DATASET_NAME = "my_dataset";
@@ -56,11 +52,12 @@ class BigQueryItemWriterBuilderTests {
.setFormatOptions(FormatOptions.csv())
.build();
BigQueryItemWriter<PersonDto> writer = new BigQueryItemWriterBuilder<PersonDto>()
BigQueryCsvItemWriter<PersonDto> writer = new BigQueryCsvItemWriterBuilder<PersonDto>()
.bigQuery(mockedBigQuery)
.rowMapper(dto -> convertDtoToCsvByteArray(csvMapper, dto))
.writeChannelConfig(writeConfiguration)
.datasetInfo(datasetInfo)
.jobConsumer(job -> this.logger.debug("Job with id: " + job.getJobId() + " is created"))
.build();
writer.afterPropertiesSet();
@@ -78,7 +75,7 @@ class BigQueryItemWriterBuilderTests {
.setFormatOptions(FormatOptions.csv())
.build();
BigQueryItemWriter<PersonDto> writer = new BigQueryItemWriterBuilder<PersonDto>()
BigQueryCsvItemWriter<PersonDto> writer = new BigQueryCsvItemWriterBuilder<PersonDto>()
.bigQuery(mockedBigQuery)
.writeChannelConfig(writeConfiguration)
.build();
@@ -88,62 +85,6 @@ class BigQueryItemWriterBuilderTests {
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);
}
@Test
void testCsvWriterWithJsonMapper() {
BigQuery mockedBigQuery = prepareMockedBigQuery();
WriteChannelConfiguration writeConfiguration = WriteChannelConfiguration
.newBuilder(TableId.of(DATASET_NAME, "json_table"))
.setAutodetect(true)
.setFormatOptions(FormatOptions.json())
.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);
@@ -168,7 +109,7 @@ class BigQueryItemWriterBuilderTests {
}
class PersonDto {
static class PersonDto {
private final String name;

122
spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/writer/builder/BigQueryJsonItemWriterBuilderTests.java Normal file
View File

@@ -0,0 +1,122 @@
/*
* Copyright 2002-2022 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.bigquery.writer.builder;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.cloud.bigquery.BigQuery;
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.bigquery.writer.BigQueryJsonItemWriter;
import org.springframework.batch.extensions.bigquery.writer.builder.BigQueryJsonItemWriterBuilder;
class BigQueryJsonItemWriterBuilderTests {
private static final String DATASET_NAME = "my_dataset";
private final Log logger = LogFactory.getLog(getClass());
/**
* Example how JSON writer is expected to be built without {@link org.springframework.context.annotation.Bean} annotation.
*/
@Test
void testJsonWriterWithRowMapper() {
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();
BigQueryJsonItemWriter<PersonDto> writer = new BigQueryJsonItemWriterBuilder<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);
}
@Test
void testCsvWriterWithJsonMapper() {
BigQuery mockedBigQuery = prepareMockedBigQuery();
WriteChannelConfiguration writeConfiguration = WriteChannelConfiguration
.newBuilder(TableId.of(DATASET_NAME, "json_table"))
.setAutodetect(true)
.setFormatOptions(FormatOptions.json())
.build();
BigQueryJsonItemWriter<PersonDto> writer = new BigQueryJsonItemWriterBuilder<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 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 final String name;
public PersonDto(String name) {
this.name = name;
}
}
}