com.redislabs.riot.redis.reader.RediSearchDocumentReader Maven / Gradle / Ivy

Go to download

package com.redislabs.riot.redis.reader;

import java.util.Iterator;
import java.util.Map;

import org.springframework.batch.item.support.AbstractItemCountingItemStreamItemReader;
import org.springframework.util.Assert;
import org.springframework.util.ClassUtils;

import com.redislabs.lettusearch.StatefulRediSearchConnection;
import com.redislabs.lettusearch.search.Document;
import com.redislabs.lettusearch.search.SearchOptions;

public class RediSearchDocumentReader extends AbstractItemCountingItemStreamItemReader> {

	private StatefulRediSearchConnection connection;
	private String index;
	private String query;
	private SearchOptions options;
	private FieldExtractor extractor;
	private Object lock = new Object();
	private Iterator> resultIterator;

	public RediSearchDocumentReader(StatefulRediSearchConnection connection, String index, String query,
			SearchOptions options, FieldExtractor extractor) {
		setName(ClassUtils.getShortName(RediSearchDocumentReader.class));
		this.connection = connection;
		this.index = index;
		this.query = query;
		this.options = options;
		this.extractor = extractor;
	}

	@Override
	protected void doOpen() throws Exception {
		Assert.state(connection == null, "Cannot open an already open ItemReader, call close first");
		resultIterator = connection.sync().search(index, query, options).iterator();
	}

	@Override
	protected void doClose() throws Exception {
		synchronized (lock) {
			connection.close();
			connection = null;
			resultIterator = null;
		}
	}

	@SuppressWarnings({ "rawtypes" })
	@Override
	protected Map doRead() throws Exception {
		synchronized (lock) {
			if (resultIterator.hasNext()) {
				Document result = resultIterator.next();
				Map fields = extractor.getFields(result.getId());
				result.putAll(fields);
				return result;
			}
		}
		return null;
	}

}