org.springframework.batch.item.redis.OperationItemWriter 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 reader and writer implementations for Redis
There is a newer version: 2.16.0
package org.springframework.batch.item.redis;

import io.lettuce.core.LettuceFutures;
import io.lettuce.core.RedisClient;
import io.lettuce.core.RedisFuture;
import io.lettuce.core.TransactionResult;
import io.lettuce.core.api.StatefulConnection;
import io.lettuce.core.api.async.BaseRedisAsyncCommands;
import io.lettuce.core.api.async.RedisTransactionalAsyncCommands;
import io.lettuce.core.cluster.RedisClusterClient;
import lombok.Setter;
import lombok.experimental.Accessors;
import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
import org.springframework.batch.item.redis.support.AbstractPipelineItemWriter;
import org.springframework.batch.item.redis.support.CommandBuilder;
import org.springframework.util.Assert;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import java.util.function.Supplier;

public class OperationItemWriter extends AbstractPipelineItemWriter {

    private final RedisOperation operation;

    public OperationItemWriter(Supplier> connectionSupplier, GenericObjectPoolConfig> poolConfig, Function, BaseRedisAsyncCommands> async, RedisOperation operation) {
        super(connectionSupplier, poolConfig, async);
        Assert.notNull(operation, "A Redis operation is required");
        this.operation = operation;
    }

    @Override
    protected void write(BaseRedisAsyncCommands commands, long timeout, List items) {
        List> futures = write(commands, items);
        commands.flushCommands();
        LettuceFutures.awaitAll(timeout, TimeUnit.MILLISECONDS, futures.toArray(new RedisFuture[0]));
    }

    protected List> write(BaseRedisAsyncCommands commands, List items) {
        List> futures = new ArrayList<>(items.size());
        for (T item : items) {
            futures.add(operation.execute(commands, item));
        }
        return futures;
    }

    public static class TransactionItemWriter extends OperationItemWriter {

        public TransactionItemWriter(Supplier> connectionSupplier, GenericObjectPoolConfig> poolConfig, Function, BaseRedisAsyncCommands> async, RedisOperation operation) {
            super(connectionSupplier, poolConfig, async, operation);
        }

        @SuppressWarnings("unchecked")
        @Override
        protected List> write(BaseRedisAsyncCommands commands, List items) {
            RedisFuture multiFuture = ((RedisTransactionalAsyncCommands) commands).multi();
            List> futures = super.write(commands, items);
            RedisFuture execFuture = ((RedisTransactionalAsyncCommands) commands).exec();
            List> allFutures = new ArrayList<>(futures.size() + 2);
            allFutures.add(multiFuture);
            allFutures.addAll(futures);
            allFutures.add(execFuture);
            return allFutures;
        }

    }

    public interface RedisOperation {

        RedisFuture execute(BaseRedisAsyncCommands commands, T item);

    }

    public static  OperationItemWriterBuilder operation(RedisOperation operation) {
        return new OperationItemWriterBuilder<>(operation);
    }

    public static class OperationItemWriterBuilder {

        private final RedisOperation operation;

        public OperationItemWriterBuilder(RedisOperation operation) {
            this.operation = operation;
        }

        public CommandOperationItemWriterBuilder client(RedisClient client) {
            return new CommandOperationItemWriterBuilder<>(client, operation);
        }

        public CommandOperationItemWriterBuilder client(RedisClusterClient client) {
            return new CommandOperationItemWriterBuilder<>(client, operation);
        }

    }


    @Setter
    @Accessors(fluent = true)
    public static class CommandOperationItemWriterBuilder extends CommandBuilder> {

        private final RedisOperation operation;
        private boolean transactional;

        public CommandOperationItemWriterBuilder(RedisClient client, RedisOperation operation) {
            super(client);
            this.operation = operation;
        }

        public CommandOperationItemWriterBuilder(RedisClusterClient client, RedisOperation operation) {
            super(client);
            this.operation = operation;
        }

        public OperationItemWriter build() {
            if (transactional) {
                return new TransactionItemWriter<>(connectionSupplier, poolConfig, async, operation);
            }
            return new OperationItemWriter<>(connectionSupplier, poolConfig, async, operation);
        }

    }


}