com.redis.spring.batch.writer.StructOperation 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.writer;

import java.util.List;
import java.util.function.Function;

import com.redis.spring.batch.KeyValue;
import com.redis.spring.batch.RedisItemWriter;
import com.redis.spring.batch.RedisItemWriter.MergePolicy;
import com.redis.spring.batch.RedisItemWriter.StreamIdPolicy;
import com.redis.spring.batch.RedisItemWriter.TtlPolicy;
import com.redis.spring.batch.writer.operation.Del;
import com.redis.spring.batch.writer.operation.ExpireAt;
import com.redis.spring.batch.writer.operation.Hset;
import com.redis.spring.batch.writer.operation.JsonSet;
import com.redis.spring.batch.writer.operation.Noop;
import com.redis.spring.batch.writer.operation.Restore;
import com.redis.spring.batch.writer.operation.RpushAll;
import com.redis.spring.batch.writer.operation.SaddAll;
import com.redis.spring.batch.writer.operation.Set;
import com.redis.spring.batch.writer.operation.TsAddAll;
import com.redis.spring.batch.writer.operation.XAddAll;
import com.redis.spring.batch.writer.operation.ZaddAll;

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

public class StructOperation implements Operation> {

    private static final XAddArgs EMPTY_XADD_ARGS = new XAddArgs();

    private final ExpireAt> expire = expire();

    private final Noop> noop = new Noop<>();

    private final Del> del = del();

    private final Hset> hset = hset();

    private final Set> set = set();

    private final JsonSet> jsonSet = jsonSet();

    private final RpushAll> rpush = rpushAll();

    private final SaddAll> sadd = saddAll();

    private final ZaddAll> zadd = zaddAll();

    private final TsAddAll> tsAdd = tsAddAll();

    private final XAddAll> xadd = xaddAll();

    private TtlPolicy ttlPolicy = RedisItemWriter.DEFAULT_TTL_POLICY;

    private MergePolicy mergePolicy = RedisItemWriter.DEFAULT_MERGE_POLICY;

    private StreamIdPolicy streamIdPolicy = RedisItemWriter.DEFAULT_STREAM_ID_POLICY;

    public StructOperation ttlPolicy(TtlPolicy ttlPolicy) {
        this.ttlPolicy = ttlPolicy;
        return this;
    }

    private XAddAll> xaddAll() {
        XAddAll> operation = new XAddAll<>();
        operation.setMessages(value());
        operation.setArgs(this::xaddArgs);
        return operation;
    }

    private TsAddAll> tsAddAll() {
        TsAddAll> operation = new TsAddAll<>();
        operation.setKey(key());
        operation.setSamples(value());
        return operation;
    }

    private ZaddAll> zaddAll() {
        ZaddAll> operation = new ZaddAll<>();
        operation.setKey(key());
        operation.setValues(value());
        return operation;
    }

    private SaddAll> saddAll() {
        SaddAll> operation = new SaddAll<>();
        operation.setKey(key());
        operation.setValues(value());
        return operation;
    }

    private RpushAll> rpushAll() {
        RpushAll> operation = new RpushAll<>();
        operation.setKey(key());
        operation.setValues(value());
        return operation;
    }

    private JsonSet> jsonSet() {
        JsonSet> operation = new JsonSet<>();
        operation.setKey(key());
        operation.setValue(value());
        return operation;
    }

    private Set> set() {
        Set> operation = new Set<>();
        operation.setKey(key());
        operation.setValue(value());
        return operation;
    }

    private Hset> hset() {
        Hset> operation = new Hset<>();
        operation.setKey(key());
        operation.setMap(value());
        return operation;
    }

    private Del> del() {
        Del> operation = new Del<>();
        operation.setKey(key());
        return operation;
    }

    private ExpireAt> expire() {
        ExpireAt> operation = new ExpireAt<>();
        operation.setKey(key());
        operation.setEpoch(KeyValue::getTtl);
        return operation;
    }

    public StructOperation mergePolicy(MergePolicy mergePolicy) {
        this.mergePolicy = mergePolicy;
        return this;
    }

    public StructOperation streamIdPolicy(StreamIdPolicy streamIdPolicy) {
        this.streamIdPolicy = streamIdPolicy;
        return this;
    }

    private static  Function, K> key() {
        return KeyValue::getKey;
    }

    @SuppressWarnings("unchecked")
    private static  Function, T> value() {
        return kv -> (T) kv.getValue();
    }

    @Override
    public void execute(BaseRedisAsyncCommands commands, KeyValue item, List> futures) {
        if (shouldSkip(item)) {
            return;
        }
        if (shouldDelete(item)) {
            del.execute(commands, item, futures);
            return;
        }
        if (isOverwrite() && !KeyValue.isString(item)) {
            del.execute(commands, item, futures);
        }
        operation(item).execute(commands, item, futures);
        if (ttlPolicy == TtlPolicy.PROPAGATE && item.getTtl() > 0) {
            expire.execute(commands, item, futures);
        }
    }

    private boolean isOverwrite() {
        return mergePolicy == MergePolicy.OVERWRITE;
    }

    public static boolean shouldSkip(KeyValue item) {
        return item == null || item.getKey() == null || (item.getValue() == null && KeyValue.hasMemoryUsage(item));
    }

    private boolean shouldDelete(KeyValue item) {
        return item.getValue() == null || Restore.TTL_KEY_DOES_NOT_EXIST == item.getTtl() || KeyValue.isNone(item);
    }

    private Operation> operation(KeyValue item) {
        switch (item.getType()) {
            case KeyValue.HASH:
                return hset;
            case KeyValue.JSON:
                return jsonSet;
            case KeyValue.LIST:
                return rpush;
            case KeyValue.SET:
                return sadd;
            case KeyValue.STREAM:
                return xadd;
            case KeyValue.STRING:
                return set;
            case KeyValue.TIMESERIES:
                return tsAdd;
            case KeyValue.ZSET:
                return zadd;
            default:
                return noop;
        }
    }

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

}