一篇文章教你学会SpringBatch使用分区读取数据库

Spring Batch 支持使用分区来读取数据库中的数据并并行处理。在讲SpringBatch使用分区读取数据库之前先介绍下EnableBatchProcessing。

目录

一、@EnableBatchProcessing

二、Spring Batch 使用分区来读取数据库

---

一、@EnableBatchProcessing

@EnableBatchProcessing 注解会自动配置一个数据源（DataSource）用于 Spring Batch 的作业存储。它使用了一个嵌入式数据库（通常是H2）作为默认的数据源，用于存储 Spring Batch 的元数据（metadata），例如作业的状态、运行历史和参数等。

如果你想要配置自定义数据源，你可以通过 BatchConfigurer 接口来覆盖默认的配置。例如：

import javax.sql.DataSource;

import org.springframework.batch.core.configuration.annotation.DefaultBatchConfigurer;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class CustomBatchConfigurer extends DefaultBatchConfigurer {

    @Override
    public void setDataSource(DataSource dataSource) {
        // 自定义数据源的设置
        super.setDataSource(dataSource);
    }
}

在上述示例中，CustomBatchConfigurer 类扩展了 DefaultBatchConfigurer，并覆盖了 setDataSource 方法，以便设置自定义数据源。这使你可以使用你自己的数据库来存储 Spring Batch 的元数据。

总之，@EnableBatchProcessing 会自动初始化一个数据源用于 Spring Batch 的元数据存储，但你可以配置自定义数据源以满足你的特定需求。在生产环境中，通常会使用一个持久的数据库来存储元数据，而在开发和测试环境中可以使用嵌入式数据库。

二、Spring Batch 使用分区来读取数据库

1、创建一个 Data 类来表示数据库中的数据：

public class Data {
    private int id;
    private String name;

    // getters and setters
}

2、创建一个 Spring Batch 配置类

import org.springframework.batch.core.Step;
import org.springframework.batch.core.configuration.annotation.EnableBatchProcessing;
import org.springframework.batch.core.configuration.annotation.JobBuilderFactory;
import org.springframework.batch.core.configuration.annotation.StepBuilderFactory;
import org.springframework.batch.core.partition.PartitionHandler;
import org.springframework.batch.core.partition.StepExecutionSplitter;
import org.springframework.batch.core.partition.support.Partitioner;
import org.springframework.batch.core.partition.support.TaskExecutorPartitionHandler;
import org.springframework.batch.item.database.JdbcPagingItemReader;
import org.springframework.batch.item.database.PagingQueryProvider;
import org.springframework.batch.item.database.support.SqlPagingQueryProviderFactoryBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.datasource.DataSourceUtils;

import javax.sql.DataSource;

@Configuration
@EnableBatchProcessing
public class BatchConfiguration {

    @Bean
    public Step partitionedStep(StepBuilderFactory stepBuilderFactory, Partitioner partitioner,
                                PartitionHandler partitionHandler) {
        return stepBuilderFactory.get("partitionedStep")
            .partitioner("workerStep", partitioner())
            .step(workerStep())
            .gridSize(4) // 设置分区数量
            .taskExecutor(new SimpleAsyncTaskExecutor())
            .partitionHandler(partitionHandler())
            .build();
    }

    @Bean
    public Step workerStep() {
        return stepBuilderFactory.get("workerStep")
            .<Data, Data>chunk(1)
            .reader(itemReader(null))
            .writer(itemWriter())
            .build();
    }

    @Bean
    public Partitioner partitioner(DataSource dataSource) {
        return gridSize -> {
            Map<String, ExecutionContext> partitions = new HashMap<>();
            PagingQueryProvider queryProvider = pagingQueryProvider(dataSource);
            JdbcPagingItemReader<Data> reader = itemReader(dataSource, queryProvider);
            PartitionHandler handler = partitionHandler(dataSource, reader);
            for (int i = 0; i < gridSize; i++) {
                ExecutionContext context = new ExecutionContext();
                context.put("partitionId", i);
                context.put("minValue", i * 100); // 设置分区查询条件
                context.put("maxValue", (i + 1) * 100 - 1);
                partitions.put("partition" + i, context);
            }
            return partitions;
        };
    }

    @Bean
    public PagingQueryProvider pagingQueryProvider(DataSource dataSource) {
        SqlPagingQueryProviderFactoryBean factoryBean = new SqlPagingQueryProviderFactoryBean();
        factoryBean.setSelectClause("SELECT id, name");
        factoryBean.setFromClause("FROM user");
        factoryBean.setWhereClause("WHERE id >= :minValue AND id <= :maxValue");
        factoryBean.setSortKey("id");
        factoryBean.setDataSource(dataSource);
        return factoryBean.getObject();
    }

    @Bean
    public JdbcPagingItemReader<Data> itemReader(DataSource dataSource, PagingQueryProvider queryProvider) {
        JdbcPagingItemReader<Data> reader = new JdbcPagingItemReader<>();
        reader.setDataSource(dataSource);
        reader.setPageSize(100); // 设置每次读取的数据量
        reader.setQueryProvider(queryProvider);
        reader.setRowMapper((resultSet, i) -> {
            Data data = new Data();
            data.setId(resultSet.getInt("id"));
            data.setName(resultSet.getString("name"));
            return data;
        });
        return reader;
    }

    @Bean
    public PartitionHandler partitionHandler(DataSource dataSource, JdbcPagingItemReader<Data> reader) {
        TaskExecutorPartitionHandler handler = new TaskExecutorPartitionHandler();
        handler.setStep(workerStep());
        handler.setGridSize(4); // 设置分区数量
        handler.setTaskExecutor(new SimpleAsyncTaskExecutor());
        handler.setJobExplorer(jobExplorer(dataSource));
        return handler;
    }

    @Bean
    public StepExecutionSplitter stepExecutionSplitter(DataSource dataSource) {
        return new StepExecutionSplitter(jobExplorer(dataSource), "partitionedStep");
    }

    @Bean
    public JdbcTemplate jdbcTemplate(DataSource dataSource) {
        return new JdbcTemplate(dataSource);
    }
}

在上面的配置中，我们定义了一个包含两个步骤的作业：partitionedStep 和 workerStep。partitionedStep 是一个分区步骤，负责使用分区读取数据库的数据。workerStep 是用于实际处理数据的步骤。我们还配置了一个分区处理器，指定分区数量和任务执行器。

3、创建一个执行器类，用于启动 Spring Batch 作业

import org.springframework.batch.core.Job;
import org.springframework.batch.core.JobParameters;
import org.springframework.batch.core.launch.JobLauncher;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

@Component
public class JobExecutor {

    @Autowired
    private JobLauncher jobLauncher;

    @Autowired
    private Job job;

    public void executeJob() throws Exception {
        jobLauncher.run(job, new JobParameters());
    }
}

这个执行器类通过注入 JobLauncher 和作业 Job 来启动 Spring Batch 作业。