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

[bq] 0.2 refactoring

Browse Source
This commit is contained in:
Volodymyr
2022-12-22 14:44:26 +02:00
committed by GitHub
parent 760cde444a
commit 038f973015
9 changed files with 187 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files

91
spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/common/BigQueryDataLoader.java Normal file
View File

@@ -0,0 +1,91 @@
/*
* 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.common;
import com.google.cloud.bigquery.BigQuery;
import com.google.cloud.bigquery.FormatOptions;
import com.google.cloud.bigquery.Job;
import com.google.cloud.bigquery.TableId;
import com.google.cloud.bigquery.WriteChannelConfiguration;
import org.springframework.batch.extensions.bigquery.writer.BigQueryCsvItemWriter;
import org.springframework.batch.extensions.bigquery.writer.BigQueryJsonItemWriter;
import org.springframework.batch.extensions.bigquery.writer.builder.BigQueryCsvItemWriterBuilder;
import org.springframework.batch.extensions.bigquery.writer.builder.BigQueryJsonItemWriterBuilder;
import org.springframework.batch.item.Chunk;
import java.util.concurrent.atomic.AtomicReference;
public class BigQueryDataLoader {
public static final Chunk<PersonDto> CHUNK = Chunk.of(
new PersonDto("Volodymyr", 27), new PersonDto("Oleksandra", 26)
);
private final BigQuery bigQuery;
public BigQueryDataLoader(BigQuery bigQuery) {
this.bigQuery = bigQuery;
}
public void loadCsvSample() throws Exception {
loadCsvSample(TestConstants.PERSONS_TABLE);
}
public void loadCsvSample(String tableName) throws Exception {
AtomicReference<Job> job = new AtomicReference<>();
WriteChannelConfiguration channelConfiguration = WriteChannelConfiguration
.newBuilder(TableId.of(TestConstants.DATASET, tableName))
.setSchema(PersonDto.getBigQuerySchema())
.setAutodetect(false)
.setFormatOptions(FormatOptions.csv())
.build();
BigQueryCsvItemWriter<PersonDto> writer = new BigQueryCsvItemWriterBuilder<PersonDto>()
.bigQuery(bigQuery)
.writeChannelConfig(channelConfiguration)
.jobConsumer(job::set)
.build();
writer.afterPropertiesSet();
writer.write(CHUNK);
job.get().waitFor();
}
public void loadJsonSample(String tableName) throws Exception {
AtomicReference<Job> job = new AtomicReference<>();
WriteChannelConfiguration channelConfiguration = WriteChannelConfiguration
.newBuilder(TableId.of(TestConstants.DATASET, tableName))
.setSchema(PersonDto.getBigQuerySchema())
.setAutodetect(false)
.setFormatOptions(FormatOptions.json())
.build();
BigQueryJsonItemWriter<PersonDto> writer = new BigQueryJsonItemWriterBuilder<PersonDto>()
.bigQuery(bigQuery)
.writeChannelConfig(channelConfiguration)
.jobConsumer(job::set)
.build();
writer.afterPropertiesSet();
writer.write(CHUNK);
job.get().waitFor();
}
}

33
spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/common/PersonDto.java Normal file
View File

@@ -0,0 +1,33 @@
/*
* 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.common;
import com.fasterxml.jackson.annotation.JsonPropertyOrder;
import com.google.cloud.bigquery.Field;
import com.google.cloud.bigquery.Schema;
import com.google.cloud.bigquery.StandardSQLTypeName;
@JsonPropertyOrder(value = {"name", "age"})
public record PersonDto(String name, Integer age) {
public static Schema getBigQuerySchema() {
Field nameField = Field.newBuilder("name", StandardSQLTypeName.STRING).build();
Field ageField = Field.newBuilder("age", StandardSQLTypeName.INT64).build();
return Schema.of(nameField, ageField);
}
}

26
spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/common/TestConstants.java Normal file
View File

@@ -0,0 +1,26 @@
/*
* 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.common;
public class TestConstants {
private TestConstants() {}
public static final String DATASET = "spring_batch_extensions";
public static final String PERSONS_TABLE = "persons";
}

30
spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/integration/writer/BigQueryCsvItemWriterTest.java
View File

@@ -18,8 +18,6 @@ package org.springframework.batch.extensions.bigquery.integration.writer;
import com.google.cloud.bigquery.BigQuery;
import com.google.cloud.bigquery.Dataset;
import com.google.cloud.bigquery.FormatOptions;
import com.google.cloud.bigquery.JobId;
import com.google.cloud.bigquery.Table;
import com.google.cloud.bigquery.TableId;
import com.google.cloud.bigquery.TableResult;
@@ -28,35 +26,23 @@ import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.TestInfo;
import org.springframework.batch.extensions.bigquery.common.BigQueryDataLoader;
import org.springframework.batch.extensions.bigquery.common.PersonDto;
import org.springframework.batch.extensions.bigquery.common.TestConstants;
import org.springframework.batch.extensions.bigquery.integration.writer.base.BaseBigQueryItemWriterTest;
import org.springframework.batch.extensions.bigquery.writer.BigQueryCsvItemWriter;
import org.springframework.batch.extensions.bigquery.writer.builder.BigQueryCsvItemWriterBuilder;
import org.springframework.batch.item.Chunk;
import java.util.concurrent.atomic.AtomicReference;
@Tag("csv")
public class BigQueryCsvItemWriterTest extends BaseBigQueryItemWriterTest {
@Test
void test1(TestInfo testInfo) throws Exception {
AtomicReference<JobId> jobId = new AtomicReference<>();
String tableName = getTableName(testInfo);
new BigQueryDataLoader(bigQuery).loadCsvSample(tableName);
Chunk<PersonDto> chunk = BigQueryDataLoader.CHUNK;
BigQueryCsvItemWriter writer = new BigQueryCsvItemWriterBuilder<PersonDto>()
.bigQuery(bigQuery)
.writeChannelConfig(generateConfiguration(testInfo, FormatOptions.csv()))
.jobConsumer(j -> jobId.set(j.getJobId()))
.build();
writer.afterPropertiesSet();
Chunk<PersonDto> chunk = Chunk.of(new PersonDto("Volodymyr", 27), new PersonDto("Oleksandra", 26));
writer.write(chunk);
waitForJobToFinish(jobId.get());
Dataset dataset = bigQuery.getDataset(DATASET);
Table table = bigQuery.getTable(TableId.of(DATASET, getTableName(testInfo)));
Dataset dataset = bigQuery.getDataset(TestConstants.DATASET);
Table table = bigQuery.getTable(TableId.of(TestConstants.DATASET, tableName));
TableId tableId = table.getTableId();
TableResult tableResult = bigQuery.listTableData(tableId, BigQuery.TableDataListOption.pageSize(2L));

30
spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/integration/writer/BigQueryJsonItemWriterTest.java
View File

@@ -18,8 +18,6 @@ package org.springframework.batch.extensions.bigquery.integration.writer;
import com.google.cloud.bigquery.BigQuery;
import com.google.cloud.bigquery.Dataset;
import com.google.cloud.bigquery.FormatOptions;
import com.google.cloud.bigquery.JobId;
import com.google.cloud.bigquery.Table;
import com.google.cloud.bigquery.TableId;
import com.google.cloud.bigquery.TableResult;
@@ -28,35 +26,23 @@ import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.TestInfo;
import org.springframework.batch.extensions.bigquery.common.BigQueryDataLoader;
import org.springframework.batch.extensions.bigquery.common.PersonDto;
import org.springframework.batch.extensions.bigquery.common.TestConstants;
import org.springframework.batch.extensions.bigquery.integration.writer.base.BaseBigQueryItemWriterTest;
import org.springframework.batch.extensions.bigquery.writer.BigQueryJsonItemWriter;
import org.springframework.batch.extensions.bigquery.writer.builder.BigQueryJsonItemWriterBuilder;
import org.springframework.batch.item.Chunk;
import java.util.concurrent.atomic.AtomicReference;
@Tag("json")
public class BigQueryJsonItemWriterTest extends BaseBigQueryItemWriterTest {
@Test
void test1(TestInfo testInfo) throws Exception {
AtomicReference<JobId> jobId = new AtomicReference<>();
String tableName = getTableName(testInfo);
new BigQueryDataLoader(bigQuery).loadJsonSample(tableName);
Chunk<PersonDto> chunk = BigQueryDataLoader.CHUNK;
BigQueryJsonItemWriter<PersonDto> writer = new BigQueryJsonItemWriterBuilder<PersonDto>()
.bigQuery(bigQuery)
.writeChannelConfig(generateConfiguration(testInfo, FormatOptions.json()))
.jobConsumer(j -> jobId.set(j.getJobId()))
.build();
writer.afterPropertiesSet();
Chunk<PersonDto> chunk = Chunk.of(new PersonDto("Viktor", 57), new PersonDto("Nina", 57));
writer.write(chunk);
waitForJobToFinish(jobId.get());
Dataset dataset = bigQuery.getDataset(DATASET);
Table table = bigQuery.getTable(TableId.of(DATASET, getTableName(testInfo)));
Dataset dataset = bigQuery.getDataset(TestConstants.DATASET);
Table table = bigQuery.getTable(TableId.of(TestConstants.DATASET, tableName));
TableId tableId = table.getTableId();
TableResult tableResult = bigQuery.listTableData(tableId, BigQuery.TableDataListOption.pageSize(2L));

35
spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/integration/writer/base/BaseBigQueryItemWriterTest.java
View File

@@ -16,16 +16,13 @@
package org.springframework.batch.extensions.bigquery.integration.writer.base;
import com.fasterxml.jackson.annotation.JsonPropertyOrder;
import com.google.cloud.RetryOption;
import com.google.cloud.bigquery.BigQuery;
import com.google.cloud.bigquery.BigQueryOptions;
import com.google.cloud.bigquery.DatasetInfo;
import com.google.cloud.bigquery.Field;
import com.google.cloud.bigquery.FormatOptions;
import com.google.cloud.bigquery.JobId;
import com.google.cloud.bigquery.JobStatus;
import com.google.cloud.bigquery.Schema;
import com.google.cloud.bigquery.StandardSQLTypeName;
import com.google.cloud.bigquery.StandardTableDefinition;
import com.google.cloud.bigquery.TableDefinition;
import com.google.cloud.bigquery.TableId;
@@ -35,33 +32,33 @@ import org.apache.commons.lang3.BooleanUtils;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.TestInfo;
import org.springframework.batch.extensions.bigquery.common.PersonDto;
import org.springframework.batch.extensions.bigquery.common.TestConstants;
import java.lang.reflect.Method;
import java.util.Objects;
public abstract class BaseBigQueryItemWriterTest {
protected static final String DATASET = "spring_extensions";
private static final String TABLE_PATTERN = "%s_%s";
protected final BigQuery bigQuery = BigQueryOptions.getDefaultInstance().getService();
private static final String TABLE_PATTERN = "%s_%s";
@BeforeEach
void prepareTest(TestInfo testInfo) {
if (Objects.isNull(bigQuery.getDataset(DATASET))) {
bigQuery.create(DatasetInfo.of(DATASET));
if (Objects.isNull(bigQuery.getDataset(TestConstants.DATASET))) {
bigQuery.create(DatasetInfo.of(TestConstants.DATASET));
}
if (Objects.isNull(bigQuery.getTable(DATASET, getTableName(testInfo)))) {
if (Objects.isNull(bigQuery.getTable(TestConstants.DATASET, getTableName(testInfo)))) {
TableDefinition tableDefinition = StandardTableDefinition.of(PersonDto.getBigQuerySchema());
bigQuery.create(TableInfo.of(TableId.of(DATASET, getTableName(testInfo)), tableDefinition));
bigQuery.create(TableInfo.of(TableId.of(TestConstants.DATASET, getTableName(testInfo)), tableDefinition));
}
}
@AfterEach
void cleanupTest(TestInfo testInfo) {
bigQuery.delete(TableId.of(DATASET, getTableName(testInfo)));
bigQuery.delete(TableId.of(TestConstants.DATASET, getTableName(testInfo)));
}
protected String getTableName(TestInfo testInfo) {
@@ -74,13 +71,14 @@ public abstract class BaseBigQueryItemWriterTest {
protected WriteChannelConfiguration generateConfiguration(TestInfo testInfo, FormatOptions formatOptions) {
return WriteChannelConfiguration
.newBuilder(TableId.of(DATASET, getTableName(testInfo)))
.newBuilder(TableId.of(TestConstants.DATASET, getTableName(testInfo)))
.setSchema(PersonDto.getBigQuerySchema())
.setAutodetect(false)
.setFormatOptions(formatOptions)
.build();
}
/** TODO check {@link com.google.cloud.bigquery.Job#waitFor(RetryOption...)} */
protected void waitForJobToFinish(JobId jobId) {
JobStatus status = bigQuery.getJob(jobId).getStatus();
@@ -89,15 +87,4 @@ public abstract class BaseBigQueryItemWriterTest {
}
}
@JsonPropertyOrder(value = {"name", "age"})
public record PersonDto(String name, Integer age) {
public static Schema getBigQuerySchema() {
Field nameField = Field.newBuilder("name", StandardSQLTypeName.STRING).build();
Field ageField = Field.newBuilder("age", StandardSQLTypeName.INT64).build();
return Schema.of(nameField, ageField);
}
}
}

2
spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/unit/base/AbstractBigQueryTest.java
View File

@@ -19,6 +19,4 @@ public abstract class AbstractBigQueryTest {
return mockedBigQuery;
}
public record PersonDto(String name) {}
}

10
spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/unit/writer/builder/BigQueryCsvItemWriterBuilderTests.java
View File

@@ -27,13 +27,15 @@ 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.springframework.batch.extensions.bigquery.common.PersonDto;
import org.springframework.batch.extensions.bigquery.common.TestConstants;
import org.springframework.batch.extensions.bigquery.unit.base.AbstractBigQueryTest;
import org.springframework.batch.extensions.bigquery.writer.BigQueryCsvItemWriter;
import org.springframework.batch.extensions.bigquery.writer.builder.BigQueryCsvItemWriterBuilder;
class BigQueryCsvItemWriterBuilderTests extends AbstractBigQueryTest {
private static final String DATASET_NAME = "my_dataset";
private static final String TABLE = "persons_csv";
private final Log logger = LogFactory.getLog(getClass());
@@ -44,10 +46,10 @@ class BigQueryCsvItemWriterBuilderTests extends AbstractBigQueryTest {
void testCsvWriterWithRowMapper() {
BigQuery mockedBigQuery = prepareMockedBigQuery();
CsvMapper csvMapper = new CsvMapper();
DatasetInfo datasetInfo = DatasetInfo.newBuilder(DATASET_NAME).setLocation("europe-west-2").build();
DatasetInfo datasetInfo = DatasetInfo.newBuilder(TestConstants.DATASET).setLocation("europe-west-2").build();
WriteChannelConfiguration writeConfiguration = WriteChannelConfiguration
.newBuilder(TableId.of(datasetInfo.getDatasetId().getDataset(), "csv_table"))
.newBuilder(TableId.of(datasetInfo.getDatasetId().getDataset(), TABLE))
.setAutodetect(true)
.setFormatOptions(FormatOptions.csv())
.build();
@@ -70,7 +72,7 @@ class BigQueryCsvItemWriterBuilderTests extends AbstractBigQueryTest {
BigQuery mockedBigQuery = prepareMockedBigQuery();
WriteChannelConfiguration writeConfiguration = WriteChannelConfiguration
.newBuilder(TableId.of(DATASET_NAME, "csv_table"))
.newBuilder(TableId.of(TestConstants.DATASET, TABLE))
.setAutodetect(true)
.setFormatOptions(FormatOptions.csv())
.build();

8
spring-batch-bigquery/src/test/java/org/springframework/batch/extensions/bigquery/unit/writer/builder/BigQueryJsonItemWriterBuilderTests.java
View File

@@ -29,14 +29,14 @@ 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.springframework.batch.extensions.bigquery.common.PersonDto;
import org.springframework.batch.extensions.bigquery.common.TestConstants;
import org.springframework.batch.extensions.bigquery.unit.base.AbstractBigQueryTest;
import org.springframework.batch.extensions.bigquery.writer.BigQueryJsonItemWriter;
import org.springframework.batch.extensions.bigquery.writer.builder.BigQueryJsonItemWriterBuilder;
class BigQueryJsonItemWriterBuilderTests extends AbstractBigQueryTest {
private static final String DATASET_NAME = "my_dataset";
private final Log logger = LogFactory.getLog(getClass());
/**
@@ -48,7 +48,7 @@ class BigQueryJsonItemWriterBuilderTests extends AbstractBigQueryTest {
ObjectMapper objectMapper = new ObjectMapper();
WriteChannelConfiguration writeConfiguration = WriteChannelConfiguration
.newBuilder(TableId.of(DATASET_NAME, "json_table"))
.newBuilder(TableId.of(TestConstants.DATASET, "persons_json"))
.setFormatOptions(FormatOptions.json())
.setSchema(Schema.of(
Field.newBuilder("name", StandardSQLTypeName.STRING).setMode(Field.Mode.REQUIRED).build()
@@ -72,7 +72,7 @@ class BigQueryJsonItemWriterBuilderTests extends AbstractBigQueryTest {
BigQuery mockedBigQuery = prepareMockedBigQuery();
WriteChannelConfiguration writeConfiguration = WriteChannelConfiguration
.newBuilder(TableId.of(DATASET_NAME, "json_table"))
.newBuilder(TableId.of(TestConstants.DATASET, "persons_json"))
.setAutodetect(true)
.setFormatOptions(FormatOptions.json())
.build();