com.redislabs.riot.gen.GeneratorImportCommand Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of riot-gen Show documentation
RIOT Core module
The newest version!
package com.redislabs.riot.gen;

import com.redislabs.mesclun.RedisModulesClient;
import com.redislabs.mesclun.StatefulRedisModulesConnection;
import com.redislabs.mesclun.search.Field;
import com.redislabs.mesclun.search.IndexInfo;
import com.redislabs.mesclun.search.RediSearchCommands;
import com.redislabs.mesclun.search.RediSearchUtils;
import com.redislabs.riot.AbstractImportCommand;
import com.redislabs.riot.ProcessorOptions;
import com.redislabs.riot.RedisOptions;
import com.redislabs.riot.RiotStepBuilder;
import lombok.extern.slf4j.Slf4j;
import org.springframework.batch.core.configuration.annotation.StepBuilderFactory;
import org.springframework.batch.core.job.flow.Flow;
import org.springframework.batch.core.step.builder.StepBuilder;
import org.springframework.batch.item.ItemProcessor;
import org.springframework.batch.item.ItemReader;
import picocli.CommandLine;
import picocli.CommandLine.Command;

import java.util.LinkedHashMap;
import java.util.Map;

@SuppressWarnings("FieldMayBeFinal")
@Slf4j
@Command(name = "import", description = "Import generated data")
public class GeneratorImportCommand extends AbstractImportCommand, Map> {

    @CommandLine.Mixin
    private GenerateOptions options = new GenerateOptions();
    @CommandLine.ArgGroup(exclusive = false, heading = "Processor options%n")
    private ProcessorOptions processorOptions = new ProcessorOptions();

    @Override
    protected Flow flow(StepBuilderFactory stepBuilderFactory) throws Exception {
        StepBuilder stepBuilder = stepBuilderFactory.get("generate-step");
        return flow(step(stepBuilder, "Generating", reader()).build());
    }

    private ItemReader> reader() {
        log.debug("Creating Faker reader with {}", options);
        FakerItemReader reader = FakerItemReader.builder().generator(generator()).start(options.getStart()).end(options.getEnd()).build();
        if (options.getSleep() > 0) {
            return new ThrottledItemReader<>(reader, options.getSleep());
        }
        return reader;
    }

    private Generator> generator() {
        Map fields = options.getFakerFields() == null ? new LinkedHashMap<>() : new LinkedHashMap<>(options.getFakerFields());
        if (options.getFakerIndex() != null) {
            fields.putAll(fieldsFromIndex(options.getFakerIndex()));
        }
        MapGenerator generator = MapGenerator.builder().locale(options.getLocale()).fields(fields).build();
        if (options.isIncludeMetadata()) {
            return new MapWithMetadataGenerator(generator);
        }
        return generator;
    }

    private String expression(Field field) {
        switch (field.getType()) {
            case TEXT:
                return "lorem.paragraph";
            case TAG:
                return "number.digits(10)";
            case GEO:
                return "address.longitude.concat(',').concat(address.latitude)";
            default:
                return "number.randomDouble(3,-1000,1000)";
        }
    }

    private Map fieldsFromIndex(String index) {
        Map fields = new LinkedHashMap<>();
        RedisModulesClient client = RedisModulesClient.create(getRedisOptions().uris().get(0));
        try (StatefulRedisModulesConnection connection = client.connect()) {
            RediSearchCommands commands = connection.sync();
            IndexInfo info = RediSearchUtils.getInfo(commands.ftInfo(index));
            for (Field field : info.getFields()) {
                fields.put(field.getName(), expression(field));
            }
        } finally {
            RedisOptions.shutdown(client);
        }
        return fields;
    }

    @Override
    protected ItemProcessor, Map> processor() throws NoSuchMethodException {
        return processorOptions.processor(getRedisOptions());
    }

    @Override
    protected  RiotStepBuilder riotStep(StepBuilder stepBuilder, String taskName) {
        RiotStepBuilder riotStepBuilder = super.riotStep(stepBuilder, taskName);
        riotStepBuilder.initialMax(() -> options.getEnd() - options.getStart());
        return riotStepBuilder;
    }

}