BATCH-2267: added ItemReader and ItemWriter implementations for elasticsearch
This commit is contained in:
Hasnain Javed
@@ -0,0 +1,89 @@
|
||||
/*
|
||||
* Copyright 2002-2014 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
|
||||
*
|
||||
* http://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.data;
|
||||
|
||||
import static org.slf4j.LoggerFactory.getLogger;
|
||||
import static org.springframework.util.Assert.state;
|
||||
import static org.springframework.util.ClassUtils.getShortName;
|
||||
|
||||
import java.util.Iterator;
|
||||
|
||||
import org.slf4j.Logger;
|
||||
import org.springframework.batch.item.ExecutionContext;
|
||||
import org.springframework.batch.item.ItemReader;
|
||||
import org.springframework.beans.factory.InitializingBean;
|
||||
import org.springframework.data.domain.Pageable;
|
||||
import org.springframework.data.elasticsearch.core.ElasticsearchOperations;
|
||||
import org.springframework.data.elasticsearch.core.query.SearchQuery;
|
||||
|
||||
/**
|
||||
* <p>
|
||||
* Restartable {@link ItemReader} that reads documents from Elasticsearch
|
||||
* via a paging technique.
|
||||
* </p>
|
||||
*
|
||||
* <p>
|
||||
* It executes the query object {@link SearchQuery} to retrieve the requested
|
||||
* documents. The query is executed using paged requests specified in the
|
||||
* {@link org.springframework.data.elasticsearch.core.query.AbstractQuery#setPageable(Pageable pageable)}.
|
||||
* Additional pages are requested as needed to provide data when the {@link #read()} method is called.
|
||||
* </p>
|
||||
*
|
||||
* <p>
|
||||
* The implementation is thread-safe between calls to
|
||||
* {@link #open(ExecutionContext)}, but remember to use <code>saveState=false</code>
|
||||
* if used in a multi-threaded client (no restart available).
|
||||
* </p>
|
||||
*
|
||||
*
|
||||
* @author Hasnain Javed
|
||||
* @since 3.x.x
|
||||
*/
|
||||
public class ElasticsearchItemReader<T> extends AbstractPaginatedDataItemReader<T> implements InitializingBean {
|
||||
|
||||
private final Logger logger;
|
||||
|
||||
private final ElasticsearchOperations elasticsearchOperations;
|
||||
|
||||
private final SearchQuery query;
|
||||
|
||||
private final Class<? extends T> targetType;
|
||||
|
||||
public ElasticsearchItemReader(ElasticsearchOperations elasticsearchOperations, SearchQuery query, Class<? extends T> targetType) {
|
||||
setName(getShortName(getClass()));
|
||||
logger = getLogger(getClass());
|
||||
this.elasticsearchOperations = elasticsearchOperations;
|
||||
this.query = query;
|
||||
this.targetType = targetType;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void afterPropertiesSet() throws Exception {
|
||||
state(elasticsearchOperations != null, "An ElasticsearchOperations implementation is required.");
|
||||
state(query != null, "A query is required.");
|
||||
state(targetType != null, "A target type to convert the input into is required.");
|
||||
}
|
||||
|
||||
@Override
|
||||
@SuppressWarnings("unchecked")
|
||||
protected Iterator<T> doPageRead() {
|
||||
|
||||
logger.debug("executing query {}", query.getQuery());
|
||||
|
||||
return (Iterator<T>)elasticsearchOperations.queryForList(query, targetType).iterator();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
/*
|
||||
* Copyright 2002-2014 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
|
||||
*
|
||||
* http://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.data;
|
||||
|
||||
import static java.lang.String.valueOf;
|
||||
import static java.util.UUID.randomUUID;
|
||||
import static org.slf4j.LoggerFactory.getLogger;
|
||||
import static org.springframework.transaction.support.TransactionSynchronizationManager.bindResource;
|
||||
import static org.springframework.transaction.support.TransactionSynchronizationManager.getResource;
|
||||
import static org.springframework.transaction.support.TransactionSynchronizationManager.hasResource;
|
||||
import static org.springframework.transaction.support.TransactionSynchronizationManager.isActualTransactionActive;
|
||||
import static org.springframework.transaction.support.TransactionSynchronizationManager.registerSynchronization;
|
||||
import static org.springframework.transaction.support.TransactionSynchronizationManager.unbindResource;
|
||||
import static org.springframework.util.Assert.state;
|
||||
import static org.springframework.util.CollectionUtils.isEmpty;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import org.slf4j.Logger;
|
||||
import org.springframework.batch.item.ItemWriter;
|
||||
import org.springframework.beans.factory.InitializingBean;
|
||||
import org.springframework.data.elasticsearch.core.ElasticsearchOperations;
|
||||
import org.springframework.data.elasticsearch.core.query.IndexQuery;
|
||||
import org.springframework.transaction.support.TransactionSynchronizationAdapter;
|
||||
|
||||
/**
|
||||
* <p>
|
||||
* A {@link ItemWriter} implementation that writes to Elasticsearch store using an implementation of Spring Data's
|
||||
* {@link ElasticsearchOperations}. Similar to MongoDB, Elasticsearch is not a transactional store,
|
||||
* the strategy for writing data is similar to {@link MongoItemWriter}. There is no roll back if an error occurs
|
||||
* </p>
|
||||
*
|
||||
* <p>
|
||||
* This writer is thread-safe once all properties are set (normal singleton behavior) so it can be used in multiple
|
||||
* concurrent transactions.
|
||||
* </p>
|
||||
*
|
||||
* @author Hasnain Javed
|
||||
* @since 3.x.x
|
||||
*
|
||||
*/
|
||||
public class ElasticsearchItemWriter implements ItemWriter<IndexQuery>, InitializingBean {
|
||||
|
||||
private final String dataKey;
|
||||
private final Logger logger;
|
||||
private ElasticsearchOperations elasticsearchOperations;
|
||||
private boolean delete;
|
||||
|
||||
public ElasticsearchItemWriter(ElasticsearchOperations elasticsearchOperations) {
|
||||
super();
|
||||
dataKey = valueOf(randomUUID());
|
||||
logger = getLogger(getClass());
|
||||
delete = false;
|
||||
this.elasticsearchOperations = elasticsearchOperations;
|
||||
}
|
||||
|
||||
/**
|
||||
* A flag for removing items given to the writer. Default value is set to false indicating that the items will be saved.
|
||||
* otherwise, the items will be removed.
|
||||
*
|
||||
* @param delete flag
|
||||
*/
|
||||
public void setDelete(boolean delete) {
|
||||
this.delete = delete;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void afterPropertiesSet() throws Exception {
|
||||
state(elasticsearchOperations != null, "An ElasticsearchOperations implementation is required.");
|
||||
}
|
||||
|
||||
@Override
|
||||
public void write(List<? extends IndexQuery> items) throws Exception {
|
||||
|
||||
if(isActualTransactionActive()) {
|
||||
addToBuffer(items);
|
||||
}else {
|
||||
writeItems(items);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Writes to Elasticsearch via the template.
|
||||
* This can be overridden by a subclass if required.
|
||||
*
|
||||
* @param items the list of items to be indexed.
|
||||
*/
|
||||
protected void writeItems(List<? extends IndexQuery> items) {
|
||||
|
||||
if(isEmpty(items)) {
|
||||
logger.warn("no items to write to elasticsearch. list is empty or null");
|
||||
}else {
|
||||
|
||||
for(IndexQuery item : items) {
|
||||
|
||||
if(delete) {
|
||||
String id = item.getId();
|
||||
logger.debug("deleting item with id {}", id);
|
||||
elasticsearchOperations.delete(item.getObject().getClass(), id);
|
||||
}else {
|
||||
String id = elasticsearchOperations.index(item);
|
||||
logger.debug("added item to elasticsearch with id {}", id);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@SuppressWarnings("unchecked")
|
||||
private void addToBuffer(List<? extends IndexQuery> items) {
|
||||
|
||||
if(hasResource(dataKey)) {
|
||||
logger.debug("appending items to buffer under key {}", dataKey);
|
||||
List<IndexQuery> buffer = (List<IndexQuery>) getResource(dataKey);
|
||||
buffer.addAll(items);
|
||||
}else {
|
||||
logger.debug("adding items to buffer under key {}", dataKey);
|
||||
bindResource(dataKey, items);
|
||||
registerSynchronization(new TransactionSynchronizationCallbackImpl());
|
||||
}
|
||||
}
|
||||
|
||||
private class TransactionSynchronizationCallbackImpl extends TransactionSynchronizationAdapter {
|
||||
@SuppressWarnings("unchecked")
|
||||
@Override
|
||||
public void beforeCommit(boolean readOnly) {
|
||||
|
||||
List<IndexQuery> items = (List<IndexQuery>) getResource(dataKey);
|
||||
if(!isEmpty(items)) {
|
||||
if(!readOnly) {
|
||||
writeItems(items);
|
||||
}else{
|
||||
logger.warn("can not write items to elasticsearch as transaction is read only");
|
||||
}
|
||||
}else {
|
||||
logger.warn("no items to write to elasticsearch. list is empty or null");
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public void afterCompletion(int status) {
|
||||
if(hasResource(dataKey)) {
|
||||
logger.debug("removing items from buffer under key {}", dataKey);
|
||||
unbindResource(dataKey);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
/*
|
||||
* Copyright 2002-2014 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
|
||||
*
|
||||
* http://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.data;
|
||||
|
||||
import static java.util.Arrays.asList;
|
||||
import static org.junit.Assert.assertEquals;
|
||||
import static org.junit.Assert.fail;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
import static org.mockito.MockitoAnnotations.initMocks;
|
||||
|
||||
import org.junit.After;
|
||||
import org.junit.Before;
|
||||
import org.junit.Test;
|
||||
import org.mockito.Mock;
|
||||
import org.springframework.data.elasticsearch.core.ElasticsearchOperations;
|
||||
import org.springframework.data.elasticsearch.core.query.NativeSearchQueryBuilder;
|
||||
import org.springframework.data.elasticsearch.core.query.SearchQuery;
|
||||
|
||||
public class ElasticsearchItemReaderTest {
|
||||
|
||||
private ElasticsearchItemReader<Object> reader;
|
||||
|
||||
@Mock
|
||||
private ElasticsearchOperations elasticsearchOperations;
|
||||
|
||||
private SearchQuery query;
|
||||
|
||||
@Before
|
||||
public void setUp() throws Exception {
|
||||
initMocks(this);
|
||||
query = new NativeSearchQueryBuilder().build();
|
||||
reader = new ElasticsearchItemReader<Object>(elasticsearchOperations, query, Object.class);
|
||||
reader.afterPropertiesSet();
|
||||
}
|
||||
|
||||
@After
|
||||
public void tearDown() {
|
||||
query = null;
|
||||
elasticsearchOperations = null;
|
||||
}
|
||||
|
||||
@Test(expected=IllegalStateException.class)
|
||||
public void shouldFailAssertionOnNullElasticsearchOperations() throws Exception {
|
||||
|
||||
try {
|
||||
new ElasticsearchItemReader<Object>(null, null, null).afterPropertiesSet();
|
||||
fail("Assertion should have thrown exception on null ElasticsearchOperations");
|
||||
}catch(IllegalStateException e) {
|
||||
assertEquals("An ElasticsearchOperations implementation is required.", e.getMessage());
|
||||
throw e;
|
||||
}
|
||||
}
|
||||
|
||||
@Test(expected=IllegalStateException.class)
|
||||
public void shouldFailAssertionOnNullQuery() throws Exception {
|
||||
|
||||
try {
|
||||
new ElasticsearchItemReader<Object>(elasticsearchOperations, null, null).afterPropertiesSet();
|
||||
fail("Assertion shold have thrown exception on null Query");
|
||||
}catch(IllegalStateException e) {
|
||||
assertEquals("A query is required.", e.getMessage());
|
||||
throw e;
|
||||
}
|
||||
}
|
||||
|
||||
@Test(expected=IllegalStateException.class)
|
||||
public void shouldFailAssertionOnNullTargetType() throws Exception {
|
||||
|
||||
try {
|
||||
new ElasticsearchItemReader<Object>(elasticsearchOperations, query, null).afterPropertiesSet();
|
||||
fail("Assertion shold have thrown exception on null Target Type");
|
||||
}catch(IllegalStateException e) {
|
||||
assertEquals("A target type to convert the input into is required.", e.getMessage());
|
||||
throw e;
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
public void shouldQueryForList() {
|
||||
|
||||
when(elasticsearchOperations.queryForList(query, Object.class)).thenReturn(asList());
|
||||
|
||||
reader.doPageRead();
|
||||
|
||||
verify(elasticsearchOperations).queryForList(query, Object.class);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,244 @@
|
||||
/*
|
||||
* Copyright 2002-2014 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
|
||||
*
|
||||
* http://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.data;
|
||||
|
||||
import static java.util.Arrays.asList;
|
||||
import static org.junit.Assert.fail;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.verifyZeroInteractions;
|
||||
import static org.mockito.MockitoAnnotations.initMocks;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
|
||||
import org.junit.After;
|
||||
import org.junit.Before;
|
||||
import org.junit.Test;
|
||||
import org.mockito.Mock;
|
||||
import org.springframework.batch.support.transaction.ResourcelessTransactionManager;
|
||||
import org.springframework.data.elasticsearch.annotations.Document;
|
||||
import org.springframework.data.elasticsearch.core.ElasticsearchOperations;
|
||||
import org.springframework.data.elasticsearch.core.query.IndexQuery;
|
||||
import org.springframework.data.elasticsearch.core.query.IndexQueryBuilder;
|
||||
import org.springframework.transaction.TransactionStatus;
|
||||
import org.springframework.transaction.support.TransactionCallback;
|
||||
import org.springframework.transaction.support.TransactionTemplate;
|
||||
|
||||
public class ElasticsearchItemWriterTest {
|
||||
|
||||
@Document(indexName="test_index", type="test_type")
|
||||
public class DummyDocument {
|
||||
|
||||
private String id;
|
||||
|
||||
public String getId() {
|
||||
return id;
|
||||
}
|
||||
|
||||
public void setId(String id) {
|
||||
this.id = id;
|
||||
}
|
||||
}
|
||||
|
||||
private ElasticsearchItemWriter writer;
|
||||
|
||||
@Mock
|
||||
private ElasticsearchOperations elasticsearchOperations;
|
||||
|
||||
private TransactionTemplate transactionTemplate;
|
||||
|
||||
private DummyDocument dummyDocument;
|
||||
|
||||
@Before
|
||||
public void setUp() throws Exception {
|
||||
initMocks(this);
|
||||
transactionTemplate = new TransactionTemplate(new ResourcelessTransactionManager());
|
||||
writer = new ElasticsearchItemWriter(elasticsearchOperations);
|
||||
writer.afterPropertiesSet();
|
||||
dummyDocument = new DummyDocument();
|
||||
}
|
||||
|
||||
@After
|
||||
public void tearDown() {
|
||||
transactionTemplate = null;
|
||||
writer = null;
|
||||
dummyDocument = null;
|
||||
}
|
||||
|
||||
@Test(expected=IllegalStateException.class)
|
||||
public void shouldFailAssertion() throws Exception {
|
||||
|
||||
new ElasticsearchItemWriter(null).afterPropertiesSet();
|
||||
fail("Assertion shold have thrown exception on null ElasticsearchOperations");
|
||||
}
|
||||
|
||||
@Test
|
||||
public void shouldNotWriteWhenNoTransactionIsActiveAndNoItem() throws Exception {
|
||||
|
||||
writer.write(null);
|
||||
verifyZeroInteractions(elasticsearchOperations);
|
||||
|
||||
writer.write(new ArrayList<IndexQuery>(0));
|
||||
verifyZeroInteractions(elasticsearchOperations);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void shouldWriteItemWhenNoTransactionIsActive() throws Exception {
|
||||
|
||||
IndexQueryBuilder builder = new IndexQueryBuilder();
|
||||
builder.withObject(dummyDocument);
|
||||
|
||||
List<IndexQuery> items = asList(builder.build());
|
||||
|
||||
writer.write(items);
|
||||
|
||||
verify(elasticsearchOperations).index(items.iterator().next());
|
||||
}
|
||||
|
||||
@Test
|
||||
public void shouldWriteItemWhenInTransaction() throws Exception {
|
||||
|
||||
IndexQueryBuilder builder = new IndexQueryBuilder();
|
||||
builder.withObject(dummyDocument);
|
||||
|
||||
final List<IndexQuery> items = asList(builder.build());
|
||||
|
||||
transactionTemplate.execute(new TransactionCallback<Void>() {
|
||||
|
||||
@Override
|
||||
public Void doInTransaction(TransactionStatus status) {
|
||||
try {
|
||||
writer.write(items);
|
||||
} catch (Exception e) {
|
||||
fail("An error occurred while writing: " + e.getMessage());
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
});
|
||||
|
||||
verify(elasticsearchOperations).index(items.iterator().next());
|
||||
}
|
||||
|
||||
@Test
|
||||
public void shouldNotWriteItemWhenTransactionFails() throws Exception {
|
||||
|
||||
IndexQueryBuilder builder = new IndexQueryBuilder();
|
||||
builder.withObject(dummyDocument);
|
||||
|
||||
final List<IndexQuery> items = asList(builder.build());
|
||||
|
||||
try {
|
||||
transactionTemplate.execute(new TransactionCallback<Void>() {
|
||||
|
||||
@Override
|
||||
public Void doInTransaction(TransactionStatus status) {
|
||||
try {
|
||||
writer.write(items);
|
||||
} catch (Exception ignore) {
|
||||
fail("unexpected error occurred");
|
||||
}
|
||||
throw new RuntimeException("rollback");
|
||||
}
|
||||
});
|
||||
} catch (RuntimeException re) {
|
||||
// ignore
|
||||
} catch (Throwable t) {
|
||||
fail("Unexpected error occurred");
|
||||
}
|
||||
|
||||
verifyZeroInteractions(elasticsearchOperations);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void shouldNotWriteItemWhenTransactionIsReadOnly() throws Exception {
|
||||
|
||||
IndexQueryBuilder builder = new IndexQueryBuilder();
|
||||
builder.withObject(dummyDocument);
|
||||
|
||||
final List<IndexQuery> items = asList(builder.build());
|
||||
|
||||
try {
|
||||
|
||||
transactionTemplate.setReadOnly(true);
|
||||
transactionTemplate.execute(new TransactionCallback<Void>() {
|
||||
|
||||
@Override
|
||||
public Void doInTransaction(TransactionStatus status) {
|
||||
try {
|
||||
writer.write(items);
|
||||
} catch (Exception ignore) {
|
||||
fail("unexpected error occurred");
|
||||
}
|
||||
return null;
|
||||
}
|
||||
});
|
||||
} catch (Throwable t) {
|
||||
fail("unexpected error occurred");
|
||||
}
|
||||
|
||||
verifyZeroInteractions(elasticsearchOperations);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void shouldRemoveItemWhenNoTransactionIsActive() throws Exception {
|
||||
|
||||
writer.setDelete(true);
|
||||
|
||||
dummyDocument.setId("123456");
|
||||
|
||||
IndexQueryBuilder builder = new IndexQueryBuilder();
|
||||
builder.withId(dummyDocument.getId());
|
||||
builder.withObject(dummyDocument);
|
||||
|
||||
final List<IndexQuery> items = asList(builder.build());
|
||||
|
||||
writer.write(items);
|
||||
|
||||
verify(elasticsearchOperations).delete(items.iterator().next().getObject().getClass(), items.iterator().next().getId());
|
||||
}
|
||||
|
||||
@Test
|
||||
public void shouldRemoveItemWhenInTransaction() throws Exception {
|
||||
|
||||
writer.setDelete(true);
|
||||
|
||||
dummyDocument.setId("123456");
|
||||
|
||||
IndexQueryBuilder builder = new IndexQueryBuilder();
|
||||
builder.withId(dummyDocument.getId());
|
||||
builder.withObject(dummyDocument);
|
||||
|
||||
final List<IndexQuery> items = asList(builder.build());
|
||||
|
||||
transactionTemplate.execute(new TransactionCallback<Void>() {
|
||||
|
||||
@Override
|
||||
public Void doInTransaction(TransactionStatus status) {
|
||||
try {
|
||||
writer.write(items);
|
||||
} catch (Exception e) {
|
||||
fail("An error occurred while writing: " + e.getMessage());
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
});
|
||||
|
||||
verify(elasticsearchOperations).delete(items.iterator().next().getObject().getClass(), items.iterator().next().getId());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user