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

[bq] Provide a default mapper for BigQueryQueryItemReader

Browse Source 
Signed-off-by: Volodymyr Perebykivskyi <vova235@gmail.com>
This commit is contained in:
Volodymyr
2025-06-12 07:29:22 +03:00
committed by GitHub
parent 22a62f8e31
commit 4114825743
6 changed files with 231 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/reader/builder/BigQueryQueryItemReaderBuilder.java
View File

@@ -39,6 +39,7 @@ public class BigQueryQueryItemReaderBuilder<T> {
private String query;
private Converter<FieldValueList, T> rowMapper;
private QueryJobConfiguration jobConfiguration;
private Class<T> targetType;
/**
* BigQuery service, responsible for API calls.
@@ -47,7 +48,7 @@ public class BigQueryQueryItemReaderBuilder<T> {
* @return {@link BigQueryQueryItemReaderBuilder}
* @see BigQueryQueryItemReader#setBigQuery(BigQuery)
*/
public BigQueryQueryItemReaderBuilder<T> bigQuery(BigQuery bigQuery) {
public BigQueryQueryItemReaderBuilder<T> bigQuery(final BigQuery bigQuery) {
this.bigQuery = bigQuery;
return this;
}
@@ -62,7 +63,7 @@ public class BigQueryQueryItemReaderBuilder<T> {
* @return {@link BigQueryQueryItemReaderBuilder}
* @see BigQueryQueryItemReader#setJobConfiguration(QueryJobConfiguration)
*/
public BigQueryQueryItemReaderBuilder<T> query(String query) {
public BigQueryQueryItemReaderBuilder<T> query(final String query) {
this.query = query;
return this;
}
@@ -74,7 +75,7 @@ public class BigQueryQueryItemReaderBuilder<T> {
* @return {@link BigQueryQueryItemReaderBuilder}
* @see BigQueryQueryItemReader#setRowMapper(Converter)
*/
public BigQueryQueryItemReaderBuilder<T> rowMapper(Converter<FieldValueList, T> rowMapper) {
public BigQueryQueryItemReaderBuilder<T> rowMapper(final Converter<FieldValueList, T> rowMapper) {
this.rowMapper = rowMapper;
return this;
}
@@ -86,21 +87,41 @@ public class BigQueryQueryItemReaderBuilder<T> {
* @return {@link BigQueryQueryItemReaderBuilder}
* @see BigQueryQueryItemReader#setJobConfiguration(QueryJobConfiguration)
*/
public BigQueryQueryItemReaderBuilder<T> jobConfiguration(QueryJobConfiguration jobConfiguration) {
public BigQueryQueryItemReaderBuilder<T> jobConfiguration(final QueryJobConfiguration jobConfiguration) {
this.jobConfiguration = jobConfiguration;
return this;
}
/**
* Specifies a target type which will be used as a result.
* Only needed when {@link BigQueryQueryItemReaderBuilder#rowMapper} is not provided.
* Take into account that only {@link Class#isRecord()} supported.
*
* @param targetType a {@link Class} that represent desired type
* @return {@link BigQueryQueryItemReaderBuilder}
*/
public BigQueryQueryItemReaderBuilder<T> targetType(final Class<T> targetType) {
this.targetType = targetType;
return this;
}
/**
* Please remember about {@link BigQueryQueryItemReader#afterPropertiesSet()}.
*
* @return {@link BigQueryQueryItemReader}
*/
public BigQueryQueryItemReader<T> build() {
BigQueryQueryItemReader<T> reader = new BigQueryQueryItemReader<>();
final BigQueryQueryItemReader<T> reader = new BigQueryQueryItemReader<>();
reader.setBigQuery(this.bigQuery);
reader.setRowMapper(this.rowMapper);
if (this.rowMapper == null) {
Assert.notNull(this.targetType, "No target type provided");
Assert.isTrue(this.targetType.isRecord(), "Only Java record supported");
reader.setRowMapper(new RecordMapper<T>().generateMapper(this.targetType));
} else {
reader.setRowMapper(this.rowMapper);
}
if (this.jobConfiguration == null) {
Assert.isTrue(StringUtils.hasText(this.query), "No query provided");

64
spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/reader/builder/RecordMapper.java Normal file
View File

@@ -0,0 +1,64 @@
/*
* Copyright 2002-2025 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.reader.builder;
import com.google.cloud.bigquery.FieldValueList;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.SimpleTypeConverter;
import org.springframework.core.convert.converter.Converter;
import org.springframework.util.Assert;
import java.lang.reflect.Constructor;
/**
* A helper class which tries to convert BigQuery response to a Java record.
*
* @param <T> Java record type
* @author Volodymyr Perebykivskyi
* @since 0.2.0
*/
public final class RecordMapper<T> {
private final SimpleTypeConverter simpleConverter = new SimpleTypeConverter();
/**
* Generates a conversion from BigQuery response to a Java record.
*
* @param targetType a Java record {@link Class}
* @return {@link Converter}
* @see org.springframework.batch.item.file.mapping.RecordFieldSetMapper
*/
public Converter<FieldValueList, T> generateMapper(Class<T> targetType) {
Constructor<T> constructor = BeanUtils.getResolvableConstructor(targetType);
Assert.isTrue(constructor.getParameterCount() > 0, "Record without fields is redundant");
String[] parameterNames = BeanUtils.getParameterNames(constructor);
Class<?>[] parameterTypes = constructor.getParameterTypes();
Object[] args = new Object[parameterNames.length];
return source -> {
if (args[0] == null) {
for (int i = 0; i < args.length; i++) {
args[i] = simpleConverter.convertIfNecessary(source.get(parameterNames[i]).getValue(), parameterTypes[i]);
}
}
return BeanUtils.instantiateClass(constructor, args);
};
}
}

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

@@ -82,7 +82,7 @@ public abstract class BigQueryBaseItemWriter<T> implements ItemWriter<T>, Initia
*
* @param datasetInfo BigQuery dataset info
*/
public void setDatasetInfo(DatasetInfo datasetInfo) {
public void setDatasetInfo(final DatasetInfo datasetInfo) {
this.datasetInfo = datasetInfo;
}
@@ -91,7 +91,7 @@ public abstract class BigQueryBaseItemWriter<T> implements ItemWriter<T>, Initia
*
* @param consumer your consumer
*/
public void setJobConsumer(Consumer<Job> consumer) {
public void setJobConsumer(final Consumer<Job> consumer) {
this.jobConsumer = consumer;
}
@@ -100,7 +100,7 @@ public abstract class BigQueryBaseItemWriter<T> implements ItemWriter<T>, Initia
*
* @param writeChannelConfig BigQuery channel configuration
*/
public void setWriteChannelConfig(WriteChannelConfiguration writeChannelConfig) {
public void setWriteChannelConfig(final WriteChannelConfiguration writeChannelConfig) {
this.writeChannelConfig = writeChannelConfig;
}
@@ -109,30 +109,30 @@ public abstract class BigQueryBaseItemWriter<T> implements ItemWriter<T>, Initia
*
* @param bigQuery BigQuery service
*/
public void setBigQuery(BigQuery bigQuery) {
public void setBigQuery(final BigQuery bigQuery) {
this.bigQuery = bigQuery;
}
@Override
public void write(Chunk<? extends T> chunk) throws Exception {
public void write(final Chunk<? extends T> chunk) throws Exception {
if (!chunk.isEmpty()) {
List<? extends T> items = chunk.getItems();
final List<? extends T> items = chunk.getItems();
doInitializeProperties(items);
if (this.logger.isDebugEnabled()) {
this.logger.debug(String.format("Mapping %d elements", items.size()));
}
ByteBuffer byteBuffer = mapDataToBigQueryFormat(items);
final ByteBuffer byteBuffer = mapDataToBigQueryFormat(items);
doWriteDataToBigQuery(byteBuffer);
}
}
private ByteBuffer mapDataToBigQueryFormat(List<? extends T> items) throws IOException {
ByteBuffer byteBuffer;
try (ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {
private ByteBuffer mapDataToBigQueryFormat(final List<? extends T> items) throws IOException {
final ByteBuffer byteBuffer;
try (final ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {
List<byte[]> data = convertObjectsToByteArrays(items);
final List<byte[]> data = convertObjectsToByteArrays(items);
for (byte[] byteArray : data) {
outputStream.write(byteArray);
@@ -147,14 +147,14 @@ public abstract class BigQueryBaseItemWriter<T> implements ItemWriter<T>, Initia
return byteBuffer;
}
private void doWriteDataToBigQuery(ByteBuffer byteBuffer) throws IOException {
private void doWriteDataToBigQuery(final ByteBuffer byteBuffer) {
if (this.logger.isDebugEnabled()) {
this.logger.debug("Writing data to BigQuery");
}
TableDataWriteChannel writeChannel = null;
try (TableDataWriteChannel writer = getWriteChannel()) {
try (final TableDataWriteChannel writer = getWriteChannel()) {
/* TableDataWriteChannel is not thread safe */
writer.write(byteBuffer);
writeChannel = writer;
@@ -209,29 +209,22 @@ public abstract class BigQueryBaseItemWriter<T> implements ItemWriter<T>, Initia
performFormatSpecificChecks();
String dataset = this.writeChannelConfig.getDestinationTable().getDataset();
final String dataset = this.writeChannelConfig.getDestinationTable().getDataset();
if (this.datasetInfo == null) {
this.datasetInfo = DatasetInfo.newBuilder(dataset).build();
} else {
Assert.isTrue(Objects.equals(this.datasetInfo.getDatasetId().getDataset(), dataset), "Dataset should be configured properly");
}
Assert.isTrue(
Objects.equals(this.datasetInfo.getDatasetId().getDataset(), dataset),
"Dataset should be configured properly"
);
createDataset();
}
private void createDataset() {
TableId tableId = this.writeChannelConfig.getDestinationTable();
String datasetToCheck = tableId.getDataset();
final TableId tableId = this.writeChannelConfig.getDestinationTable();
final String datasetToCheck = tableId.getDataset();
if (datasetToCheck != null) {
Dataset foundDataset = this.bigQuery.getDataset(datasetToCheck);
if (foundDataset == null && this.datasetInfo != null) {
this.bigQuery.create(this.datasetInfo);
}
if (datasetToCheck != null && this.bigQuery.getDataset(datasetToCheck) == null && this.datasetInfo != null) {
this.bigQuery.create(this.datasetInfo);
}
}
@@ -270,7 +263,7 @@ public abstract class BigQueryBaseItemWriter<T> implements ItemWriter<T>, Initia
* @param table BigQuery table
* @return {@code true} if BigQuery {@link Table} has schema already described
*/
protected boolean tableHasDefinedSchema(Table table) {
protected boolean tableHasDefinedSchema(final Table table) {
return Optional
.ofNullable(table)
.map(Table::getDefinition)

2
spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/writer/BigQueryItemWriterException.java
View File

@@ -22,6 +22,8 @@ import org.springframework.batch.item.ItemWriterException;
/**
* Unchecked {@link Exception} indicating that an error has occurred on during {@link ItemWriter#write(Chunk)}.
* @author Volodymyr Perebykivskyi
* @since 0.2.0
*/
public class BigQueryItemWriterException extends ItemWriterException {

57
spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/unit/reader/builder/BigQueryItemReaderBuilderTest.java
View File

@@ -22,6 +22,9 @@ import com.google.cloud.bigquery.QueryJobConfiguration;
import com.google.cloud.bigquery.TableId;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
import org.springframework.batch.extensions.bigquery.common.PersonDto;
import org.springframework.batch.extensions.bigquery.common.TestConstants;
import org.springframework.batch.extensions.bigquery.reader.BigQueryQueryItemReader;
@@ -30,6 +33,7 @@ import org.springframework.batch.extensions.bigquery.unit.base.AbstractBigQueryT
import org.springframework.core.convert.converter.Converter;
import java.lang.invoke.MethodHandles;
import java.util.stream.Stream;
class BigQueryItemReaderBuilderTest extends AbstractBigQueryTest {
@@ -65,7 +69,41 @@ class BigQueryItemReaderBuilderTest extends AbstractBigQueryTest {
}
@Test
void testBuild_WithJobConfiguration() throws IllegalAccessException, NoSuchFieldException {
void testBuild_WithoutRowMapper() throws IllegalAccessException, NoSuchFieldException {
BigQuery mockedBigQuery = prepareMockedBigQuery();
MethodHandles.Lookup handle = MethodHandles.privateLookupIn(BigQueryQueryItemReader.class, MethodHandles.lookup());
QueryJobConfiguration expectedJobConfiguration = QueryJobConfiguration
.newBuilder("SELECT p.name, p.age FROM spring_batch_extensions.persons p LIMIT 1")
.build();
BigQueryQueryItemReader<PersonDto> reader = new BigQueryQueryItemReaderBuilder<PersonDto>()
.bigQuery(mockedBigQuery)
.jobConfiguration(expectedJobConfiguration)
.targetType(PersonDto.class)
.build();
Assertions.assertNotNull(reader);
BigQuery actualBigQuery = (BigQuery) handle
.findVarHandle(BigQueryQueryItemReader.class, "bigQuery", BigQuery.class)
.get(reader);
Converter<FieldValueList, PersonDto> actualRowMapper = (Converter<FieldValueList, PersonDto>) handle
.findVarHandle(BigQueryQueryItemReader.class, "rowMapper", Converter.class)
.get(reader);
QueryJobConfiguration actualJobConfiguration = (QueryJobConfiguration) handle
.findVarHandle(BigQueryQueryItemReader.class, "jobConfiguration", QueryJobConfiguration.class)
.get(reader);
Assertions.assertEquals(mockedBigQuery, actualBigQuery);
Assertions.assertNotNull(actualRowMapper);
Assertions.assertEquals(expectedJobConfiguration, actualJobConfiguration);
}
@Test
void testBuild() throws IllegalAccessException, NoSuchFieldException {
BigQuery mockedBigQuery = prepareMockedBigQuery();
MethodHandles.Lookup handle = MethodHandles.privateLookupIn(BigQueryQueryItemReader.class, MethodHandles.lookup());
@@ -99,8 +137,19 @@ class BigQueryItemReaderBuilderTest extends AbstractBigQueryTest {
Assertions.assertEquals(jobConfiguration, actualJobConfiguration);
}
@Test
void testBuild_NoQueryProvided() {
Assertions.assertThrows(IllegalArgumentException.class, new BigQueryQueryItemReaderBuilder<>()::build);
@ParameterizedTest
@MethodSource("brokenBuilders")
void testBuild_Exception(String expectedMessage, BigQueryQueryItemReaderBuilder<?> builder) {
IllegalArgumentException ex = Assertions.assertThrows(IllegalArgumentException.class, builder::build);
Assertions.assertEquals(expectedMessage, ex.getMessage());
}
private static Stream<Arguments> brokenBuilders() {
final class HumanDto {}
return Stream.of(
Arguments.of("No target type provided", new BigQueryQueryItemReaderBuilder<PersonDto>()),
Arguments.of("Only Java record supported", new BigQueryQueryItemReaderBuilder<HumanDto>().targetType(HumanDto.class)),
Arguments.of("No query provided", new BigQueryQueryItemReaderBuilder<PersonDto>().rowMapper(source -> null))
);
}
}

64
spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/unit/reader/builder/RecordMapperTest.java Normal file
View File

@@ -0,0 +1,64 @@
/*
* Copyright 2002-2025 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.unit.reader.builder;
import com.google.cloud.bigquery.Field;
import com.google.cloud.bigquery.FieldValue;
import com.google.cloud.bigquery.FieldValueList;
import com.google.cloud.bigquery.StandardSQLTypeName;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
import org.springframework.batch.extensions.bigquery.common.PersonDto;
import org.springframework.batch.extensions.bigquery.common.TestConstants;
import org.springframework.batch.extensions.bigquery.reader.builder.RecordMapper;
import org.springframework.core.convert.converter.Converter;
import java.util.List;
class RecordMapperTest {
@Test
void testGenerateMapper() {
RecordMapper<PersonDto> mapper = new RecordMapper<>();
List<PersonDto> expected = TestConstants.CHUNK.getItems();
Field name = Field.of(TestConstants.NAME, StandardSQLTypeName.STRING);
Field age = Field.of(TestConstants.AGE, StandardSQLTypeName.INT64);
PersonDto person1 = expected.get(0);
FieldValue value1 = FieldValue.of(FieldValue.Attribute.PRIMITIVE, person1.name());
FieldValue value2 = FieldValue.of(FieldValue.Attribute.PRIMITIVE, person1.age());
FieldValueList row = FieldValueList.of(List.of(value1, value2), name, age);
Converter<FieldValueList, PersonDto> converter = mapper.generateMapper(PersonDto.class);
Assertions.assertNotNull(converter);
PersonDto actual = converter.convert(row);
Assertions.assertEquals(expected.get(0).name(), actual.name());
Assertions.assertEquals(expected.get(0).age(), actual.age());
}
@Test
void testGenerateMapper_EmptyRecord() {
record TestRecord(){}
IllegalArgumentException ex = Assertions.assertThrows(IllegalArgumentException.class, () -> new RecordMapper<TestRecord>().generateMapper(TestRecord.class));
Assertions.assertEquals("Record without fields is redundant", ex.getMessage());
}
}