refactor: Consolidated generator classes

core/riot-core/src/main/java/com/redis/riot/core/FakerItemReader.java

@@ -1,55 +1,109 @@
 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.
  *
  * @author Julien Ruaux
  */
 public class FakerItemReader extends AbstractItemCountingItemStreamItemReader<Map<String, Object>> {
 
-	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<String, Expression> 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<Map<String, Object>> generator;
+	private EvaluationContext context;
 
-	public FakerItemReader(Generator<Map<String, Object>> 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<String, Object> doRead() throws Exception {
-		return generator.next(start + ((getCurrentItemCount() - 1) % count));
+		Map<String, Object> 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<String, Expression> 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
+
 	}
 
 }

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,35 +41,33 @@ 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<String, Object>, Map<String, Object>> step = step(context.getRedisClient(), reader);
 		StepProgressMonitor monitor = progressMonitor("Generating");
 		monitor.withInitialMax(options.getCount());
 		monitor.register(step);
 		return job(commandName()).start(step.build()).build();
 	}
 
-	private void addFieldsFromIndex(CommandContext context, String index, Map<String, String> fields) {
-		try (StatefulRedisModulesConnection<String, String> connection = RedisModulesUtils
-				.connection(context.getRedisClient())) {
-			RediSearchCommands<String, String> commands = connection.sync();
-			IndexInfo info = RedisModulesUtils.indexInfo(commands.ftInfo(index));
-			for (Field<String> field : info.getFields()) {
-				fields.put(field.getName(), expression(field));
+	private Map<String, String> fields(CommandContext context) {
+		Map<String, String> fields = new LinkedHashMap<>();
+		fields.putAll(options.getFields());
+		options.getRedisearchIndex().ifPresent(index -> {
+			try (StatefulRedisModulesConnection<String, String> connection = RedisModulesUtils
+					.connection(context.getRedisClient())) {
+				RediSearchCommands<String, String> commands = connection.sync();
+				IndexInfo info = RedisModulesUtils.indexInfo(commands.ftInfo(index));
+				for (Field<String> field : info.getFields()) {
+					fields.put(field.getName(), expression(field));
+				}
 			}
-		}
-	}
-
-	private Generator<Map<String, Object>> generator(CommandContext context) {
-		Map<String, String> 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<String> field) {

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 = "<int>")
-	protected int start = DEFAULT_START;
 	@Option(names = "--count", description = "Number of items to generate (default: ${DEFAULT-VALUE}).", paramLabel = "<int>")
 	private int count = DEFAULT_COUNT;
+	@Option(names = "--index", description = "Index range (default: ${DEFAULT-VALUE}).", paramLabel = "<range>")
+	private IntRange indexRange = DEFAULT_INDEX_RANGE;
 	@Parameters(arity = "0..*", description = "SpEL expressions in the form field1=\"exp\" field2=\"exp\"...", paramLabel = "SPEL")
 	private Map<String, String> fields = new LinkedHashMap<>();
 	@Option(names = "--infer", description = "Introspect given RediSearch index to infer Faker fields.", paramLabel = "<name>")
@@ -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<String> getRedisearchIndex() {
 		return redisearchIndex;
 	}
 
-	public void setFakerIndex(String fakerIndex) {
-		this.redisearchIndex = Optional.of(fakerIndex);
+	public void setRedisearchIndex(Optional<String> 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
+				+ "]";
 	}
 
 }

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 = "<int>")
 	private int count = DEFAULT_COUNT;