diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/database/JpaCursorItemReader.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/database/JpaCursorItemReader.java new file mode 100644 index 000000000..24afe03e3 --- /dev/null +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/database/JpaCursorItemReader.java @@ -0,0 +1,153 @@ +/* + * Copyright 2020 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.item.database; + +import java.util.Iterator; +import java.util.Map; + +import javax.persistence.EntityManager; +import javax.persistence.EntityManagerFactory; +import javax.persistence.Query; + +import org.springframework.batch.item.ExecutionContext; +import org.springframework.batch.item.ItemStreamException; +import org.springframework.batch.item.database.orm.JpaQueryProvider; +import org.springframework.batch.item.support.AbstractItemCountingItemStreamItemReader; +import org.springframework.beans.factory.InitializingBean; +import org.springframework.dao.DataAccessResourceFailureException; +import org.springframework.util.Assert; +import org.springframework.util.ClassUtils; + +/** + * {@link org.springframework.batch.item.ItemStreamReader} implementation based + * on JPA {@link Query#getResultStream()}. It executes the JPQL query when + * initialized and iterates over the result set as {@link #read()} method is called, + * returning an object corresponding to the current row. The query can be set + * directly using {@link #setQueryString(String)}, or using a query provider via + * {@link #setQueryProvider(JpaQueryProvider)}. + * + * The implementation is not thread-safe. + * + * @author Mahmoud Ben Hassine + * @param type of items to read + * @since 4.3 + */ +public class JpaCursorItemReader extends AbstractItemCountingItemStreamItemReader + implements InitializingBean { + + private EntityManagerFactory entityManagerFactory; + private EntityManager entityManager; + private String queryString; + private JpaQueryProvider queryProvider; + private Map parameterValues; + private Iterator iterator; + + /** + * Create a new {@link JpaCursorItemReader}. + */ + public JpaCursorItemReader() { + setName(ClassUtils.getShortName(JpaCursorItemReader.class)); + } + + /** + * Set the JPA entity manager factory. + * + * @param entityManagerFactory JPA entity manager factory + */ + public void setEntityManagerFactory(EntityManagerFactory entityManagerFactory) { + this.entityManagerFactory = entityManagerFactory; + } + + /** + * Set the JPA query provider. + * + * @param queryProvider JPA query provider + */ + public void setQueryProvider(JpaQueryProvider queryProvider) { + this.queryProvider = queryProvider; + } + + /** + * Set the JPQL query string. + * + * @param queryString JPQL query string + */ + public void setQueryString(String queryString) { + this.queryString = queryString; + } + + /** + * Set the parameter values to be used for the query execution. + * + * @param parameterValues the values keyed by parameter names used in + * the query string. + */ + public void setParameterValues(Map parameterValues) { + this.parameterValues = parameterValues; + } + + @Override + public void afterPropertiesSet() throws Exception { + Assert.notNull(this.entityManagerFactory, "EntityManagerFactory is required"); + if (this.queryProvider == null) { + Assert.hasLength(this.queryString, "Query string is required when queryProvider is null"); + } + } + + @Override + @SuppressWarnings("unchecked") + protected void doOpen() throws Exception { + this.entityManager = this.entityManagerFactory.createEntityManager(); + if (this.entityManager == null) { + throw new DataAccessResourceFailureException("Unable to create an EntityManager"); + } + if (this.queryProvider != null) { + this.queryProvider.setEntityManager(this.entityManager); + } + Query query = createQuery(); + if (this.parameterValues != null) { + this.parameterValues.forEach(query::setParameter); + } + this.iterator = query.getResultStream().iterator(); + } + + private Query createQuery() { + if (this.queryProvider == null) { + return this.entityManager.createQuery(this.queryString); + } + else { + return this.queryProvider.createQuery(); + } + } + + @Override + protected T doRead() { + return this.iterator.hasNext() ? this.iterator.next() : null; + } + + @Override + public void update(ExecutionContext executionContext) throws ItemStreamException { + super.update(executionContext); + this.entityManager.clear(); + } + + @Override + protected void doClose() { + if (this.entityManager != null) { + this.entityManager.close(); + } + } +} diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/database/builder/JpaCursorItemReaderBuilder.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/database/builder/JpaCursorItemReaderBuilder.java new file mode 100644 index 000000000..0a143d4f7 --- /dev/null +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/database/builder/JpaCursorItemReaderBuilder.java @@ -0,0 +1,182 @@ +/* + * Copyright 2020 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.item.database.builder; + +import java.util.Map; + +import javax.persistence.EntityManagerFactory; + +import org.springframework.batch.item.ExecutionContext; +import org.springframework.batch.item.ItemStreamSupport; +import org.springframework.batch.item.database.JpaCursorItemReader; +import org.springframework.batch.item.database.orm.JpaQueryProvider; +import org.springframework.batch.item.support.AbstractItemCountingItemStreamItemReader; +import org.springframework.util.Assert; + +/** + * Builder for {@link JpaCursorItemReader}. + * + * @author Mahmoud Ben Hassine + * + * @since 4.3 + */ +public class JpaCursorItemReaderBuilder { + + private EntityManagerFactory entityManagerFactory; + private String queryString; + private JpaQueryProvider queryProvider; + private Map parameterValues; + private boolean saveState = true; + private String name; + private int maxItemCount = Integer.MAX_VALUE; + private int currentItemCount; + + /** + * Configure if the state of the {@link ItemStreamSupport} + * should be persisted within the {@link ExecutionContext} + * for restart purposes. + * + * @param saveState defaults to true + * @return The current instance of the builder. + */ + public JpaCursorItemReaderBuilder saveState(boolean saveState) { + this.saveState = saveState; + + return this; + } + + /** + * The name used to calculate the key within the {@link ExecutionContext}. + * Required if {@link #saveState(boolean)} is set to true. + * + * @param name name of the reader instance + * @return The current instance of the builder. + * @see ItemStreamSupport#setName(String) + */ + public JpaCursorItemReaderBuilder name(String name) { + this.name = name; + + return this; + } + + /** + * Configure the max number of items to be read. + * + * @param maxItemCount the max items to be read + * @return The current instance of the builder. + * @see AbstractItemCountingItemStreamItemReader#setMaxItemCount(int) + */ + public JpaCursorItemReaderBuilder maxItemCount(int maxItemCount) { + this.maxItemCount = maxItemCount; + + return this; + } + + /** + * Index for the current item. Used on restarts to indicate where to start from. + * + * @param currentItemCount current index + * @return this instance for method chaining + * @see AbstractItemCountingItemStreamItemReader#setCurrentItemCount(int) + */ + public JpaCursorItemReaderBuilder currentItemCount(int currentItemCount) { + this.currentItemCount = currentItemCount; + + return this; + } + + /** + * A map of parameter values to be set on the query. The key of the map is + * the name of the parameter to be set with the value being the value to be set. + * + * @param parameterValues map of values + * @return this instance for method chaining + * @see JpaCursorItemReader#setParameterValues(Map) + */ + public JpaCursorItemReaderBuilder parameterValues(Map parameterValues) { + this.parameterValues = parameterValues; + + return this; + } + + /** + * A query provider. This should be set only if {@link #queryString(String)} + * have not been set. + * + * @param queryProvider the query provider + * @return this instance for method chaining + * @see JpaCursorItemReader#setQueryProvider(JpaQueryProvider) + */ + public JpaCursorItemReaderBuilder queryProvider(JpaQueryProvider queryProvider) { + this.queryProvider = queryProvider; + + return this; + } + + /** + * The JPQL query string to execute. This should only be set if + * {@link #queryProvider(JpaQueryProvider)} has not been set. + * + * @param queryString the JPQL query + * @return this instance for method chaining + * @see JpaCursorItemReader#setQueryString(String) + */ + public JpaCursorItemReaderBuilder queryString(String queryString) { + this.queryString = queryString; + + return this; + } + + /** + * The {@link EntityManagerFactory} to be used for executing the configured + * {@link #queryString}. + * + * @param entityManagerFactory {@link EntityManagerFactory} used to create + * {@link javax.persistence.EntityManager} + * @return this instance for method chaining + */ + public JpaCursorItemReaderBuilder entityManagerFactory(EntityManagerFactory entityManagerFactory) { + this.entityManagerFactory = entityManagerFactory; + + return this; + } + + /** + * Returns a fully constructed {@link JpaCursorItemReader}. + * + * @return a new {@link JpaCursorItemReader} + */ + public JpaCursorItemReader build() { + Assert.notNull(this.entityManagerFactory, "An EntityManagerFactory is required"); + if (this.saveState) { + Assert.hasText(this.name, "A name is required when saveState is set to true"); + } + if (this.queryProvider == null) { + Assert.hasLength(this.queryString, "Query string is required when queryProvider is null"); + } + + JpaCursorItemReader reader = new JpaCursorItemReader<>(); + reader.setEntityManagerFactory(this.entityManagerFactory); + reader.setQueryProvider(this.queryProvider); + reader.setQueryString(this.queryString); + reader.setParameterValues(this.parameterValues); + reader.setCurrentItemCount(this.currentItemCount); + reader.setMaxItemCount(this.maxItemCount); + reader.setSaveState(this.saveState); + reader.setName(this.name); + return reader; + } +} diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/JpaCursorItemReaderCommonTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/JpaCursorItemReaderCommonTests.java new file mode 100644 index 000000000..8eca55293 --- /dev/null +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/JpaCursorItemReaderCommonTests.java @@ -0,0 +1,56 @@ +/* + * Copyright 2020 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.item.database; + +import org.springframework.batch.item.ExecutionContext; +import org.springframework.batch.item.ItemReader; +import org.springframework.batch.item.sample.Foo; +import org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean; +import org.springframework.orm.jpa.vendor.HibernateJpaVendorAdapter; + +/** + * @author Mahmoud Ben Hassine + */ +public class JpaCursorItemReaderCommonTests extends + AbstractDatabaseItemStreamItemReaderTests { + + @Override + protected ItemReader getItemReader() throws Exception { + LocalContainerEntityManagerFactoryBean factoryBean = + new LocalContainerEntityManagerFactoryBean(); + factoryBean.setDataSource(getDataSource()); + factoryBean.setPersistenceUnitName("bar"); + factoryBean.setJpaVendorAdapter(new HibernateJpaVendorAdapter()); + factoryBean.afterPropertiesSet(); + + String jpqlQuery = "from Foo"; + JpaCursorItemReader itemReader = new JpaCursorItemReader<>(); + itemReader.setQueryString(jpqlQuery); + itemReader.setEntityManagerFactory(factoryBean.getObject()); + itemReader.afterPropertiesSet(); + itemReader.setSaveState(true); + return itemReader; + } + + @Override + protected void pointToEmptyInput(ItemReader tested) throws Exception { + JpaCursorItemReader reader = (JpaCursorItemReader) tested; + reader.close(); + reader.setQueryString("from Foo foo where foo.id = -1"); + reader.afterPropertiesSet(); + reader.open(new ExecutionContext()); + } +} diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/builder/JpaCursorItemReaderBuilderTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/builder/JpaCursorItemReaderBuilderTests.java new file mode 100644 index 000000000..e831ee4d6 --- /dev/null +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/database/builder/JpaCursorItemReaderBuilderTests.java @@ -0,0 +1,263 @@ +/* + * Copyright 2020 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.item.database.builder; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import javax.persistence.EntityManagerFactory; +import javax.sql.DataSource; + +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import org.springframework.batch.item.ExecutionContext; +import org.springframework.batch.item.database.JpaCursorItemReader; +import org.springframework.batch.item.database.orm.JpaNamedQueryProvider; +import org.springframework.batch.item.database.orm.JpaNativeQueryProvider; +import org.springframework.batch.item.sample.Foo; +import org.springframework.context.ConfigurableApplicationContext; +import org.springframework.context.annotation.AnnotationConfigApplicationContext; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.core.io.ClassPathResource; +import org.springframework.core.io.Resource; +import org.springframework.jdbc.datasource.embedded.EmbeddedDatabaseFactory; +import org.springframework.jdbc.datasource.init.DataSourceInitializer; +import org.springframework.jdbc.datasource.init.ResourceDatabasePopulator; +import org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean; +import org.springframework.orm.jpa.vendor.HibernateJpaVendorAdapter; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.fail; + +/** + * @author Mahmoud Ben Hassine + */ +public class JpaCursorItemReaderBuilderTests { + + private EntityManagerFactory entityManagerFactory; + private ConfigurableApplicationContext context; + + @Before + public void setUp() { + this.context = new AnnotationConfigApplicationContext(JpaCursorItemReaderBuilderTests.TestDataSourceConfiguration.class); + this.entityManagerFactory = (EntityManagerFactory) context.getBean("entityManagerFactory"); + } + + @After + public void tearDown() { + if(this.context != null) { + this.context.close(); + } + } + + @Test + public void testConfiguration() throws Exception { + JpaCursorItemReader reader = new JpaCursorItemReaderBuilder() + .name("fooReader") + .entityManagerFactory(this.entityManagerFactory) + .currentItemCount(2) + .maxItemCount(4) + .queryString("select f from Foo f ") + .build(); + + reader.afterPropertiesSet(); + + ExecutionContext executionContext = new ExecutionContext(); + + reader.open(executionContext); + Foo item1 = reader.read(); + Foo item2 = reader.read(); + assertNull(reader.read()); + reader.update(executionContext); + reader.close(); + + assertEquals(3, item1.getId()); + assertEquals("bar3", item1.getName()); + assertEquals(3, item1.getValue()); + assertEquals(4, item2.getId()); + assertEquals("bar4", item2.getName()); + assertEquals(4, item2.getValue()); + + assertEquals(2, executionContext.size()); + } + + @Test + public void testConfigurationNoSaveState() throws Exception { + Map parameters = new HashMap<>(); + parameters.put("value", 2); + + JpaCursorItemReader reader = new JpaCursorItemReaderBuilder() + .name("fooReader") + .entityManagerFactory(this.entityManagerFactory) + .queryString("select f from Foo f where f.id > :value") + .parameterValues(parameters) + .saveState(false) + .build(); + + reader.afterPropertiesSet(); + + ExecutionContext executionContext = new ExecutionContext(); + + reader.open(executionContext); + + int i = 0; + while(reader.read() != null) { + i++; + } + + reader.update(executionContext); + reader.close(); + + assertEquals(3, i); + assertEquals(0, executionContext.size()); + } + + @Test + public void testConfigurationNamedQueryProvider() throws Exception { + JpaNamedQueryProvider namedQueryProvider = new JpaNamedQueryProvider<>(); + namedQueryProvider.setNamedQuery("allFoos"); + namedQueryProvider.setEntityClass(Foo.class); + namedQueryProvider.afterPropertiesSet(); + + JpaCursorItemReader reader = new JpaCursorItemReaderBuilder() + .name("fooReader") + .entityManagerFactory(this.entityManagerFactory) + .queryProvider(namedQueryProvider) + .build(); + + reader.afterPropertiesSet(); + + ExecutionContext executionContext = new ExecutionContext(); + reader.open(executionContext); + + Foo foo; + List foos = new ArrayList<>(); + + while((foo = reader.read()) != null) { + foos.add(foo); + } + + reader.update(executionContext); + reader.close(); + + int id = 0; + for (Foo testFoo:foos) { + assertEquals(++id, testFoo.getId()); + } + } + + @Test + public void testConfigurationNativeQueryProvider() throws Exception { + + JpaNativeQueryProvider provider = new JpaNativeQueryProvider<>(); + provider.setEntityClass(Foo.class); + provider.setSqlQuery("select * from T_FOOS"); + provider.afterPropertiesSet(); + + JpaCursorItemReader reader = new JpaCursorItemReaderBuilder() + .name("fooReader") + .entityManagerFactory(this.entityManagerFactory) + .queryProvider(provider) + .build(); + + reader.afterPropertiesSet(); + + ExecutionContext executionContext = new ExecutionContext(); + + reader.open(executionContext); + + int i = 0; + while(reader.read() != null) { + i++; + } + + reader.update(executionContext); + reader.close(); + + assertEquals(5, i); + } + + @Test + public void testValidation() { + try { + new JpaCursorItemReaderBuilder().build(); + fail("An EntityManagerFactory is required"); + } + catch (IllegalArgumentException iae) { + assertEquals("An EntityManagerFactory is required", iae.getMessage()); + } + + try { + new JpaCursorItemReaderBuilder() + .entityManagerFactory(this.entityManagerFactory) + .saveState(true) + .build(); + fail("A name is required when saveState is set to true"); + } + catch (IllegalArgumentException iae) { + assertEquals("A name is required when saveState is set to true", iae.getMessage()); + } + + try { + new JpaCursorItemReaderBuilder() + .entityManagerFactory(this.entityManagerFactory) + .saveState(false) + .build(); + fail("Query string is required when queryProvider is null"); + } + catch (IllegalArgumentException iae) { + assertEquals("Query string is required when queryProvider is null", iae.getMessage()); + } + } + + @Configuration + public static class TestDataSourceConfiguration { + + @Bean + public DataSource dataSource() { + return new EmbeddedDatabaseFactory().getDatabase(); + } + + @Bean + public DataSourceInitializer initializer(DataSource dataSource) { + DataSourceInitializer dataSourceInitializer = new DataSourceInitializer(); + dataSourceInitializer.setDataSource(dataSource); + + Resource create = new ClassPathResource("org/springframework/batch/item/database/init-foo-schema-hsqldb.sql"); + dataSourceInitializer.setDatabasePopulator(new ResourceDatabasePopulator(create)); + + return dataSourceInitializer; + } + + @Bean + public LocalContainerEntityManagerFactoryBean entityManagerFactory() { + LocalContainerEntityManagerFactoryBean entityManagerFactoryBean = + new LocalContainerEntityManagerFactoryBean(); + + entityManagerFactoryBean.setDataSource(dataSource()); + entityManagerFactoryBean.setPersistenceUnitName("bar"); + entityManagerFactoryBean.setJpaVendorAdapter(new HibernateJpaVendorAdapter()); + + return entityManagerFactoryBean; + } + } +}