com.redis.spring.batch.util.ConnectionUtils 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
There is a newer version: 4.0.7
package com.redis.spring.batch.util;

import java.io.IOException;
import java.io.InputStream;
import java.util.function.Supplier;

import org.springframework.batch.item.ItemStreamException;
import org.springframework.util.FileCopyUtils;

import com.redis.lettucemod.RedisModulesClient;
import com.redis.lettucemod.cluster.RedisModulesClusterClient;
import com.redis.spring.batch.reader.KeyValueItemProcessor;

import io.lettuce.core.AbstractRedisClient;
import io.lettuce.core.ReadFrom;
import io.lettuce.core.api.StatefulConnection;
import io.lettuce.core.api.StatefulRedisConnection;
import io.lettuce.core.api.sync.RedisScriptingCommands;
import io.lettuce.core.cluster.api.StatefulRedisClusterConnection;
import io.lettuce.core.codec.RedisCodec;
import io.lettuce.core.codec.StringCodec;

public abstract class ConnectionUtils {

    private ConnectionUtils() {
    }

    public static Supplier> supplier(AbstractRedisClient client) {
        return supplier(client, StringCodec.UTF8);
    }

    public static Supplier> supplier(AbstractRedisClient client, ReadFrom readFrom) {
        return supplier(client, StringCodec.UTF8, readFrom);
    }

    public static  Supplier> supplier(AbstractRedisClient client, RedisCodec codec) {
        return supplier(client, codec, null);
    }

    public static  Supplier> supplier(AbstractRedisClient client, RedisCodec codec,
            ReadFrom readFrom) {
        if (client instanceof RedisModulesClusterClient) {
            return () -> {
                StatefulRedisClusterConnection connection = ((RedisModulesClusterClient) client).connect(codec);
                if (readFrom != null) {
                    connection.setReadFrom(readFrom);
                }
                return connection;
            };
        }
        return () -> ((RedisModulesClient) client).connect(codec);
    }

    @SuppressWarnings("unchecked")
    public static  T sync(StatefulConnection connection) {
        if (connection instanceof StatefulRedisClusterConnection) {
            return (T) ((StatefulRedisClusterConnection) connection).sync();
        }
        return (T) ((StatefulRedisConnection) connection).sync();
    }

    @SuppressWarnings("unchecked")
    public static  T async(StatefulConnection connection) {
        if (connection instanceof StatefulRedisClusterConnection) {
            return (T) ((StatefulRedisClusterConnection) connection).async();
        }
        return (T) ((StatefulRedisConnection) connection).async();
    }

    @SuppressWarnings("unchecked")
    public static String loadScript(AbstractRedisClient client, String filename) {
        byte[] bytes;
        try (InputStream inputStream = KeyValueItemProcessor.class.getClassLoader().getResourceAsStream(filename)) {
            bytes = FileCopyUtils.copyToByteArray(inputStream);
        } catch (IOException e) {
            throw new ItemStreamException("Could not read LUA script file " + filename);
        }
        try (StatefulConnection connection = supplier(client, StringCodec.UTF8).get()) {
            return ((RedisScriptingCommands) sync(connection)).scriptLoad(bytes);
        }
    }

}