com.redis.spring.batch.reader.KeyValueReadOperation 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.reader;

import java.io.IOException;
import java.util.Iterator;
import java.util.List;
import java.util.function.Function;
import java.util.function.UnaryOperator;

import org.springframework.batch.item.Chunk;
import org.springframework.util.CollectionUtils;
import org.springframework.util.unit.DataSize;

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.CodecUtils;

import io.lettuce.core.AbstractRedisClient;
import io.lettuce.core.RedisFuture;
import io.lettuce.core.api.async.BaseRedisAsyncCommands;
import io.lettuce.core.codec.RedisCodec;

public class KeyValueReadOperation implements Operation> {

	public enum Type {
		DUMP, STRUCT
	}

	private static final String FILENAME = "keyvalue.lua";

	private final Evalsha evalsha;

	private final Function toStringValueFunction;

	private UnaryOperator> postOperator = UnaryOperator.identity();

	public KeyValueReadOperation(AbstractRedisClient client, RedisCodec codec, DataSize memLimit, int samples,
			Type type) throws IOException {
		this.evalsha = new Evalsha<>(FILENAME, client, codec);
		this.evalsha.setKeyFunction(Function.identity());
		this.evalsha.setArgs(memLimit.toBytes(), samples, type.name().toLowerCase());
		this.toStringValueFunction = CodecUtils.toStringValueFunction(codec);
	}

	public void setPostOperator(UnaryOperator> operator) {
		this.postOperator = operator;
	}

	@Override
	public void execute(BaseRedisAsyncCommands commands, Chunk inputs,
			Chunk>> outputs) {
		Chunk>> evalshaOutputs = new Chunk<>();
		evalsha.execute(commands, inputs, evalshaOutputs);
		for (RedisFuture> future : evalshaOutputs) {
			outputs.add(new MappingRedisFuture<>(future.toCompletableFuture(), this::toKeyValue));
		}
	}

	@SuppressWarnings("unchecked")
	protected KeyValue toKeyValue(List