All Downloads are FREE. Search and download functionalities are using the official Maven repository.

org.redisson.spring.data.connection.RedissonReactiveZSetCommands Maven / Gradle / Ivy

There is a newer version: 3.39.0
Show newest version
/**
 * Copyright (c) 2013-2024 Nikita Koksharov
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.redisson.spring.data.connection;

import org.reactivestreams.Publisher;
import org.redisson.client.codec.ByteArrayCodec;
import org.redisson.client.codec.DoubleCodec;
import org.redisson.client.codec.LongCodec;
import org.redisson.client.codec.StringCodec;
import org.redisson.client.protocol.RedisCommand;
import org.redisson.client.protocol.RedisCommands;
import org.redisson.client.protocol.RedisStrictCommand;
import org.redisson.client.protocol.convertor.DoubleNullSafeReplayConvertor;
import org.redisson.client.protocol.decoder.ListMultiDecoder2;
import org.redisson.client.protocol.decoder.ObjectSetReplayDecoder;
import org.redisson.reactive.CommandReactiveExecutor;
import org.springframework.data.domain.Range;
import org.springframework.data.domain.Sort.Direction;
import org.springframework.data.redis.connection.DefaultTuple;
import org.springframework.data.redis.connection.ReactiveRedisConnection.CommandResponse;
import org.springframework.data.redis.connection.ReactiveRedisConnection.KeyCommand;
import org.springframework.data.redis.connection.ReactiveRedisConnection.NumericResponse;
import org.springframework.data.redis.connection.ReactiveZSetCommands;
import org.springframework.data.redis.connection.RedisZSetCommands.Tuple;
import org.springframework.util.Assert;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.math.BigDecimal;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

/**
 *
 * @author Nikita Koksharov
 *
 */
public class RedissonReactiveZSetCommands extends RedissonBaseReactive implements ReactiveZSetCommands {

    RedissonReactiveZSetCommands(CommandReactiveExecutor executorService) {
        super(executorService);
    }

    private static final RedisCommand ZADD_FLOAT = new RedisCommand<>("ZADD", new DoubleNullSafeReplayConvertor());

    @Override
    public Flux> zAdd(Publisher commands) {
        return execute(commands, command -> {

            Assert.notNull(command.getKey(), "Key must not be null!");
            Assert.notEmpty(command.getTuples(), "Tuples must not be empty or null!");

            byte[] keyBuf = toByteArray(command.getKey());

            List params = new ArrayList(command.getTuples().size()*2+1);
            params.add(keyBuf);
            if (command.isIncr() || command.isUpsert() || command.isReturnTotalChanged()) {
                if (command.isUpsert()) {
                    params.add("NX");
                } else {
                    params.add("XX");
                }
                if (command.isReturnTotalChanged()) {
                    params.add("CH");
                }
                if (command.isIncr()) {
                    params.add("INCR");
                }
            }

            for (Tuple entry : command.getTuples()) {
                params.add(BigDecimal.valueOf(entry.getScore()).toPlainString());
                params.add(entry.getValue());
            }

            Mono m;
            if (command.isIncr()) {
                m = write(keyBuf, DoubleCodec.INSTANCE, ZADD_FLOAT, params.toArray());
            } else {
                m = write(keyBuf, StringCodec.INSTANCE, RedisCommands.ZADD, params.toArray());
            }
            return m.map(v -> new NumericResponse<>(command, v));
        });
    }

    @Override
    public Flux> zRem(Publisher commands) {
        return execute(commands, command -> {

            Assert.notNull(command.getKey(), "Key must not be null!");
            Assert.notNull(command.getValues(), "Values must not be null!");

            List args = new ArrayList(command.getValues().size() + 1);
            args.add(toByteArray(command.getKey()));
            args.addAll(command.getValues().stream().map(v -> toByteArray(v)).collect(Collectors.toList()));

            Mono m = write((byte[])args.get(0), StringCodec.INSTANCE, RedisCommands.ZREM_LONG, args.toArray());
            return m.map(v -> new NumericResponse<>(command, v));
        });
    }

    @Override
    public Flux> zIncrBy(Publisher commands) {
        return execute(commands, command -> {

            Assert.notNull(command.getKey(), "Key must not be null!");
            Assert.notNull(command.getValue(), "Member must not be null!");
            Assert.notNull(command.getIncrement(), "Increment value must not be null!");

            byte[] keyBuf = toByteArray(command.getKey());
            byte[] valueBuf = toByteArray(command.getValue());
            Mono m = write(keyBuf, DoubleCodec.INSTANCE, RedisCommands.ZINCRBY, keyBuf, new BigDecimal(command.getIncrement().doubleValue()).toPlainString(), valueBuf);
            return m.map(v -> new NumericResponse<>(command, v));
        });
    }

    @Override
    public Flux> zRank(Publisher commands) {
        return execute(commands, command -> {

            Assert.notNull(command.getKey(), "Key must not be null!");
            Assert.notNull(command.getValue(), "Member must not be null!");

            byte[] keyBuf = toByteArray(command.getKey());
            byte[] valueBuf = toByteArray(command.getValue());
            RedisCommand cmd = RedisCommands.ZRANK;
            if (command.getDirection() == Direction.DESC) {
                cmd = RedisCommands.ZREVRANK;
            }
            Mono m = read(keyBuf, DoubleCodec.INSTANCE, cmd, keyBuf, valueBuf);
            return m.map(v -> new NumericResponse<>(command, v));
        });
    }

    private static final RedisCommand> ZRANGE_ENTRY = new RedisCommand<>("ZRANGE", new ScoredSortedSetReplayDecoder());

    private static final RedisCommand> ZRANGE_ENTRY_V2 = new RedisCommand>("ZRANGE",
            new ListMultiDecoder2(new ObjectSetReplayDecoder(), new ScoredSortedSetReplayDecoderV2()));

    private static final RedisCommand> ZRANGE = new RedisCommand<>("ZRANGE", new ObjectSetReplayDecoder());
    private static final RedisCommand> ZREVRANGE_ENTRY = new RedisCommand<>("ZREVRANGE", new ScoredSortedSetReplayDecoder());

    private static final RedisCommand> ZREVRANGE_ENTRY_V2 = new RedisCommand("ZREVRANGE",
            new ListMultiDecoder2(new ObjectSetReplayDecoder(), new ScoredSortedSetReplayDecoderV2()));

    private static final RedisCommand> ZREVRANGE = new RedisCommand<>("ZREVRANGE", new ObjectSetReplayDecoder());

    @Override
    public Flux>> zRange(Publisher commands) {
        return execute(commands, command -> {

            Assert.notNull(command.getKey(), "Key must not be null!");
            Assert.notNull(command.getRange(), "Range must not be null!");

            byte[] keyBuf = toByteArray(command.getKey());

            long start = command.getRange().getLowerBound().getValue().orElse(0L);
            long end = command.getRange().getUpperBound().getValue().get();

            Flux flux;
            if (command.getDirection() == Direction.ASC) {
                if (command.isWithScores()) {
                    RedisCommand> cmd = ZRANGE_ENTRY;
                    if (executorService.getServiceManager().isResp3()) {
                        cmd = ZRANGE_ENTRY_V2;
                    }
                    Mono> m = read(keyBuf, ByteArrayCodec.INSTANCE, cmd,
                                keyBuf, start, end, "WITHSCORES");
                    flux = m.flatMapMany(e -> Flux.fromIterable(e));
                } else {
                    Mono> m = read(keyBuf, ByteArrayCodec.INSTANCE, ZRANGE, keyBuf, start, end);
                    flux = m.flatMapMany(e -> Flux.fromIterable(e).map(b -> new DefaultTuple(b, Double.NaN)));
                }
            } else {
                if (command.isWithScores()) {
                    RedisCommand> cmd = ZREVRANGE_ENTRY;
                    if (executorService.getServiceManager().isResp3()) {
                        cmd = ZREVRANGE_ENTRY_V2;
                    }
                    Mono> m = read(keyBuf, ByteArrayCodec.INSTANCE, cmd,
                                keyBuf, start, end, "WITHSCORES");
                    flux = m.flatMapMany(e -> Flux.fromIterable(e));
                } else {
                    Mono> m = read(keyBuf, ByteArrayCodec.INSTANCE, ZREVRANGE, keyBuf, start, end);
                    flux = m.flatMapMany(e -> Flux.fromIterable(e).map(b -> new DefaultTuple(b, Double.NaN)));
                }
            }
            return Mono.just(new CommandResponse<>(command, flux));
        });
    }

    private static final RedisCommand> ZRANGEBYSCORE = new RedisCommand>("ZRANGEBYSCORE", new ScoredSortedSetReplayDecoder());

    private static final RedisCommand> ZRANGEBYSCORE_V2 = new RedisCommand>("ZRANGEBYSCORE",
            new ListMultiDecoder2(new ObjectSetReplayDecoder(), new ScoredSortedSetReplayDecoderV2()));

    private static final RedisCommand> ZREVRANGEBYSCORE = new RedisCommand>("ZREVRANGEBYSCORE", new ScoredSortedSetReplayDecoder());

    private static final RedisCommand> ZREVRANGEBYSCORE_V2 = new RedisCommand>("ZREVRANGEBYSCORE",
            new ListMultiDecoder2(new ObjectSetReplayDecoder(), new ScoredSortedSetReplayDecoderV2()));

    @Override
    public Flux>> zRangeByScore(
            Publisher commands) {
        return execute(commands, command -> {

            Assert.notNull(command.getKey(), "Key must not be null!");
            Assert.notNull(command.getRange(), "Range must not be null!");

            byte[] keyBuf = toByteArray(command.getKey());

            String start = toLowerBound(command.getRange());
            String end = toUpperBound(command.getRange());

            List args = new ArrayList();
            args.add(keyBuf);
            if (command.getDirection() == Direction.ASC) {
                args.add(start);
            } else {
                args.add(end);
            }
            if (command.getDirection() == Direction.ASC) {
                args.add(end);
            } else {
                args.add(start);
            }
            if (command.isWithScores()) {
                args.add("WITHSCORES");
            }
            if (command.getLimit().isPresent() && !command.getLimit().get().isUnlimited()) {
                args.add("LIMIT");
                args.add(command.getLimit().get().getOffset());
                args.add(command.getLimit().get().getCount());
            }

            Flux flux;
            if (command.getDirection() == Direction.ASC) {
                if (command.isWithScores()) {
                    RedisCommand> cmd = ZRANGEBYSCORE;
                    if (executorService.getServiceManager().isResp3()) {
                        cmd = ZRANGEBYSCORE_V2;
                    }

                    Mono> m = read(keyBuf, ByteArrayCodec.INSTANCE, cmd, args.toArray());
                    flux = m.flatMapMany(e -> Flux.fromIterable(e));
                } else {
                    Mono> m = read(keyBuf, ByteArrayCodec.INSTANCE, RedisCommands.ZRANGEBYSCORE, args.toArray());
                    flux = m.flatMapMany(e -> Flux.fromIterable(e).map(b -> new DefaultTuple(b, Double.NaN)));
                }
            } else {
                if (command.isWithScores()) {
                    RedisCommand> cmd = ZREVRANGEBYSCORE;
                    if (executorService.getServiceManager().isResp3()) {
                        cmd = ZREVRANGEBYSCORE_V2;
                    }
                    Mono> m = read(keyBuf, ByteArrayCodec.INSTANCE, cmd, args.toArray());
                    flux = m.flatMapMany(e -> Flux.fromIterable(e));
                } else {
                    Mono> m = read(keyBuf, ByteArrayCodec.INSTANCE, RedisCommands.ZREVRANGEBYSCORE, args.toArray());
                    flux = m.flatMapMany(e -> Flux.fromIterable(e).map(b -> new DefaultTuple(b, Double.NaN)));
                }
            }

            return Mono.just(new CommandResponse<>(command, flux));
        });
    }

    private static final RedisStrictCommand ZCOUNT = new RedisStrictCommand("ZCOUNT");
    
    String toLowerBound(Range range) {
        StringBuilder s = new StringBuilder();
        if (!range.getLowerBound().isInclusive()) {
            s.append("(");
        }
        if (!range.getLowerBound().getValue().isPresent() || range.getLowerBound().getValue().get().toString().isEmpty()) {
            s.append("-inf");
        } else {
            s.append(range.getLowerBound().getValue().get());
        }
        return s.toString();
    }

    String toUpperBound(Range range) {
        StringBuilder s = new StringBuilder();
        if (!range.getUpperBound().isInclusive()) {
            s.append("(");
        }
        if (!range.getUpperBound().getValue().isPresent() || range.getUpperBound().getValue().get().toString().isEmpty()) {
            s.append("+inf");
        } else {
            s.append(range.getUpperBound().getValue().get());
        }
        return s.toString();
    }
    
    String toLexLowerBound(Range range, Object defaultValue) {
        StringBuilder s = new StringBuilder();
        if (range.getLowerBound().isInclusive()) {
            s.append("[");
        } else {
            s.append("(");
        }
        if (!range.getLowerBound().getValue().isPresent() || range.getLowerBound().getValue().get().toString().isEmpty()) {
            s.append(defaultValue);
        } else {
            s.append(range.getLowerBound().getValue().get());
        }
        return s.toString();
    }

    String toLexUpperBound(Range range, Object defaultValue) {
        StringBuilder s = new StringBuilder();
        if (range.getUpperBound().isInclusive()) {
            s.append("[");
        } else {
            s.append("(");
        }
        if (!range.getUpperBound().getValue().isPresent() || range.getUpperBound().getValue().get().toString().isEmpty()) {
            s.append(defaultValue);
        } else {
            s.append(range.getUpperBound().getValue().get());
        }
        return s.toString();
    }

    @Override
    public Flux> zCount(Publisher commands) {
        return execute(commands, command -> {

            Assert.notNull(command.getKey(), "Key must not be null!");
            Assert.notNull(command.getRange(), "Range must not be null!");

            byte[] keyBuf = toByteArray(command.getKey());
            Mono m = read(keyBuf, StringCodec.INSTANCE, ZCOUNT, 
                                keyBuf, toLowerBound(command.getRange()),
                                toUpperBound(command.getRange()));
            return m.map(v -> new NumericResponse<>(command, v));
        });
    }

    @Override
    public Flux> zCard(Publisher commands) {
        return execute(commands, command -> {

            Assert.notNull(command.getKey(), "Key must not be null!");
            
            byte[] keyBuf = toByteArray(command.getKey());
            Mono m = read(keyBuf, StringCodec.INSTANCE, RedisCommands.ZCARD, keyBuf);
            return m.map(v -> new NumericResponse<>(command, v));
        });
    }

    @Override
    public Flux> zScore(Publisher commands) {
        return execute(commands, command -> {

            Assert.notNull(command.getKey(), "Key must not be null!");
            Assert.notNull(command.getValue(), "Value must not be null!");

            byte[] keyBuf = toByteArray(command.getKey());
            byte[] valueBuf = toByteArray(command.getValue());
            Mono m = read(keyBuf, StringCodec.INSTANCE, RedisCommands.ZSCORE, keyBuf, valueBuf);
            return m.map(v -> new NumericResponse<>(command, v));
        });
    }

    private static final RedisStrictCommand ZREMRANGEBYRANK = new RedisStrictCommand("ZREMRANGEBYRANK");
    
    @Override
    public Flux> zRemRangeByRank(
            Publisher commands) {
        return execute(commands, command -> {

            Assert.notNull(command.getKey(), "Key must not be null!");
            Assert.notNull(command.getRange(), "Range must not be null!");
            
            byte[] keyBuf = toByteArray(command.getKey());
            Mono m = write(keyBuf, StringCodec.INSTANCE, ZREMRANGEBYRANK, 
                                keyBuf, command.getRange().getLowerBound().getValue().orElse(0L),
                                command.getRange().getUpperBound().getValue().get());
            return m.map(v -> new NumericResponse<>(command, v));
        });
    }
    
    private static final RedisStrictCommand ZREMRANGEBYSCORE = new RedisStrictCommand("ZREMRANGEBYSCORE");

    @Override
    public Flux> zRemRangeByScore(
            Publisher commands) {
        return execute(commands, command -> {

            Assert.notNull(command.getKey(), "Key must not be null!");
            Assert.notNull(command.getRange(), "Range must not be null!");

            byte[] keyBuf = toByteArray(command.getKey());
            Mono m = write(keyBuf, StringCodec.INSTANCE, ZREMRANGEBYSCORE, 
                                keyBuf, toLowerBound(command.getRange()),
                                toUpperBound(command.getRange()));
            return m.map(v -> new NumericResponse<>(command, v));
        });
    }

    private static final RedisStrictCommand ZUNIONSTORE = new RedisStrictCommand("ZUNIONSTORE");
    
    @Override
    public Flux> zUnionStore(Publisher commands) {
        return execute(commands, command -> {

            Assert.notNull(command.getKey(), "Destination key must not be null!");
            Assert.notEmpty(command.getSourceKeys(), "Source keys must not be null or empty!");

            byte[] keyBuf = toByteArray(command.getKey());
            List args = new ArrayList(command.getSourceKeys().size() * 2 + 5);
            args.add(keyBuf);
            args.add(command.getSourceKeys().size());
            args.addAll(command.getSourceKeys().stream().map(e -> toByteArray(e)).collect(Collectors.toList()));
            if (!command.getWeights().isEmpty()) {
                args.add("WEIGHTS");
                for (Double weight : command.getWeights()) {
                    args.add(BigDecimal.valueOf(weight).toPlainString());
                }
            }
            if (command.getAggregateFunction().isPresent()) {
                args.add("AGGREGATE");
                args.add(command.getAggregateFunction().get().name());
            }
            Mono m = write(keyBuf, LongCodec.INSTANCE, ZUNIONSTORE, args.toArray());
            return m.map(v -> new NumericResponse<>(command, v));
        });
    }
    
    private static final RedisStrictCommand ZINTERSTORE = new RedisStrictCommand("ZINTERSTORE");

    @Override
    public Flux> zInterStore(Publisher commands) {
        return execute(commands, command -> {

            Assert.notNull(command.getKey(), "Destination key must not be null!");
            Assert.notEmpty(command.getSourceKeys(), "Source keys must not be null or empty!");

            byte[] keyBuf = toByteArray(command.getKey());
            List args = new ArrayList(command.getSourceKeys().size() * 2 + 5);
            args.add(keyBuf);
            args.add(command.getSourceKeys().size());
            args.addAll(command.getSourceKeys().stream().map(e -> toByteArray(e)).collect(Collectors.toList()));
            if (!command.getWeights().isEmpty()) {
                args.add("WEIGHTS");
                for (Double weight : command.getWeights()) {
                    args.add(BigDecimal.valueOf(weight).toPlainString());
                }
            }
            if (command.getAggregateFunction().isPresent()) {
                args.add("AGGREGATE");
                args.add(command.getAggregateFunction().get().name());
            }
            Mono m = write(keyBuf, LongCodec.INSTANCE, ZINTERSTORE, args.toArray());
            return m.map(v -> new NumericResponse<>(command, v));
        });
    }
    
    private static final RedisCommand> ZRANGEBYLEX = new RedisCommand>("ZRANGEBYLEX", new ObjectSetReplayDecoder());
    private static final RedisCommand> ZREVRANGEBYLEX = new RedisCommand>("ZREVRANGEBYLEX", new ObjectSetReplayDecoder());

    @Override
    public Flux>> zRangeByLex(
            Publisher commands) {
        return execute(commands, command -> {

            Assert.notNull(command.getKey(), "Key must not be null!");
            Assert.notNull(command.getRange(), "Range must not be null!");

            byte[] keyBuf = toByteArray(command.getKey());

            String start = null;
            String end = null;
            if (command.getDirection() == Direction.ASC) {
                start = toLexLowerBound(command.getRange(), "-");
                end = toLexUpperBound(command.getRange(), "+");
            } else {
                start = toLexUpperBound(command.getRange(), "-");
                end = toLexLowerBound(command.getRange(), "+");
            }
            
            Mono> m;
            if (!command.getLimit().isUnlimited()) {
                if (command.getDirection() == Direction.ASC) {
                    m = read(keyBuf, ByteArrayCodec.INSTANCE, ZRANGEBYLEX, 
                                keyBuf, start, end, "LIMIT", command.getLimit().getOffset(), command.getLimit().getCount());
                } else {
                    m = read(keyBuf, ByteArrayCodec.INSTANCE, ZREVRANGEBYLEX, 
                                keyBuf, start, end, "LIMIT", command.getLimit().getOffset(), command.getLimit().getCount());
                }
            } else {
                if (command.getDirection() == Direction.ASC) {
                    m = read(keyBuf, ByteArrayCodec.INSTANCE, ZRANGEBYLEX, 
                                keyBuf, start, end);
                } else {
                    m = read(keyBuf, ByteArrayCodec.INSTANCE, ZREVRANGEBYLEX, 
                                keyBuf, start, end);
                }
            }
            Flux flux = m.flatMapMany(e -> Flux.fromIterable(e).map(v -> ByteBuffer.wrap(v)));
            return Mono.just(new CommandResponse<>(command, flux));
        });
    }

}