org.springframework.batch.item.redis.DataStructureItemWriter 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.*;
import io.lettuce.core.api.StatefulConnection;
import io.lettuce.core.api.async.*;
import io.lettuce.core.cluster.RedisClusterClient;
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.batch.item.redis.support.DataStructure;
import org.springframework.core.convert.converter.Converter;

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

@SuppressWarnings("unchecked")
public class DataStructureItemWriter extends AbstractPipelineItemWriter {

    private final Converter key;
    private final Converter value;
    private final Converter dataType;
    private final Converter absoluteTTL;

    public DataStructureItemWriter(Supplier> connectionSupplier, GenericObjectPoolConfig> poolConfig, Function, BaseRedisAsyncCommands> async, Converter key, Converter value, Converter dataType, Converter absoluteTTL) {
        super(connectionSupplier, poolConfig, async);
        this.key = key;
        this.value = value;
        this.dataType = dataType;
        this.absoluteTTL = absoluteTTL;
    }

    @Override
    protected void write(BaseRedisAsyncCommands commands, long timeout, List items) {
        try {
            List> futures = new ArrayList<>(items.size());
            for (T item : items) {
                String key = this.key.convert(item);
                Object value = this.value.convert(item);
                if (value == null) {
                    futures.add(((RedisKeyAsyncCommands) commands).del(key));
                    continue;
                }
                String type = this.dataType.convert(item);
                if (type == null) {
                    continue;
                }
                switch (type.toLowerCase()) {
                    case DataStructure.STRING:
                        futures.add(((RedisStringAsyncCommands) commands).set(key, (String) value));
                        break;
                    case DataStructure.LIST:
                        futures.add(((RedisListAsyncCommands) commands).rpush(key, ((Collection) value).toArray(new String[0])));
                        break;
                    case DataStructure.SET:
                        futures.add(((RedisSetAsyncCommands) commands).sadd(key, ((Collection) value).toArray(new String[0])));
                        break;
                    case DataStructure.ZSET:
                        futures.add(((RedisSortedSetAsyncCommands) commands).zadd(key, ((Collection>) value).toArray(new ScoredValue[0])));
                        break;
                    case DataStructure.HASH:
                        futures.add(((RedisHashAsyncCommands) commands).hset(key, (Map) value));
                        break;
                    case DataStructure.STREAM:
                        Collection> messages = (Collection>) value;
                        for (StreamMessage message : messages) {
                            futures.add(((RedisStreamAsyncCommands) commands).xadd(key, new XAddArgs().id(message.getId()), message.getBody()));
                        }
                        break;
                }
                Long ttl = absoluteTTL.convert(item);
                if (ttl == null) {
                    continue;
                }
                if (ttl > 0) {
                    futures.add(((RedisKeyAsyncCommands) commands).pexpireat(key, ttl));
                }
            }
            commands.flushCommands();
            LettuceFutures.awaitAll(timeout, TimeUnit.MILLISECONDS, futures.toArray(new RedisFuture[0]));
        } finally {
            commands.setAutoFlushCommands(true);
        }
    }

    public static DataStructureItemWriterBuilder client(RedisClient client) {
        return new DataStructureItemWriterBuilder(client);
    }

    public static DataStructureItemWriterBuilder client(RedisClusterClient client) {
        return new DataStructureItemWriterBuilder(client);
    }

    public static class DataStructureItemWriterBuilder extends CommandBuilder {

        public DataStructureItemWriterBuilder(RedisClusterClient client) {
            super(client);
        }

        public DataStructureItemWriterBuilder(RedisClient client) {
            super(client);
        }

        public DataStructureItemWriter build() {
            return new DataStructureItemWriter<>(connectionSupplier, poolConfig, async, DataStructure::getKey, DataStructure::getValue, DataStructure::getType, DataStructure::getAbsoluteTTL);
        }

    }

}