com.redis.spring.batch.writer.operation.StructWrite Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of spring-batch-redis Show documentation
Spring Batch Redis
The newest version!
package com.redis.spring.batch.writer.operation;

import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.function.Predicate;
import java.util.stream.Collector;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

import org.springframework.batch.item.Chunk;

import com.redis.lettucemod.timeseries.AddOptions;
import com.redis.lettucemod.timeseries.DuplicatePolicy;
import com.redis.spring.batch.common.DataType;
import com.redis.spring.batch.common.KeyValue;
import com.redis.spring.batch.common.Operation;
import com.redis.spring.batch.util.Predicates;

import io.lettuce.core.RedisFuture;
import io.lettuce.core.StreamMessage;
import io.lettuce.core.XAddArgs;
import io.lettuce.core.api.async.BaseRedisAsyncCommands;

public class StructWrite implements Operation, Object> {

	private final Collector, ?, Map>>> groupByType = Collectors
			.groupingBy(KeyValue::getType);

	private Predicate> existPredicate() {
		return KeyValue::exists;
	}

	@SuppressWarnings("unchecked")
	private final Predicate> expirePredicate = Predicates.and(existPredicate(), k -> k.getTtl() > 0);

	private Predicate> deletePredicate = Predicates.negate(existPredicate());

	private final Operation, Object> deleteOperation = deleteOperation();

	private final Operation, Object> expireOperation = expireOperation();

	private final Operation, Object> hashOperation = hashOperation();

	private final Operation, Object> jsonOperation = jsonOperation();

	private final Operation, Object> listOperation = listOperation();

	private final Operation, Object> setOperation = setOperation();

	private final Operation, Object> streamOperation = streamOperation();

	private final Operation, Object> stringOperation = stringOperation();

	private final Operation, Object> timeseriesOperation = timeseriesOperation();

	private final Operation, Object> zsetOperation = zsetOperation();

	private final Operation, Object> noOperation = noOperation();

	private Noop> noOperation() {
		return new Noop<>();
	}

	public void setOverwrite(boolean overwrite) {
		if (overwrite) {
			deletePredicate = Predicates.isTrue();
		}
	}

	@Override
	public void execute(BaseRedisAsyncCommands commands, Chunk> items,
			Chunk> futures) {

		Chunk> toDelete = new Chunk<>(
				StreamSupport.stream(items.spliterator(), false).filter(deletePredicate).collect(Collectors.toList()));
		deleteOperation.execute(commands, toDelete, futures);
		Map>> toWrite = StreamSupport.stream(items.spliterator(), false)
				.filter(KeyValue::exists).collect(groupByType);
		for (Entry>> entry : toWrite.entrySet()) {
			operation(entry.getKey()).execute(commands, new Chunk<>(entry.getValue()), futures);
		}
		List> toExpire = StreamSupport.stream(items.spliterator(), false).filter(expirePredicate)
				.collect(Collectors.toList());
		expireOperation.execute(commands, new Chunk<>(toExpire), futures);
	}

	private Operation, Object> operation(DataType type) {
		switch (type) {
		case HASH:
			return hashOperation;
		case JSON:
			return jsonOperation;
		case LIST:
			return listOperation;
		case SET:
			return setOperation;
		case STREAM:
			return streamOperation;
		case STRING:
			return stringOperation;
		case TIMESERIES:
			return timeseriesOperation;
		case ZSET:
			return zsetOperation;
		default:
			return noOperation;
		}
	}

	private Hset> hashOperation() {
		return new Hset<>(KeyValue::getKey, this::value);
	}

	private Set> stringOperation() {
		return new Set<>(KeyValue::getKey, this::value);
	}

	private XAddArgs xaddArgs(StreamMessage message) {
		XAddArgs args = new XAddArgs();
		if (message.getId() != null) {
			args.id(message.getId());
		}
		return args;
	}

	private XAddAll> streamOperation() {
		XAddAll> operation = new XAddAll<>();
		operation.setMessagesFunction(this::value);
		operation.setArgsFunction(this::xaddArgs);
		return operation;
	}

	private TsAddAll> timeseriesOperation() {
		TsAddAll> operation = new TsAddAll<>();
		operation.setKeyFunction(KeyValue::getKey);
		operation.setOptions(AddOptions.builder().policy(DuplicatePolicy.LAST).build());
		operation.setSamplesFunction(this::value);
		return operation;
	}

	private ZaddAll> zsetOperation() {
		return new ZaddAll<>(KeyValue::getKey, this::value);
	}

	private SaddAll> setOperation() {
		return new SaddAll<>(KeyValue::getKey, this::value);
	}

	private RpushAll> listOperation() {
		return new RpushAll<>(KeyValue::getKey, this::value);
	}

	private JsonSet> jsonOperation() {
		return new JsonSet<>(KeyValue::getKey, this::value);
	}

	private Operation, Object> deleteOperation() {
		return new Del<>(KeyValue::getKey);
	}

	private ExpireAt> expireOperation() {
		ExpireAt> operation = new ExpireAt<>(KeyValue::getKey);
		operation.setEpochFunction(KeyValue::getTtl);
		return operation;
	}

	@SuppressWarnings("unchecked")
	private  O value(KeyValue struct) {
		return (O) struct.getValue();
	}

}