diff --git a/core/riot-core/src/main/java/com/redis/riot/core/FakerItemReader.java b/core/riot-core/src/main/java/com/redis/riot/core/FakerItemReader.java index 9f16af86c..792cadcf8 100644 --- a/core/riot-core/src/main/java/com/redis/riot/core/FakerItemReader.java +++ b/core/riot-core/src/main/java/com/redis/riot/core/FakerItemReader.java @@ -1,12 +1,25 @@ package com.redis.riot.core; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.Locale; import java.util.Map; +import java.util.Map.Entry; import org.springframework.batch.item.ItemReader; import org.springframework.batch.item.support.AbstractItemCountingItemStreamItemReader; +import org.springframework.expression.EvaluationContext; +import org.springframework.expression.Expression; +import org.springframework.expression.spel.standard.SpelExpressionParser; +import org.springframework.expression.spel.support.SimpleEvaluationContext; +import org.springframework.expression.spel.support.SimpleEvaluationContext.Builder; import org.springframework.util.Assert; import org.springframework.util.ClassUtils; +import com.redis.spring.batch.common.IntRange; + +import net.datafaker.Faker; + /** * {@link ItemReader} that generates HashMaps using Faker. * @@ -14,42 +27,83 @@ */ public class FakerItemReader extends AbstractItemCountingItemStreamItemReader> { - public static final int DEFAULT_START = 1; - public static final int DEFAULT_COUNT = 1000; + private static final String FIELD_THREAD = "thread"; + public static final String FIELD_INDEX = "index"; + public static final Locale DEFAULT_LOCALE = Locale.getDefault(); + + private final SpelExpressionParser parser = new SpelExpressionParser(); + + private final Map expressions = new LinkedHashMap<>(); + private IntRange indexRange = IntRange.from(1); + private Locale locale = DEFAULT_LOCALE; + private boolean includeMetadata; - private int start = DEFAULT_START; - private int count = DEFAULT_COUNT; - private final Generator> generator; + private EvaluationContext context; - public FakerItemReader(Generator> generator) { + public FakerItemReader() { setName(ClassUtils.getShortName(getClass())); - Assert.notNull(generator, "A generator is required"); - setMaxItemCount(count); - this.generator = generator; } - public void setStart(int start) { - this.start = start; + public FakerItemReader withIndexRange(IntRange range) { + this.indexRange = range; + return this; } - public void setCount(int count) { - this.count = count; - setMaxItemCount(count); + public FakerItemReader withLocale(Locale locale) { + this.locale = locale; + return this; + } + + public FakerItemReader withIncludeMetadata(boolean include) { + this.includeMetadata = include; + return this; + } + + public FakerItemReader withField(String field, String expression) { + this.expressions.put(field, parser.parseExpression(expression)); + return this; + } + + public FakerItemReader withFields(String... fields) { + Assert.isTrue(fields.length % 2 == 0, + "fields.length must be a multiple of 2 and contain a sequence of field1, expression1, field2, expression2, fieldN, expressionN"); + for (int i = 0; i < fields.length; i += 2) { + withField(fields[i], fields[i + 1]); + } + return this; } @Override protected Map doRead() throws Exception { - return generator.next(start + ((getCurrentItemCount() - 1) % count)); + Map map = new HashMap<>(); + int index = index(); + if (includeMetadata) { + map.put(FIELD_INDEX, index); + map.put(FIELD_THREAD, Thread.currentThread().getId()); + } + context.setVariable(FIELD_INDEX, index); + for (Entry expression : expressions.entrySet()) { + map.put(expression.getKey(), expression.getValue().getValue(context)); + } + return map; + } + + private int index() { + return (indexRange.getMin() + getCurrentItemCount() - 1) % indexRange.getMax(); } @Override protected void doOpen() throws Exception { - // nothing to see here, move along + Faker faker = new Faker(locale); + Builder contextBuilder = SimpleEvaluationContext.forPropertyAccessors(new ReflectivePropertyAccessor()); + contextBuilder.withInstanceMethods(); + contextBuilder.withRootObject(faker); + this.context = contextBuilder.build(); } @Override protected void doClose() throws Exception { - // nothing to see here, move along + } } diff --git a/core/riot-core/src/main/java/com/redis/riot/core/Generator.java b/core/riot-core/src/main/java/com/redis/riot/core/Generator.java deleted file mode 100644 index dd18cf50f..000000000 --- a/core/riot-core/src/main/java/com/redis/riot/core/Generator.java +++ /dev/null @@ -1,6 +0,0 @@ -package com.redis.riot.core; - -public interface Generator { - - T next(int index); -} diff --git a/core/riot-core/src/main/java/com/redis/riot/core/MapGenerator.java b/core/riot-core/src/main/java/com/redis/riot/core/MapGenerator.java deleted file mode 100644 index b404029e4..000000000 --- a/core/riot-core/src/main/java/com/redis/riot/core/MapGenerator.java +++ /dev/null @@ -1,71 +0,0 @@ -package com.redis.riot.core; - -import java.util.HashMap; -import java.util.LinkedHashMap; -import java.util.Locale; -import java.util.Map; -import java.util.Map.Entry; - -import org.springframework.expression.EvaluationContext; -import org.springframework.expression.Expression; -import org.springframework.expression.spel.standard.SpelExpressionParser; -import org.springframework.expression.spel.support.SimpleEvaluationContext; - -import net.datafaker.Faker; - -public class MapGenerator implements Generator> { - - public static final String FIELD_INDEX = "index"; - - private final EvaluationContext context; - private final Map expressions; - - public MapGenerator(EvaluationContext context, Map expressions) { - this.context = context; - this.expressions = expressions; - } - - @Override - public Map next(int index) { - Map map = new HashMap<>(); - context.setVariable(FIELD_INDEX, index); - for (Entry expression : expressions.entrySet()) { - map.put(expression.getKey(), expression.getValue().getValue(context)); - } - return map; - } - - public static MapGeneratorBuilder builder() { - return new MapGeneratorBuilder(); - } - - public static class MapGeneratorBuilder { - - private final SpelExpressionParser parser = new SpelExpressionParser(); - private Locale locale = Locale.getDefault(); - private final Map expressions = new LinkedHashMap<>(); - - public MapGeneratorBuilder locale(Locale locale) { - this.locale = locale; - return this; - } - - public MapGeneratorBuilder field(String field, String expression) { - this.expressions.put(field, parser.parseExpression(expression)); - return this; - } - - public MapGeneratorBuilder fields(Map fields) { - fields.forEach((k, v) -> expressions.put(k, parser.parseExpression(v))); - return this; - } - - public MapGenerator build() { - Faker faker = new Faker(locale); - SimpleEvaluationContext context = new SimpleEvaluationContext.Builder(new ReflectivePropertyAccessor()) - .withInstanceMethods().withRootObject(faker).build(); - return new MapGenerator(context, expressions); - } - - } -} diff --git a/core/riot-core/src/main/java/com/redis/riot/core/MapWithMetadataGenerator.java b/core/riot-core/src/main/java/com/redis/riot/core/MapWithMetadataGenerator.java deleted file mode 100644 index 13c695c06..000000000 --- a/core/riot-core/src/main/java/com/redis/riot/core/MapWithMetadataGenerator.java +++ /dev/null @@ -1,22 +0,0 @@ -package com.redis.riot.core; - -import java.util.Map; - -public class MapWithMetadataGenerator implements Generator> { - - private static final String FIELD_THREAD = "thread"; - - private final Generator> parent; - - public MapWithMetadataGenerator(Generator> parent) { - this.parent = parent; - } - - @Override - public Map next(int index) { - Map map = parent.next(index); - map.put(MapGenerator.FIELD_INDEX, index); - map.put(FIELD_THREAD, Thread.currentThread().getId()); - return map; - } -} diff --git a/plugins/riot/src/main/java/com/redis/riot/cli/FakerImport.java b/plugins/riot/src/main/java/com/redis/riot/cli/FakerImport.java index 90ac7e689..62b110753 100644 --- a/plugins/riot/src/main/java/com/redis/riot/cli/FakerImport.java +++ b/plugins/riot/src/main/java/com/redis/riot/cli/FakerImport.java @@ -18,9 +18,6 @@ import com.redis.riot.cli.common.FakerImportOptions; import com.redis.riot.cli.common.StepProgressMonitor; import com.redis.riot.core.FakerItemReader; -import com.redis.riot.core.Generator; -import com.redis.riot.core.MapGenerator; -import com.redis.riot.core.MapWithMetadataGenerator; import picocli.CommandLine.Command; import picocli.CommandLine.Mixin; @@ -44,9 +41,12 @@ public void setOptions(FakerImportOptions options) { @Override protected Job job(CommandContext context) { log.log(Level.FINE, "Creating Faker reader with {0}", options); - FakerItemReader reader = new FakerItemReader(generator(context)); - reader.setCurrentItemCount(options.getStart() - 1); + FakerItemReader reader = new FakerItemReader(); reader.setMaxItemCount(options.getCount()); + reader.withIndexRange(options.getIndexRange()); + fields(context).forEach(reader::withField); + reader.withLocale(options.getLocale()); + reader.withIncludeMetadata(options.isIncludeMetadata()); SimpleStepBuilder, Map> step = step(context.getRedisClient(), reader); StepProgressMonitor monitor = progressMonitor("Generating"); monitor.withInitialMax(options.getCount()); @@ -54,25 +54,20 @@ protected Job job(CommandContext context) { return job(commandName()).start(step.build()).build(); } - private void addFieldsFromIndex(CommandContext context, String index, Map fields) { - try (StatefulRedisModulesConnection connection = RedisModulesUtils - .connection(context.getRedisClient())) { - RediSearchCommands commands = connection.sync(); - IndexInfo info = RedisModulesUtils.indexInfo(commands.ftInfo(index)); - for (Field field : info.getFields()) { - fields.put(field.getName(), expression(field)); + private Map fields(CommandContext context) { + Map fields = new LinkedHashMap<>(); + fields.putAll(options.getFields()); + options.getRedisearchIndex().ifPresent(index -> { + try (StatefulRedisModulesConnection connection = RedisModulesUtils + .connection(context.getRedisClient())) { + RediSearchCommands commands = connection.sync(); + IndexInfo info = RedisModulesUtils.indexInfo(commands.ftInfo(index)); + for (Field field : info.getFields()) { + fields.put(field.getName(), expression(field)); + } } - } - } - - private Generator> generator(CommandContext context) { - Map fields = new LinkedHashMap<>(options.getFields()); - options.getRedisearchIndex().ifPresent(index -> addFieldsFromIndex(context, index, fields)); - MapGenerator generator = MapGenerator.builder().locale(options.getLocale()).fields(fields).build(); - if (options.isIncludeMetadata()) { - return new MapWithMetadataGenerator(generator); - } - return generator; + }); + return fields; } private String expression(Field field) { diff --git a/plugins/riot/src/main/java/com/redis/riot/cli/common/FakerImportOptions.java b/plugins/riot/src/main/java/com/redis/riot/cli/common/FakerImportOptions.java index 71720df60..140e7d196 100644 --- a/plugins/riot/src/main/java/com/redis/riot/cli/common/FakerImportOptions.java +++ b/plugins/riot/src/main/java/com/redis/riot/cli/common/FakerImportOptions.java @@ -7,20 +7,22 @@ import org.springframework.util.Assert; +import com.redis.spring.batch.common.IntRange; + import picocli.CommandLine.Option; import picocli.CommandLine.Parameters; public class FakerImportOptions { - public static final int DEFAULT_START = 1; public static final int DEFAULT_COUNT = 1000; public static final Locale DEFAULT_LOCALE = Locale.ENGLISH; public static final boolean DEFAULT_INCLUDE_METADATA = false; + private static final IntRange DEFAULT_INDEX_RANGE = IntRange.from(1); - @Option(names = "--start", description = "Start index (default: ${DEFAULT-VALUE}).", paramLabel = "") - protected int start = DEFAULT_START; @Option(names = "--count", description = "Number of items to generate (default: ${DEFAULT-VALUE}).", paramLabel = "") private int count = DEFAULT_COUNT; + @Option(names = "--index", description = "Index range (default: ${DEFAULT-VALUE}).", paramLabel = "") + private IntRange indexRange = DEFAULT_INDEX_RANGE; @Parameters(arity = "0..*", description = "SpEL expressions in the form field1=\"exp\" field2=\"exp\"...", paramLabel = "SPEL") private Map fields = new LinkedHashMap<>(); @Option(names = "--infer", description = "Introspect given RediSearch index to infer Faker fields.", paramLabel = "") @@ -30,12 +32,12 @@ public class FakerImportOptions { @Option(names = "--metadata", description = "Include metadata (index, partition).") private boolean includeMetadata = DEFAULT_INCLUDE_METADATA; - public int getStart() { - return start; + public IntRange getIndexRange() { + return indexRange; } - public void setStart(int start) { - this.start = start; + public void setIndexRange(IntRange indexRange) { + this.indexRange = indexRange; } public int getCount() { @@ -59,8 +61,8 @@ public Optional getRedisearchIndex() { return redisearchIndex; } - public void setFakerIndex(String fakerIndex) { - this.redisearchIndex = Optional.of(fakerIndex); + public void setRedisearchIndex(Optional redisearchIndex) { + this.redisearchIndex = redisearchIndex; } public Locale getLocale() { @@ -82,8 +84,9 @@ public void setIncludeMetadata(boolean includeMetadata) { @Override public String toString() { - return "FakerGeneratorOptions [fields=" + fields + ", redisearchIndex=" + redisearchIndex + ", locale=" + locale - + ", includeMetadata=" + includeMetadata + ", count=" + count + "]"; + return "FakerImportOptions [indexRange=" + indexRange + ", count=" + count + ", fields=" + fields + + ", redisearchIndex=" + redisearchIndex + ", locale=" + locale + ", includeMetadata=" + includeMetadata + + "]"; } } diff --git a/plugins/riot/src/main/java/com/redis/riot/cli/common/GenerateOptions.java b/plugins/riot/src/main/java/com/redis/riot/cli/common/GenerateOptions.java index 8e3b84b3d..a7686c939 100644 --- a/plugins/riot/src/main/java/com/redis/riot/cli/common/GenerateOptions.java +++ b/plugins/riot/src/main/java/com/redis/riot/cli/common/GenerateOptions.java @@ -27,7 +27,7 @@ public class GenerateOptions { public static final int DEFAULT_COUNT = 1000; - public static final IntRange DEFAULT_KEY_RANGE = IntRange.between(1, 1000); + public static final IntRange DEFAULT_KEY_RANGE = IntRange.from(1); @Option(names = "--count", description = "Number of items to generate (default: ${DEFAULT-VALUE}).", paramLabel = "") private int count = DEFAULT_COUNT;