From 4114825743f5464d19d80b494a0378599130ef56 Mon Sep 17 00:00:00 2001 From: Volodymyr Date: Thu, 12 Jun 2025 07:29:22 +0300 Subject: [PATCH] [bq] Provide a default mapper for BigQueryQueryItemReader Signed-off-by: Volodymyr Perebykivskyi --- .../BigQueryQueryItemReaderBuilder.java | 33 ++++++++-- .../bigquery/reader/builder/RecordMapper.java | 64 +++++++++++++++++++ .../writer/BigQueryBaseItemWriter.java | 49 ++++++-------- .../writer/BigQueryItemWriterException.java | 2 + .../BigQueryItemReaderBuilderTest.java | 57 +++++++++++++++-- .../unit/reader/builder/RecordMapperTest.java | 64 +++++++++++++++++++ 6 files changed, 231 insertions(+), 38 deletions(-) create mode 100644 spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/reader/builder/RecordMapper.java create mode 100644 spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/unit/reader/builder/RecordMapperTest.java diff --git a/spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/reader/builder/BigQueryQueryItemReaderBuilder.java b/spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/reader/builder/BigQueryQueryItemReaderBuilder.java index 63f417d..8f24b12 100644 --- a/spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/reader/builder/BigQueryQueryItemReaderBuilder.java +++ b/spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/reader/builder/BigQueryQueryItemReaderBuilder.java @@ -39,6 +39,7 @@ public class BigQueryQueryItemReaderBuilder { private String query; private Converter rowMapper; private QueryJobConfiguration jobConfiguration; + private Class targetType; /** * BigQuery service, responsible for API calls. @@ -47,7 +48,7 @@ public class BigQueryQueryItemReaderBuilder { * @return {@link BigQueryQueryItemReaderBuilder} * @see BigQueryQueryItemReader#setBigQuery(BigQuery) */ - public BigQueryQueryItemReaderBuilder bigQuery(BigQuery bigQuery) { + public BigQueryQueryItemReaderBuilder bigQuery(final BigQuery bigQuery) { this.bigQuery = bigQuery; return this; } @@ -62,7 +63,7 @@ public class BigQueryQueryItemReaderBuilder { * @return {@link BigQueryQueryItemReaderBuilder} * @see BigQueryQueryItemReader#setJobConfiguration(QueryJobConfiguration) */ - public BigQueryQueryItemReaderBuilder query(String query) { + public BigQueryQueryItemReaderBuilder query(final String query) { this.query = query; return this; } @@ -74,7 +75,7 @@ public class BigQueryQueryItemReaderBuilder { * @return {@link BigQueryQueryItemReaderBuilder} * @see BigQueryQueryItemReader#setRowMapper(Converter) */ - public BigQueryQueryItemReaderBuilder rowMapper(Converter rowMapper) { + public BigQueryQueryItemReaderBuilder rowMapper(final Converter rowMapper) { this.rowMapper = rowMapper; return this; } @@ -86,21 +87,41 @@ public class BigQueryQueryItemReaderBuilder { * @return {@link BigQueryQueryItemReaderBuilder} * @see BigQueryQueryItemReader#setJobConfiguration(QueryJobConfiguration) */ - public BigQueryQueryItemReaderBuilder jobConfiguration(QueryJobConfiguration jobConfiguration) { + public BigQueryQueryItemReaderBuilder 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 targetType(final Class targetType) { + this.targetType = targetType; + return this; + } + /** * Please remember about {@link BigQueryQueryItemReader#afterPropertiesSet()}. * * @return {@link BigQueryQueryItemReader} */ public BigQueryQueryItemReader build() { - BigQueryQueryItemReader reader = new BigQueryQueryItemReader<>(); + final BigQueryQueryItemReader 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().generateMapper(this.targetType)); + } else { + reader.setRowMapper(this.rowMapper); + } if (this.jobConfiguration == null) { Assert.isTrue(StringUtils.hasText(this.query), "No query provided"); diff --git a/spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/reader/builder/RecordMapper.java b/spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/reader/builder/RecordMapper.java new file mode 100644 index 0000000..eaabfeb --- /dev/null +++ b/spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/reader/builder/RecordMapper.java @@ -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 Java record type + * @author Volodymyr Perebykivskyi + * @since 0.2.0 + */ +public final class RecordMapper { + + 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 generateMapper(Class targetType) { + Constructor 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); + }; + } +} diff --git a/spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/writer/BigQueryBaseItemWriter.java b/spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/writer/BigQueryBaseItemWriter.java index aaea894..91dc7a2 100644 --- a/spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/writer/BigQueryBaseItemWriter.java +++ b/spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/writer/BigQueryBaseItemWriter.java @@ -82,7 +82,7 @@ public abstract class BigQueryBaseItemWriter implements ItemWriter, 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 implements ItemWriter, Initia * * @param consumer your consumer */ - public void setJobConsumer(Consumer consumer) { + public void setJobConsumer(final Consumer consumer) { this.jobConsumer = consumer; } @@ -100,7 +100,7 @@ public abstract class BigQueryBaseItemWriter implements ItemWriter, 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 implements ItemWriter, Initia * * @param bigQuery BigQuery service */ - public void setBigQuery(BigQuery bigQuery) { + public void setBigQuery(final BigQuery bigQuery) { this.bigQuery = bigQuery; } @Override - public void write(Chunk chunk) throws Exception { + public void write(final Chunk chunk) throws Exception { if (!chunk.isEmpty()) { - List items = chunk.getItems(); + final List 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 items) throws IOException { - ByteBuffer byteBuffer; - try (ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) { + private ByteBuffer mapDataToBigQueryFormat(final List items) throws IOException { + final ByteBuffer byteBuffer; + try (final ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) { - List data = convertObjectsToByteArrays(items); + final List data = convertObjectsToByteArrays(items); for (byte[] byteArray : data) { outputStream.write(byteArray); @@ -147,14 +147,14 @@ public abstract class BigQueryBaseItemWriter implements ItemWriter, 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 implements ItemWriter, 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 implements ItemWriter, 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) diff --git a/spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/writer/BigQueryItemWriterException.java b/spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/writer/BigQueryItemWriterException.java index 47a6fd6..bf17348 100644 --- a/spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/writer/BigQueryItemWriterException.java +++ b/spring-batch-bigquery/src/main/java/org/springframework/batch/extensions/bigquery/writer/BigQueryItemWriterException.java @@ -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 { diff --git a/spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/unit/reader/builder/BigQueryItemReaderBuilderTest.java b/spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/unit/reader/builder/BigQueryItemReaderBuilderTest.java index ffb0320..44cb5a7 100644 --- a/spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/unit/reader/builder/BigQueryItemReaderBuilderTest.java +++ b/spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/unit/reader/builder/BigQueryItemReaderBuilderTest.java @@ -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 reader = new BigQueryQueryItemReaderBuilder() + .bigQuery(mockedBigQuery) + .jobConfiguration(expectedJobConfiguration) + .targetType(PersonDto.class) + .build(); + + Assertions.assertNotNull(reader); + + BigQuery actualBigQuery = (BigQuery) handle + .findVarHandle(BigQueryQueryItemReader.class, "bigQuery", BigQuery.class) + .get(reader); + + Converter actualRowMapper = (Converter) 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 brokenBuilders() { + final class HumanDto {} + return Stream.of( + Arguments.of("No target type provided", new BigQueryQueryItemReaderBuilder()), + Arguments.of("Only Java record supported", new BigQueryQueryItemReaderBuilder().targetType(HumanDto.class)), + Arguments.of("No query provided", new BigQueryQueryItemReaderBuilder().rowMapper(source -> null)) + ); } } diff --git a/spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/unit/reader/builder/RecordMapperTest.java b/spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/unit/reader/builder/RecordMapperTest.java new file mode 100644 index 0000000..7ecb018 --- /dev/null +++ b/spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/unit/reader/builder/RecordMapperTest.java @@ -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 mapper = new RecordMapper<>(); + List 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 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().generateMapper(TestRecord.class)); + Assertions.assertEquals("Record without fields is redundant", ex.getMessage()); + } + +}